"""Tests of the static blocks: Restrepo replication, KKT conditions, national accounting identity."""
import os
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from agi_ge.calibration import calibrate  # noqa: E402
from agi_ge.production import housing_rent, income_identity, static_period, value_added_by_sector  # noqa: E402
from agi_ge.tasks import Technology, restrepo_benchmark, technology_from_params  # noqa: E402

NU = [0.3, 0.25, 0.25, 0.2, 0, 0]
GAMMA = [0, 0, 0, 0, 0.05, 0.05]
SKILL = [0, 0, 1, 1, 0, 1]
ALPHA = np.array([1.0, 2.0, 4.0, 8.0, 3.0, 6.0])
H = [0.6, 0.4]


@pytest.fixture(scope="module")
def bench():
    return restrepo_benchmark(NU, GAMMA, SKILL, -1.0, ALPHA, H, [0.0, 1.0, 10.0, 100.0, 1e3, 1e4, 1e5])


def test_kkt_conditions(bench):
    tech = Technology(nu=np.array(NU), gamma=np.array(GAMMA), skill=np.array(SKILL), rho_B=-1.0,
                      nu_N=0, rho_N=0, nu_T=0, nu_E=0, rho_T=0)
    for s in bench:
        assert tech.check_kkt(s, H, s["Q_total"], ALPHA) < 1e-6


def test_restrepo_P1_all_bottlenecks_automated(bench):
    s = bench[-1]
    assert np.all(s["Q"][:4] > 0)                       # compute used in every bottleneck task


def test_restrepo_P2_P3_output_linear_in_compute(bench):
    s1, s2 = bench[-2], bench[-1]
    A = s2["A"]
    # Y = A (Q + sum CEU_tilde H): slope in Q equals A, intercept equals the compute-equivalent labor
    slope = (s2["Y"] - s1["Y"]) / (s2["Q_total"] - s1["Q_total"])
    assert abs(slope / A - 1) < 1e-6
    intercept = s2["Y"] - A * s2["Q_total"]
    assert abs(intercept - np.sum(s2["w"] * np.array(H))) < 1e-6


def test_restrepo_P4_wages_equal_A_times_CEU(bench):
    s = bench[-1]
    A = s["A"]
    # skill H: its best automated task has alpha = 8 = CEU(H) -> w_H = A * 8
    assert abs(s["w"][1] / (A * 8.0) - 1) < 1e-6
    # skill L: best automated task alpha = 2 (accessory task alpha=3 is not worth automating): discounted CEU
    assert abs(s["w"][0] / (A * 2.0) - 1) < 1e-6
    assert s["w"][0] < A * 3.0


def test_restrepo_P5_labor_share_vanishes(bench):
    ls = [s["labor_share"] for s in bench]
    assert ls[0] == pytest.approx(1.0)
    assert ls[-1] < 1e-3
    assert bench[-1]["compute_share"] > 0.999


def test_restrepo_P6_total_wages_higher_with_AGI(bench):
    pre = np.sum(bench[0]["w"] * np.array(H))
    post = np.sum(bench[-1]["w"] * np.array(H))
    assert post > pre


@pytest.fixture(scope="module")
def calibrated():
    p, info = calibrate()
    return p, info, technology_from_params(p)


def test_income_identity_holds(calibrated):
    p, info, tech = calibrated
    al = p.alpha_path()
    for t, KQ in [(0, 1.0), (20, 50.0), (60, 5e3)]:
        per = static_period(p, tech, p.KN0, KQ, p.KE0, p.KH0, al[t], p.Z0, p.eps_path()[t], gf=1.0)
        housing_rent(p, per, 0.6)
        assert abs(income_identity(p, per)) < 1e-6 * per["GDP"]
        va = value_added_by_sector(p, per)
        assert abs(va["total"] - per["GDP"]) < 1e-6 * per["GDP"]


def test_income_identity_fails_when_accounting_is_broken(calibrated):
    """A deliberately inconsistent accounting (compute income counted twice) must be detected."""
    p, info, tech = calibrated
    per = static_period(p, tech, p.KN0, 1.0, p.KE0, p.KH0, p.alpha_path()[0], p.Z0, p.eps0, gf=1.0)
    housing_rent(p, per, 0.6)
    per_bad = dict(per)
    per_bad["GDP"] = per["GDP"] + per["A"] * per["Qtot"]      # counting compute services as final output
    assert abs(income_identity(p, per_bad)) > 1e-3


def test_calibration_shares(calibrated):
    p, info, tech = calibrated
    per = static_period(p, tech, p.KN0, p.KQ0, p.KE0, p.KH0, p.alpha_path()[0], p.Z0, p.eps0, gf=1.0)
    housing_rent(p, per, info["C"])
    assert per["wage_bill"] / per["GDP"] == pytest.approx(0.568303, abs=2e-3)
    assert per["PE"] * per["E"] / per["GDP"] == pytest.approx(0.04, abs=2e-3)
    assert per["util"] == 1.0


def test_energy_bottleneck_lowers_utilization(calibrated):
    p, info, tech = calibrated
    al = p.alpha_path()[60]
    per = static_period(p, tech, p.KN0, 1e7, p.KE0, p.KH0, al, p.Z0, p.eps_path()[60], gf=1.0)
    assert per["util"] < 1.0
    assert abs(per["A"] - p.eps_path()[60] * per["PE"]) < 1e-8 * per["A"]   # idle capacity: A = energy cost