"""High-level driver: calibrate a scenario, solve its perfect-foresight transition, save results.

Strategy (docs/replication.md): (1) calibrate the 2026 base point consistently with the scenario's
anticipated technology paths; (2) homotopy in the size of the technology shock on a short horizon;
(3) horizon continuation up to the full horizon; (4) store paths, residual diagnostics and a horizon check.
"""
from __future__ import annotations

import json
import os
import time

import numpy as np

from .calibration import calibrate
from .transition import KINDS, TwoClassModel

SCENARIO_RATE_KEYS = ("g_pQ", "g_eps", "g_alpha")


def _params_for(overrides, lam, T, param_overrides=None):
    ov = dict(overrides or {})
    base = {k: ov.get(k, None) for k in SCENARIO_RATE_KEYS}
    p, info = calibrate(overrides=ov)
    # scale the technology rates by lambda (calibration stays consistent with the scaled rates)
    scaled = {k: lam * (base[k] if base[k] is not None else getattr(p, k)) for k in SCENARIO_RATE_KEYS}
    p, info = calibrate(overrides={**ov, **scaled})
    p = p.replace(T=T, **(param_overrides or {}))
    if lam < 1.0 and p.alpha_shock_year is not None:
        # for a jump scenario, homotopy scales the log distance of the jump instead
        p = p.replace(alpha_shock_year=None, g_alpha=lam * 0.5)
    return p, info


def extend_guess(u_prev, T_prev, T_new, nv=6):
    U = u_prev.reshape(T_prev, nv)
    extra = T_new - T_prev
    Ul = np.vstack([U, np.tile(U[-1], (extra, 1))])
    g = U[-1, :5] - U[-2, :5]
    for j in range(extra):
        Ul[T_prev + j, :5] = U[-1, :5] + g * (j + 1)
    return Ul.ravel()


def solve_scenario(name, overrides=None, param_overrides=None, T=275, T_short=None,
                   lambdas=(1.0,), fallback_lambdas=(0.0, 0.25, 0.5, 0.75, 1.0), horizons=(),
                   tol=1e-9, verbose=True, outdir=None, horizon_check_years=50):
    t0 = time.time()
    log = []
    u, T_prev = None, None
    T_short = T if T_short is None else T_short      # homotopy on the full horizon by default (terminal near BGP)
    # --- homotopy
    for lam in lambdas:
        p, info = _params_for(overrides, lam, T_short, param_overrides)
        m = TwoClassModel(p)
        if u is None:
            u = m.smart_guess()
        if verbose:
            print(f"[{name}] homotopy lambda={lam:.2f} T={T_short}", flush=True)
        u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)
        log.append(dict(stage="homotopy", lam=lam, T=T_short, converged=bool(ok), res=float(hist[-1])))
    if not ok and fallback_lambdas and tuple(lambdas) != tuple(fallback_lambdas):
        # direct solve failed: homotopy in the shock size from the frozen economy
        if verbose:
            print(f"[{name}] direct solve not converged (|res|={hist[-1]:.2e}); switching to homotopy", flush=True)
        u = None
        for lam in fallback_lambdas:
            p, info = _params_for(overrides, lam, T_short, param_overrides)
            m = TwoClassModel(p)
            if u is None:
                u = m.smart_guess()
            if verbose:
                print(f"[{name}] homotopy lambda={lam:.2f} T={T_short}", flush=True)
            u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)
            log.append(dict(stage="homotopy-fallback", lam=lam, T=T_short, converged=bool(ok), res=float(hist[-1])))
    T_prev = T_short
    # --- horizon continuation
    for Th in [h for h in horizons if T_short < h < T] + [T]:
        p, info = _params_for(overrides, 1.0, Th, param_overrides)
        m = TwoClassModel(p)
        u = extend_guess(u, T_prev, Th)
        if verbose:
            print(f"[{name}] horizon {Th}", flush=True)
        u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)
        log.append(dict(stage="horizon", lam=1.0, T=Th, converged=bool(ok), res=float(hist[-1])))
        T_prev = Th
    P = m.evaluate(u, full=True)
    # --- horizon (truncation) check
    hc = None
    if horizon_check_years:
        pl, _ = _params_for(overrides, 1.0, T + horizon_check_years, param_overrides)
        ml = TwoClassModel(pl)
        ul, okl, histl = ml.solve(extend_guess(u, T, T + horizon_check_years), tol=tol, verbose=False, maxit=40)
        Kl, rl, _ = ml.unpack(ul)
        K0, r0, _ = m.unpack(u)
        n = T - horizon_check_years
        hc = dict(converged=bool(okl), res=float(histl[-1]),
                  max_log_dev_K={k: float(np.max(np.abs(np.log(Kl[k][1:n + 1] / K0[k][1:n + 1])))) for k in KINDS},
                  max_dev_r=float(np.max(np.abs(rl[1:n + 1] - r0[1:n + 1]))), years_compared=int(n))
    out = dict(name=name, overrides=overrides or {}, param_overrides=param_overrides or {}, T=T,
               log=log, horizon_check=hc, seconds=time.time() - t0, converged=bool(ok), final_res=float(hist[-1]),
               calibration_info={k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in info.items()})
    if outdir:
        os.makedirs(outdir, exist_ok=True)
        np.save(os.path.join(outdir, f"{name}_u.npy"), u)
        p.to_json(os.path.join(outdir, f"{name}_params.json"))
        save_paths(P, os.path.join(outdir, f"{name}_paths.npz"))
        with open(os.path.join(outdir, f"{name}_meta.json"), "w") as f:
            json.dump(out, f, indent=2, default=float)
    return p, m, u, P, out


PATH_KEYS = ("years", "r", "Y", "GDP", "wage", "wL", "wH", "A", "RT", "RTH", "RTp", "RTHp", "RH", "rentR", "PE", "util",
             "Qtot", "E", "EQ", "C_total", "C_K", "C_W", "X_K", "X_W", "P_idx", "transfers", "tax", "fund", "PT", "PTH",
             "W", "inv_goods", "foreign_flow", "labor_share", "income_residual", "Li", "Qi", "Xi", "HS", "alpha", "pQ",
             "eps", "Z", "gf", "T_eff", "TH_eff", "residuals", "goods_gap")


def save_paths(P, path):
    d = {}
    for k in PATH_KEYS:
        d[k] = np.asarray(P[k])
    for k in KINDS:
        d[f"K_{k}"] = P["K"][k]
        d[f"x_{k}"] = P["x"][k]
        d[f"I_{k}"] = P["I"][k]
        d[f"q_{k}"] = P["q"][k]
        d[f"R_{k}"] = P["R"][k]
        d[f"V_{k}"] = P["V"][k]
        d[f"D_{k}"] = P["D"][k]
        d[f"adj_{k}"] = P["adj"][k]
    np.savez_compressed(path, **d)


def load_paths(path):
    z = np.load(path)
    return {k: z[k] for k in z.files}