"""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 and out/exclusions.csv. Deterministic (no randomness)."""
import json, hashlib, os
import numpy as np
import pandas as pd
from scipy.optimize import brentq

DATA, OUT = "data", "out"
end = pd.read_csv(os.path.join(DATA, "endowments.csv"))
par = pd.read_csv(os.path.join(DATA, "parameters.csv"))
n_rows = {"endowments.csv": int(len(end)), "parameters.csv": int(len(par))}
P = dict(zip(par["parameter"], par["value"].astype(float)))
issues = []

# --- validation ---
assert sorted(end["agent"]) == ["A", "B"], "need exactly agents A and B"
assert not end.isna().any().any() and not par.isna().any().any()
e = end.set_index("agent")[["good1", "good2"]].astype(float)
alpha = {"A": P["alpha_A"], "B": P["alpha_B"]}
assert all(0 < a < 1 for a in alpha.values()), "alpha must lie in (0,1)"
assert (e.values > 0).all(), "endowments must be strictly positive"

# aggregate endowments: derive from primary per-agent data, compare with summary sheet
agg1, agg2 = float(e["good1"].sum()), float(e["good2"].sum())
for g, agg in (("good1", agg1), ("good2", agg2)):
    rec = P[f"aggregate_endowment_{g}"]
    if not np.isclose(rec, agg, rtol=0, atol=1e-12):
        issues.append({
            "issue": f"parameters.csv aggregate_endowment_{g} = {rec} does not equal the sum of individual "
                     f"endowments in endowments.csv ({agg}); discrepancy {rec - agg:+g}.",
            "handling": f"Used the sum of the primary per-agent endowments ({agg}) in the market-clearing "
                        f"condition; the summary-sheet value {rec} was not used. Clearing against {rec} would be "
                        f"inconsistent with the agents' budget constraints (Walras' law would fail by p*({agg - rec:+g}))."})

# --- demands ---
def wealth(i, p): return p * e.loc[i, "good1"] + e.loc[i, "good2"]
def x1(i, p): return alpha[i] * wealth(i, p) / p
def x2(i, p): return (1 - alpha[i]) * wealth(i, p)
def z1(p): return sum(x1(i, p) for i in "AB") - agg1
def z2(p): return sum(x2(i, p) for i in "AB") - agg2

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

# --- independent numerical root of z1(p) (baseline check) ---
lo, hi = 1e-6, 1e6
assert z1(lo) > 0 > z1(hi)
p_num = brentq(z1, lo, hi, xtol=1e-14, rtol=1e-15, maxiter=500)

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": x1("A", p_cf), "x2_A": x2("A", p_cf),
    "x1_B": x1("B", p_cf), "x2_B": x2("B", p_cf),
    "excess_demand_good1": z1(p_cf),
    "excess_demand_good2": z2(p_cf),
    "walras_law_residual": p_cf * z1(p_cf) + z2(p_cf),
    "walras_law_residual_offeq_max_abs": max(abs(p * z1(p) + z2(p)) for p in np.linspace(0.1, 10, 100)),
    "aggregate_endowment_good1": agg1,
    "aggregate_endowment_good2": agg2,
    "aggregate_endowment_good1_recorded": P["aggregate_endowment_good1"],
    "aggregate_endowment_good2_recorded": P["aggregate_endowment_good2"],
    "price_ratio_numerical": p_num,
    "price_abs_diff_closed_vs_numerical": abs(p_cf - p_num),
    "alpha_A": alpha["A"], "alpha_B": alpha["B"],
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": n_rows,
    "n_analysed": int(len(e)),
    "random_seed": None,
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame([{"step": 1, "rule": "keep agents A and B with complete, strictly positive endowments",
               "n_before": len(end), "n_after": len(e), "n_removed": len(end) - len(e),
               "rationale": "no rows removed; all inputs valid"}]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
print(f"p closed-form = {p_cf:.12g}; p numerical = {p_num:.12g}; z1 = {res['excess_demand_good1']:.3g}; "
      f"Walras residual = {res['walras_law_residual']:.3g}; data issues = {len(issues)}")