"""F5-conf-01: Walrasian equilibrium of a 2x2 Cobb-Douglas exchange economy.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd
from scipy.optimize import brentq

end = pd.read_csv("data/endowments.csv")
par = pd.read_csv("data/parameters.csv")
P = dict(zip(par["parameter"], par["value"]))
issues = []

# --- validation ---
assert set(end["agent"]) == {"A", "B"} and len(end) == 2, "expect one row per agent A,B"
assert end["agent"].is_unique
if (end[["good1", "good2"]] <= 0).any().any():
    raise ValueError("non-positive endowment")
alpha = {"A": float(P["alpha_A"]), "B": float(P["alpha_B"])}
for k, a in alpha.items():
    assert 0 < a < 1, f"alpha_{k} not in (0,1)"
e = {r.agent: (float(r.good1), float(r.good2)) for r in end.itertuples()}
E1 = sum(v[0] for v in e.values()); E2 = sum(v[1] for v in e.values())
E1_sheet = float(P["aggregate_endowment_good1"]); E2_sheet = float(P["aggregate_endowment_good2"])
if not np.isclose(E1, E1_sheet):
    issues.append({"issue": f"parameters.csv aggregate_endowment_good1 = {E1_sheet:g} but the sum of individual "
                            f"endowments of good 1 in endowments.csv is {E1:g} (8 + 4).",
                   "handling": f"Used {E1:g}, the sum of the primary individual endowments; an aggregate that differs "
                               "from the sum of individual endowments is inconsistent with a pure exchange economy "
                               "(no production), so the summary-sheet value was treated as a transcription error and "
                               "not used. Market clearing against the sheet value is reported in the report only as a diagnostic."})
if not np.isclose(E2, E2_sheet):
    issues.append({"issue": f"aggregate_endowment_good2 sheet {E2_sheet:g} != sum {E2:g}",
                   "handling": f"Used {E2:g} (sum of individual endowments)."})

# --- demands (p2 = 1) ---
def demand(i, p):
    w = p * e[i][0] + e[i][1]
    return alpha[i] * w / p, (1 - alpha[i]) * w

def z(p, E1=E1, E2=E2):
    x1 = sum(demand(i, p)[0] for i in e); x2 = sum(demand(i, p)[1] for i in e)
    return x1 - E1, x2 - E2

# --- closed form: sum_i alpha_i (p e1_i + e2_i) = p E1  =>  p = sum alpha_i e2_i / sum (1-alpha_i) e1_i
p_cf = sum(alpha[i] * e[i][1] for i in e) / sum((1 - alpha[i]) * e[i][0] for i in e)

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

z1, z2 = z(p_cf)
walras = p_cf * z1 + z2
# Walras law also checked off-equilibrium on a grid
grid = np.logspace(-2, 2, 41)
walras_grid_max = float(max(abs(p * z(p)[0] + z(p)[1]) for p in grid))
# diagnostic: clearing vs sheet aggregate
z1_sheet = z(p_cf, E1=E1_sheet)[0]

xA, xB = demand("A", p_cf), demand("B", p_cf)
res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": xA[0], "x2_A": xA[1], "x1_B": xB[0], "x2_B": xB[1],
    "excess_demand_good1": z1,
    "walras_law_residual": walras,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "price_ratio_numerical": p_num,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"endowments.csv": int(len(end)), "parameters.csv": int(len(par))},
    "n_analysed": len(e),
    "random_seed": None,
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))
print("excess demand good2 at p_cf:", z2)
print("|p_cf - p_num|:", abs(p_cf - p_num))
print("max |Walras residual| on grid p in [0.01,100]:", walras_grid_max)
print("diagnostic z1 vs sheet E1=14 at p_cf:", z1_sheet)