"""Task-based production block: zero-AI limit, uniqueness of the automation threshold, Euler's theorem, shares."""
import numpy as np
import pytest
from hypothesis import given, settings, strategies as st

from age.economy.tasks import TaskTech


def test_zero_ai_limit_is_cobb_douglas():
    t = TaskTech(alpha=0.33)
    K, L = 3.0, 1.0
    Y = t.output(0.0, K, 0.0, L)
    assert np.isclose(Y, K ** 0.33 * L ** 0.67)
    pr = t.prices(0.0, K, 0.0, L)
    assert pr['iota'] == 0.0 and pr['pQ'] == 0.0
    assert np.isclose(pr['share_L'], 0.67)


@settings(max_examples=40, deadline=None)
@given(QP=st.floats(1e-4, 50.0), L=st.floats(0.2, 3.0), psi1=st.floats(0.5, 6.0))
def test_threshold_unique_and_monotone(QP, L, psi1):
    t = TaskTech(psi1=psi1, iota_bar=0.9)
    i1 = t.iota_of(QP, L)
    i2 = t.iota_of(QP * 1.5, L)
    assert 0.0 < i1 <= 0.9
    assert i2 >= i1 - 1e-12          # more compute -> more automation (capped)
    if i1 < 0.9 - 1e-9:
        # interior: threshold condition psi(iota) = iota/(1-iota) L/QP holds
        assert np.isclose(t.psi(i1), i1 / (1 - i1) * L / QP, rtol=1e-8)


def test_euler_theorem_and_shares():
    """Factor payments exhaust output (CRS in K, L, QP): RK K + w L + pQ QP = Y."""
    t = TaskTech(alpha=0.33, iota_bar=0.6)
    K, L, QP = 4.0, 1.0, 0.7
    pr = t.prices(0.1, K, QP, L)
    total = pr['RK'] * K + pr['w'] * L + pr['pQ'] * QP
    assert np.isclose(total, pr['Y'], rtol=1e-10)
    assert np.isclose(pr['share_K'] + pr['share_Q'] + pr['share_L'], 1.0)


def test_marginal_products_numerically():
    t = TaskTech(alpha=0.33, iota_bar=0.9)
    K, L, QP, z = 4.0, 1.0, 0.5, 0.0
    pr = t.prices(z, K, QP, L)
    h = 1e-6
    dYdK = (t.output(z, K + h, QP, L) - t.output(z, K - h, QP, L)) / (2 * h)
    dYdL = (t.output(z, K, QP, L + h) - t.output(z, K, QP, L - h)) / (2 * h)
    dYdQ = (t.output(z, K, QP + h, L) - t.output(z, K, QP - h, L)) / (2 * h)
    # envelope: iota re-optimised inside output(); marginal products equal competitive prices
    assert np.isclose(dYdK, pr['RK'], rtol=1e-5)
    assert np.isclose(dYdL, pr['w'], rtol=1e-4)
    assert np.isclose(dYdQ, pr['pQ'], rtol=1e-4)


def test_compute_demand_inverts_prices():
    t = TaskTech(alpha=0.33, iota_bar=0.9)
    K, L, QP, z = 4.0, 1.0, 0.5, 0.0
    pr = t.prices(z, K, QP, L)
    QP_back = t.production_compute_demand(pr['pQ'], pr['w'], L, z, K)
    assert np.isclose(QP_back, QP, rtol=1e-6)