"""F5-conf-01: Walrasian equilibrium of a 2x2 Cobb-Douglas exchange economy.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads data/ only; deterministic (no randomness)."""
import json, os
import 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"].astype(float)))

issues = []
# --- validation / exclusions ---
assert end["agent"].is_unique and set(end["agent"]) == {"A", "B"}
valid = end[(end.good1 > 0) & (end.good2 > 0)]
agents = list(valid["agent"])
alpha = {i: P[f"alpha_{i}"] for i in agents}
assert all(0 < a < 1 for a in alpha.values())
e1 = dict(zip(valid.agent, valid.good1.astype(float)))
e2 = dict(zip(valid.agent, valid.good2.astype(float)))

E1 = sum(e1.values()); E2 = sum(e2.values())
rec1 = P["aggregate_endowment_good1"]; rec2 = P["aggregate_endowment_good2"]
if abs(rec1 - E1) > 1e-12:
    issues.append({"issue": f"parameters.csv aggregate_endowment_good1={rec1} differs from sum of individual endowments in endowments.csv ({E1}).",
                   "handling": f"Used {E1} (sum of primary individual endowments) in market clearing; the recorded {rec1} is inconsistent with agents' resources and was not used. Price implied by {rec1} reported as diagnostic only."})
if abs(rec2 - E2) > 1e-12:
    issues.append({"issue": f"parameters.csv aggregate_endowment_good2={rec2} differs from sum of individual endowments ({E2}).",
                   "handling": f"Used {E2}."})

def demand(p, i):
    m = p * e1[i] + e2[i]
    return alpha[i] * m / p, (1 - alpha[i]) * m

def z(p, E1_=None):
    E1_ = E1 if E1_ is None else E1_
    z1 = sum(demand(p, i)[0] for i in agents) - E1_
    z2 = sum(demand(p, i)[1] for i in agents) - E2
    return z1, z2

# --- 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] * e2[i] for i in agents) / sum((1 - alpha[i]) * e1[i] for i in agents)
alloc = {i: demand(p_cf, i) for i in agents}
z1, z2 = z(p_cf)

# --- independent numerical root of z1(p) (does not use closed form) ---
p_num = brentq(lambda p: z(p)[0], 1e-8, 1e8, xtol=1e-15, rtol=1e-15, maxiter=500)

# diagnostic: root if the (inconsistent) recorded aggregate were used
p_diag = brentq(lambda p: z(p, rec1)[0], 1e-8, 1e8, xtol=1e-15, rtol=1e-15, maxiter=500)

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": alloc["A"][0], "x2_A": alloc["A"][1],
    "x1_B": alloc["B"][0], "x2_B": alloc["B"][1],
    "excess_demand_good1": z1,
    "excess_demand_good2": z2,
    "walras_law_residual": p_cf * z1 + z2,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "price_ratio_numerical": p_num,
    "abs_diff_closed_form_vs_numerical": abs(p_cf - p_num),
    "diagnostic_price_if_recorded_aggregate_used": p_diag,
    "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,
}
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))