"""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 used."""
import json, math, os
import pandas as pd
from scipy.optimize import brentq

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
end = pd.read_csv(os.path.join(ROOT, "data", "endowments.csv"))
par = pd.read_csv(os.path.join(ROOT, "data", "parameters.csv"))
P = dict(zip(par["parameter"], par["value"].astype(float)))

issues = []
# --- validation ---
assert not end.isna().any().any() and not par["value"].isna().any()
assert end["agent"].is_unique and set(end["agent"]) == {"A", "B"}
end = end.set_index("agent")
agents = ["A", "B"]
e1 = {a: float(end.loc[a, "good1"]) for a in agents}
e2 = {a: float(end.loc[a, "good2"]) for a in agents}
alpha = {a: P[f"alpha_{a}"] for a in agents}
valid = [a for a in agents if e1[a] > 0 and e2[a] > 0 and 0 < alpha[a] < 1]
assert valid == agents, "invalid agent parameters"

E1 = sum(e1.values()); E2 = sum(e2.values())
E1_sheet = P["aggregate_endowment_good1"]; E2_sheet = P["aggregate_endowment_good2"]
for g, s, r in [(1, E1, E1_sheet), (2, E2, E2_sheet)]:
    if not math.isclose(s, r, rel_tol=1e-12):
        issues.append({
            "issue": f"parameters.csv aggregate_endowment_good{g} = {r} but sum of individual "
                     f"endowments in endowments.csv = {s}.",
            "handling": f"Used the sum of individual endowments ({s}) as the aggregate, since budgets are "
                        f"defined by individual endowments and endowments.csv is the primary data; using {r} "
                        f"would violate Walras' law. Summary-sheet value reported only as a sensitivity in report.md."})

# --- demand and excess demand (p2 = 1) ---
def demand(a, p):
    m = p * e1[a] + e2[a]
    return alpha[a] * m / p, (1 - alpha[a]) * m

def z(p, agg1=E1, agg2=E2):
    d = [demand(a, p) for a in agents]
    return sum(x[0] for x in d) - agg1, sum(x[1] for x in d) - agg2

# --- closed form ---
p_cf = sum(alpha[a] * e2[a] for a in agents) / sum((1 - alpha[a]) * e1[a] for a in agents)
# --- independent numerical root of z1(p) (does not use closed form) ---
p_num = brentq(lambda p: z(p)[0], 1e-6, 1e6, xtol=1e-14, rtol=1e-15, maxiter=500)

xA = demand("A", p_cf); xB = demand("B", p_cf)
z1, z2 = z(p_cf)
walras = p_cf * z1 + z2

# sensitivity: price that would clear good 1 if summary-sheet aggregate were imposed
p_sheet = brentq(lambda p: z(p, E1_sheet, E2_sheet)[0], 1e-6, 1e6, xtol=1e-14, rtol=1e-15)
z_sheet = z(p_sheet, E1_sheet, E2_sheet)

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,
    "excess_demand_good2": z2,
    "walras_law_residual": walras,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "price_ratio_numerical": p_num,
    "abs_diff_closed_form_vs_numerical": abs(p_cf - p_num),
    "baseline_price_ratio_numerical": p_num,
    "sensitivity_summary_sheet_E1": E1_sheet,
    "sensitivity_price_clearing_good1_with_sheet_E1": p_sheet,
    "sensitivity_excess_demand_good2_with_sheet_E1": z_sheet[1],
    "sensitivity_walras_residual_with_sheet_E1": p_sheet * z_sheet[0] + z_sheet[1],
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"endowments.csv": int(len(end)), "parameters.csv": int(len(par))},
    "n_analysed": len(valid),
    "random_seed": None,
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))