"""F5-conf-01: Walrasian equilibrium of a 2x2 Cobb-Douglas exchange economy.
Reads only data/, writes out/results.json and out/exclusions.csv deterministically. No randomness."""
import json, os, hashlib
import pandas as pd
from scipy.optimize import brentq

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
D = os.path.join(ROOT, "data"); O = os.path.join(ROOT, "out")

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

# --- audit ---
assert sorted(end["agent"]) == ["A", "B"] and not end["agent"].duplicated().any()
assert end[["good1", "good2"]].notna().all().all() and (end[["good1", "good2"]] > 0).all().all()
for a in ("alpha_A", "alpha_B"):
    assert 0 < P[a] < 1, a
E = end.set_index("agent")[["good1", "good2"]].astype(float)
alpha = {"A": P["alpha_A"], "B": P["alpha_B"]}
sum1, sum2 = float(E["good1"].sum()), float(E["good2"].sum())
rec1, rec2 = P["aggregate_endowment_good1"], P["aggregate_endowment_good2"]
if abs(rec1 - sum1) > 1e-12:
    issues.append({"issue": f"parameters.csv aggregate_endowment_good1 = {rec1:g} but individual endowments in endowments.csv sum to {sum1:g} (8 + 4); the summary sheet is internally inconsistent with the primary data.",
                   "handling": f"Used the sum of individual endowments ({sum1:g}) as the aggregate endowment of good 1 in the market-clearing condition, because endowments.csv is the primary data and the aggregate is by definition the sum of individual endowments; the recorded value {rec1:g} is ignored (with it the market cannot clear: excess demand at the equilibrium price would be {sum1 - rec1:g})."})
if abs(rec2 - sum2) > 1e-12:
    issues.append({"issue": f"aggregate_endowment_good2 recorded {rec2:g} vs sum {sum2:g}",
                   "handling": "Used the sum of individual endowments."})
W1, W2 = sum1, sum2

# --- closed form ---
# Cobb-Douglas demand: x1_i = a_i m_i / p, x2_i = (1-a_i) m_i, m_i = p e1_i + e2_i.
# Clearing good 1: sum 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 "AB") / sum((1 - alpha[i]) * E.loc[i, "good1"] for i in "AB")

def demand(p):
    out = {}
    for i in "AB":
        m = p * E.loc[i, "good1"] + E.loc[i, "good2"]
        out[i] = (alpha[i] * m / p, (1 - alpha[i]) * m)
    return out

def z1(p):
    d = demand(p); return d["A"][0] + d["B"][0] - W1

def z2(p):
    d = demand(p); return d["A"][1] + d["B"][1] - W2

# --- independent numerical root of z1(p) (does not use the closed form) ---
p_num = brentq(z1, 1e-6, 1e6, xtol=1e-14, rtol=1e-15, maxiter=500)

d = demand(p_cf)
res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": float(p_cf),
    "x1_A": float(d["A"][0]), "x2_A": float(d["A"][1]),
    "x1_B": float(d["B"][0]), "x2_B": float(d["B"][1]),
    "excess_demand_good1": float(z1(p_cf)),
    "walras_law_residual": float(p_cf * z1(p_cf) + z2(p_cf)),
    "aggregate_endowment_good1": W1,
    "aggregate_endowment_good2": W2,
    "price_ratio_numerical": float(p_num),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": n_rows,
    "n_analysed": int(len(E)),
    "random_seed": None,
    # supplementary diagnostics
    "excess_demand_good2": float(z2(p_cf)),
    "abs_diff_closed_form_vs_numerical": float(abs(p_cf - p_num)),
    "excess_demand_good1_if_recorded_aggregate_used": float(z1(p_cf) + W1 - rec1),
    "recorded_aggregate_endowment_good1": float(rec1),
    "input_sha256": {f: hashlib.sha256(open(os.path.join(D, f), "rb").read()).hexdigest() for f in sorted(n_rows)},
}
with open(os.path.join(O, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame([{"step": 1, "rule": "agents with missing/non-positive endowments", "n_before": len(E), "n_after": len(E), "n_removed": 0,
               "rationale": "none found; both agents analysed"}]).to_csv(os.path.join(O, "exclusions.csv"), index=False)
print(json.dumps({k: res[k] for k in ["price_ratio_p1_over_p2", "price_ratio_numerical", "excess_demand_good1", "walras_law_residual"]}))