"""F5-conf-01: Walrasian equilibrium of a 2x2 Cobb-Douglas exchange economy.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json deterministically (no randomness).
"""
import json
from pathlib import Path
import numpy as np
import pandas as pd
from scipy.optimize import brentq

ROOT = Path(__file__).resolve().parent.parent
end = pd.read_csv(ROOT / "data" / "endowments.csv")
par = pd.read_csv(ROOT / "data" / "parameters.csv")
P = dict(zip(par["parameter"], par["value"].astype(float)))

issues = []
# --- validation -----------------------------------------------------------
assert set(end["agent"]) == {"A", "B"} and not end["agent"].duplicated().any()
assert end[["good1", "good2"]].notna().all().all()
assert (end[["good1", "good2"]] > 0).all().all(), "endowments must be strictly positive"
aA, aB = P["alpha_A"], P["alpha_B"]
assert 0 < aA < 1 and 0 < aB < 1
e = end.set_index("agent")[["good1", "good2"]].astype(float)
alpha = {"A": aA, "B": aB}
agents = ["A", "B"]

W1 = float(e["good1"].sum())
W2 = float(e["good2"].sum())
for g, W, key in [("good1", W1, "aggregate_endowment_good1"), ("good2", W2, "aggregate_endowment_good2")]:
    rec = P[key]
    if not np.isclose(rec, W):
        issues.append({
            "issue": f"parameters.csv {key} = {rec:g} disagrees with the sum of individual endowments in endowments.csv ({W:g}).",
            "handling": f"Used {W:g} (sum of primary individual endowments) so that aggregate supply is consistent with "
                        f"agents' budget constraints; summary-sheet value {rec:g} discarded (reported only as a sensitivity in report.md).",
        })

# --- demand / excess demand ------------------------------------------------
def demand(i, p):
    m = p * e.loc[i, "good1"] + e.loc[i, "good2"]
    return alpha[i] * m / p, (1 - alpha[i]) * m

def z(p, W1=W1, W2=W2):
    x1 = sum(demand(i, p)[0] for i in agents)
    x2 = sum(demand(i, p)[1] for i in agents)
    return x1 - W1, x2 - W2

# --- closed form: sum_i a_i (p e1_i + e2_i) = p W1  =>  p = sum a_i e2_i / sum (1-a_i) e1_i
p_cf = sum(alpha[i] * e.loc[i, "good2"] for i in agents) / sum((1 - alpha[i]) * e.loc[i, "good1"] for i in agents)

# --- independent numerical check: root of z1(p) via Brent on a wide bracket -
p_num = brentq(lambda p: z(p)[0], 1e-6, 1e6, xtol=1e-14, rtol=1e-15, maxiter=500)

x = {i: demand(i, p_cf) for i in agents}
z1, z2 = z(p_cf)
walras = p_cf * z1 + z2
# Walras' law holds identically: check on a grid of off-equilibrium prices too
grid = np.logspace(-2, 2, 41)
walras_grid_max = float(max(abs(p * z(p)[0] + z(p)[1]) for p in grid))

# sensitivity: price that would clear good 1 if the summary-sheet aggregate were used
sens = {}
if not np.isclose(P["aggregate_endowment_good1"], W1):
    Ws = P["aggregate_endowment_good1"]
    ps = brentq(lambda p: z(p, W1=Ws)[0], 1e-6, 1e6, xtol=1e-14)
    zs1, zs2 = z(ps, W1=Ws)
    sens = {"sensitivity_price_if_summary_aggregate_good1": ps,
            "sensitivity_walras_residual_if_summary_aggregate_good1": ps * zs1 + zs2}

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": x["A"][0], "x2_A": x["A"][1], "x1_B": x["B"][0], "x2_B": x["B"][1],
    "excess_demand_good1": z1,
    "excess_demand_good2": z2,
    "walras_law_residual": walras,
    "walras_law_residual_max_abs_offeq_grid": walras_grid_max,
    "aggregate_endowment_good1": W1,
    "aggregate_endowment_good2": W2,
    "price_ratio_numerical": p_num,
    "abs_diff_closed_form_vs_numerical": abs(p_cf - p_num),
    "income_A": p_cf * e.loc["A", "good1"] + e.loc["A", "good2"],
    "income_B": p_cf * e.loc["B", "good1"] + e.loc["B", "good2"],
    **sens,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"endowments.csv": int(len(end)), "parameters.csv": int(len(par))},
    "n_analysed": len(agents),
    "random_seed": None,
}
res = {k: (float(v) if isinstance(v, (np.floating, np.integer)) and not isinstance(v, bool) else v) for k, v in res.items()}
(ROOT / "out" / "results.json").write_text(json.dumps(res, indent=2) + "\n")
print(f"p closed form = {p_cf:.12g}; p numerical = {p_num:.12g}; z1 = {z1:.3g}; Walras residual = {walras:.3g}")