"""Calibration of the 2026 base point (US) from data/processed/calibration_targets.csv.

Units: GDP in 2026 = 1 (per efficiency unit of labor/land/resource, gf = 1); workers H_L + H_H = 1.
The base point is a balanced-growth path of the frozen-AI economy (alpha fixed at its 2026 value). CES weights
are set so that the base-point cost shares equal the targets (share_j = nu_j x_j^rho / sum, hence
nu_j ∝ share_j x_j^{-rho}); the aggregate scale Z_0 then normalizes final output. Every target has a
`source` and a `status` ("data" if read from the data file, "assumed" if a documented assumption,
"scenario" if a forward-looking scenario input). The result is checked by re-solving the BGP (steady.py).
"""
from __future__ import annotations

import csv
import json
import os
from dataclasses import dataclass, field, asdict
from typing import Dict

import numpy as np

from .ces import ces, rho_of_sigma
from .params import Params, TaskSpec
from .steady import bgp_interest_rate, solve_bgp, target_rentals

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
TARGETS_CSV = os.path.join(ROOT, "data", "processed", "calibration_targets.csv")


@dataclass
class Target:
    value: float
    source: str
    status: str = "assumed"      # data | assumed | scenario


def _default_targets() -> Dict[str, Target]:
    """Documented assumptions used when a data target is not retrieved. Overwritten by the data file."""
    A = "assumed (see docs/assumptions_register.md)"
    L = "literature value (see docs/assumptions_register.md)"
    return {
        "labor_share": Target(0.58, A),
        "housing_va_share": Target(0.12, A + "; BEA housing services ~12% of GDP, not retrieved as a series"),
        "KN_over_GDP": Target(1.22, A),                 # overwritten by data; K_N is residual anyway
        "KH_over_GDP": Target(1.17, A),
        "KE_over_GDP": Target(0.30, A),
        "VQ_over_GDP": Target(0.06, A),                 # market value of compute capital / GDP (derived, see CSV_DERIVED)
        "energy_share": Target(0.04, A),                # energy sector gross output / GDP
        "datacenter_energy_share": Target(0.044, A),    # share of electricity used by data centers
        "land_rent_share": Target(0.03, A),             # production land rent / GDP
        "resource_share_energy": Target(0.30, A),       # eta (implied in calibration)
        "housing_land_share": Target(0.35, A),          # nu_H (implied in calibration)
        "delta_N": Target(0.089, A),
        "delta_H": Target(0.023, A),
        "delta_E": Target(0.05, A),
        "delta_Q": Target(0.25, A),
        "r_star": Target(0.035, A),                     # real return on wealth in the frozen-AI BGP
        "g_Z": Target(0.015, A),                        # trend growth per efficiency unit
        "gamma_u": Target(2.0, L),
        "share_L_workers": Target(0.60, A),             # share of workers with skill L
        "skill_premium": Target(1.7, A),                # w_H / w_L per worker
        "task_wage_shares_L": Target(0.40, A),          # share of L wage bill in task 1 (rest in task 2)
        "task_wage_shares_H": Target(0.50, A),          # share of H wage bill in task 3 (rest in task 4)
        "sigma_B": Target(0.5, L),
        "sigma_N": Target(0.7, L),
        "sigma_T": Target(0.6, L),
        "phi_N": Target(3.0, L), "phi_Q": Target(1.0, L), "phi_E": Target(3.0, L), "phi_H": Target(3.0, L),
        "g_pQ": Target(0.15, "scenario"), "g_eps": Target(0.10, "scenario"), "g_alpha": Target(0.10, "scenario"),
        "pQ_floor_ratio": Target(1e-2, "scenario"), "eps_floor_ratio": Target(1e-2, "scenario"),
    }


# calibration_targets.csv target_id -> target name (value used directly)
CSV_MAP = {
    "T01_labor_share_PWT_US_latest": "labor_share",
    "T06_net_stock_private_nonres_fixed_assets_GDP_latest": "KN_over_GDP",
    "T06_depreciation_rate_private_nonres_latest": "delta_N",
    "T10_datacenter_share_of_US_electricity_2023": ("datacenter_energy_share", 0.01),   # percent -> ratio
}

# derived data targets: name -> (function of the CSV rows, description)
def _derived(rows):
    out = {}
    g = lambda k: float(rows[k]["value"])
    try:
        tot, nonres = g("T06_net_stock_private_fixed_assets_GDP_latest"), g("T06_net_stock_private_nonres_fixed_assets_GDP_latest")
        out["KH_over_GDP"] = Target(tot - nonres, "BEA fixed assets via FRED: total minus nonresidential private net stock / GDP (2024)", "data")
        dtot, dnr = g("T06_depreciation_rate_private_total_latest"), g("T06_depreciation_rate_private_nonres_latest")
        out["delta_H"] = Target((dtot * tot - dnr * nonres) / (tot - nonres), "implied residential depreciation = (CFC_total - CFC_nonres)/K_res (BEA via FRED, 2024)", "data")
        ipe = g("T05_info_processing_equipment_GDP_latest")
        out["VQ_over_GDP"] = Target(ipe / (0.25 + 0.015 + 0.15), "info-processing-equipment investment/GDP (FRED Y034RC1Q027SBEA, 2026Q2) capitalized at delta_Q + g_Z + g_pQ", "data-derived")
    except KeyError:
        pass
    return out


def load_targets(path=TARGETS_CSV) -> Dict[str, Target]:
    tg = _default_targets()
    if os.path.exists(path):
        with open(path, newline="") as f:
            rows = {r["target_id"]: r for r in csv.DictReader(f)}
        for tid, spec in CSV_MAP.items():
            row = rows.get(tid)
            if row and row.get("status") == "RETRIEVED" and row.get("value") not in (None, ""):
                key, scale = (spec, 1.0) if isinstance(spec, str) else spec
                tg[key] = Target(float(row["value"]) * scale, f"{row.get('source')} {row.get('series_id')} ({row.get('year_or_period')})", "data")
        tg.update(_derived(rows))
    return tg


def calibrate(tg: Dict[str, Target] | None = None, verbose=False, overrides: Dict[str, float] | None = None):
    """overrides: target values to replace (e.g. scenario inputs g_pQ, sigma_N, r_star); the base point is
    re-derived so that 2026 remains consistent with the data under the scenario's anticipated price path."""
    tg = dict(load_targets() if tg is None else tg)
    for k, val in (overrides or {}).items():
        tg[k] = Target(val, f"override ({tg[k].source})" if k in tg else "override", tg[k].status if k in tg else "scenario")
    v = {k: t.value for k, t in tg.items()}
    g = v["g_Z"]
    gamma_u = v["gamma_u"]
    r = v["r_star"]
    beta = (1 + g) ** gamma_u / (1 + r)
    GDP = 1.0
    hv = v["housing_va_share"] * GDP
    Y = GDP - hv
    delta = {"N": v["delta_N"], "Q": v["delta_Q"], "E": v["delta_E"], "H": v["delta_H"]}
    phi = {"N": v["phi_N"], "Q": v["phi_Q"], "E": v["phi_E"], "H": v["phi_H"]}
    K = {"N": v["KN_over_GDP"] * GDP, "H": v["KH_over_GDP"] * GDP, "E": v["KE_over_GDP"] * GDP}
    VQ = v["VQ_over_GDP"] * GDP
    # --- rentals at the BGP (per unit of price); compute: user cost with the anticipated price decline
    #     R_Q,0 = p_0 [ (1+r) e^{g_p} - (1 - delta_Q) ]  (capital bought at t=-1 at price p_0 e^{g_p})
    Rst = target_rentals(Params(delta=delta, phi=phi, g_Z=g, beta=beta, gamma_u=gamma_u, pQ0=1.0), r, 1.0)
    Rst["Q"] = (1 + r) * np.exp(v["g_pQ"]) - (1 - delta["Q"])
    # --- energy block: E_0 = 1 ; P_E E = energy_share
    energy_val = v["energy_share"] * GDP
    E0 = 1.0
    PE = energy_val / E0
    RE = 1.0
    eta_implied = 1 - Rst["E"] * K["E"] / energy_val
    if verbose:
        print(f"  resource share in energy implied by K_E target: {eta_implied:.3f} (assumed {v['resource_share_energy']:.3f}) -> using implied")
    eta = float(np.clip(eta_implied, 0.05, 0.9))
    ZE = E0 / (K["E"] ** (1 - eta) * RE ** eta)
    # --- labor: workers, wages
    HL = v["share_L_workers"]
    HH = 1 - HL
    wage_bill = v["labor_share"] * GDP
    wL = wage_bill / (HL + HH * v["skill_premium"])
    wH = wL * v["skill_premium"]
    # --- tasks: wage shares across tasks within skill; task 1 also uses compute
    sL1 = v["task_wage_shares_L"]
    sH3 = v["task_wage_shares_H"]
    LB = {0: sL1 * HL, 1: (1 - sL1) * HL, 2: sH3 * HH, 3: (1 - sH3) * HH}   # labor in each task
    # --- compute block: K_Q = 1 by normalization of compute units; p^Q_0 = VQ; the compute Euler gives R_Q;
    #     the services price A = R_Q + eps_0 P_E; the task-1 tie A = w_L/alpha_1 gives alpha_1(2026); the
    #     compute share of task 1 is then implied (reported, not targeted).
    KQ = 1.0
    pQ0 = VQ / KQ
    eps0 = v["datacenter_energy_share"] * E0 / KQ
    RQ = Rst["Q"] * pQ0
    A = RQ + eps0 * PE
    alpha10 = wL / A
    Q_eff = KQ / alpha10                   # compute in task 1, worker equivalents
    c1 = Q_eff / (LB[0] + Q_eff)
    if verbose:
        print(f"  implied compute share of task-1 input in 2026: {c1:.3f}; alpha_1(2026) = {alpha10:.1f}")
    EY = E0 - eps0 * KQ
    # --- structures: K_N is residual so that factor payments add up to GDP = 1 exactly
    land_val = v["land_rent_share"] * GDP
    energy_in_Y = PE * (E0 - eps0 * KQ)
    RN = Rst["N"]
    valN = Y - wage_bill - A * KQ - land_val - energy_in_Y
    K["N"] = valN / RN
    if verbose:
        print(f"  K_N/GDP implied by income shares: {K['N']:.3f} (target {v['KN_over_GDP']:.3f})")
    # --- consumption / housing
    inv = sum((delta[k] + g) * K[k] + 0.5 * phi[k] * g ** 2 * K[k] for k in K) + pQ0 * ((delta["Q"] + g) * KQ + 0.5 * phi["Q"] * g ** 2 * KQ)
    C = Y - inv
    theta_H = hv / (hv + C)
    X = np.array([LB[0] + Q_eff, LB[1], LB[2], LB[3]])
    Fval = np.array([wL, wL, wH, wH])                           # marginal values of tasks
    task_val = Fval * X                                          # = wages + compute income in task 1
    rhoB = rho_of_sigma(v["sigma_B"])
    nuB = task_val * X ** (-rhoB)
    nuB = nuB / nuB.sum()
    XB = ces(nuB, X, rhoB)
    PB = task_val.sum() / XB                                     # price of the composite
    # --- V nest: structures vs tasks
    rhoN = rho_of_sigma(v["sigma_N"])
    valN, valB = RN * K["N"], PB * XB
    wN = np.array([valN, valB]) * np.array([K["N"], XB]) ** (-rhoN)
    wN = wN / wN.sum()
    nu_N = float(wN[0])
    V = ces([nu_N, 1 - nu_N], [K["N"], XB], rhoN)
    PV = (valN + valB) / V
    # --- top nest: land, energy, V ; land rent share target
    rhoT = rho_of_sigma(v["sigma_T"])
    T_land = 1.0
    valE = PE * EY
    # Y (gross output of final goods) must equal valT + valE + valV + accessory(=0 at base)
    valV = valN + valB
    Y_implied = land_val + valE + valV
    wT = np.array([land_val, valE, valV]) * np.array([T_land, EY, V]) ** (-rhoT)
    wT = wT / wT.sum()
    nu_T, nu_E = float(wT[0]), float(wT[1])
    Gval = ces(wT, [T_land, EY, V], rhoT)
    Z0 = Y_implied / Gval
    # --- housing: land share implied by the structures' rental and the housing value-added target
    nu_H_implied = 1 - Rst["H"] * K["H"] / hv
    if verbose:
        print(f"  land share in housing implied by K_H target: {nu_H_implied:.3f} (assumed {v['housing_land_share']:.3f}) -> using implied")
    nu_H = float(np.clip(nu_H_implied, 0.05, 0.9))
    TH = 1.0
    # --- tasks of the quantitative model: four bottleneck classes. Accessory tasks (linear add-on) are kept
    #     for the Restrepo replication only: with an energy constraint a linear accessory block makes the
    #     value of compute bounded away from zero and lets compute absorb all energy (docs/research_log.md).
    tasks = TaskSpec(names=["routine", "cognitive-L", "cognitive-H", "expert"],
                     skill=[0, 0, 1, 1], nu=nuB.tolist(), gamma=[0.0, 0.0, 0.0, 0.0],
                     alpha_inf=[1.0, 2.0, 4.0, 8.0])
    alpha0_mult = [alpha10 / 1.0, 50.0, 400.0, 3000.0]   # task 1 anchored to 2026 data; others are scenario inputs
    p = Params(tasks=tasks, sigma_B=v["sigma_B"], H=[HL, HH], sigma_N=v["sigma_N"], nu_N=nu_N, sigma_T=v["sigma_T"],
               nu_T=nu_T, nu_E=nu_E, Z0=Z0, g_Z=g, ZE=ZE, eta=eta, RE=RE, T_land=T_land, TH_land=TH, nu_H=nu_H,
               theta_H=theta_H, delta=delta, phi=phi, pQ0=pQ0, g_pQ=v["g_pQ"], pQ_floor_ratio=v["pQ_floor_ratio"],
               eps0=eps0, g_eps=v["g_eps"], eps_floor_ratio=v["eps_floor_ratio"], beta=beta, gamma_u=gamma_u,
               alpha0_mult=alpha0_mult, g_alpha=v["g_alpha"], KN0=K["N"], KQ0=KQ, KE0=K["E"], KH0=K["H"])
    p.source = {k: f"{t.status}: {t.source}" for k, t in tg.items()}
    p.source["eta"] = "calibrated: implied by K_E/GDP and energy share"
    p.source["KN0"] = "calibrated: residual so that factor payments add up to GDP (K_N/GDP target not imposed)"
    p.source["nu_H"] = "calibrated: implied by K_H/GDP, delta_H, r* and the housing value-added share"
    p.source["theta_H"] = "calibrated: housing VA share and goods consumption at the base point"
    p.source["alpha_1_2026"] = "calibrated: value of compute capital, compute Euler and task-1 tie A = w_L/alpha_1"
    p.source["compute_in_task1_implied"] = f"{c1:.4f}"
    p.source["Y_implied"] = f"{Y_implied:.4f} (final goods gross output; GDP = 1 by normalization)"
    info = dict(Y=Y_implied, C=C, theta_H=theta_H, A=A, alpha10=alpha10, KQ=KQ, pQ0=pQ0, eps0=eps0, wL=wL, wH=wH,
                nuB=nuB.tolist(), nu_N=nu_N, nu_T=nu_T, nu_E=nu_E, Z0=Z0, eta=eta, beta=beta, r=r, RQ=RQ, PE=PE)
    return p, info


def verify_calibration(p, info, tol=1e-6, verbose=True):
    """Solve the frozen-AI BGP from the calibrated stocks: the stocks must be a fixed point."""
    b = solve_bgp(p.replace(alpha_freeze=True), verbose=verbose)
    dev = {k: abs(np.log(b["K"][k] / {"N": p.KN0, "Q": p.KQ0, "E": p.KE0, "H": p.KH0}[k])) for k in b["K"]}
    dev["pQ0"] = abs(np.log(b["pQ0"] / p.pQ0))
    per = b["period"]
    out = dict(max_dev=max(dev.values()), dev=dev, GDP=per["GDP"], labor_share=per["wage_bill"] / per["GDP"],
               compute_income_share=per["RQ"] * per["KQ"] / per["GDP"], energy_share=per["PE"] * per["E"] / per["GDP"],
               land_rent_share=per["RT"] * per["T_eff"] / per["GDP"], housing_share=per["RH"] * per["HS"] / per["GDP"],
               A=per["A"], util=per["util"], wL=per["w"][0], wH=per["w"][1])
    if verbose:
        print("  calibration check:", {k: (round(float(x), 5) if not isinstance(x, dict) else x) for k, x in out.items()})
    return out, b


def save_parameters(p, info, path):
    d = p.to_dict()
    d["_calibration_info"] = {k: (v if not hasattr(v, "tolist") else v.tolist()) for k, v in info.items()}
    with open(path, "w") as f:
        json.dump(d, f, indent=2, default=float)