"""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
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 of inputs ---
assert set(end["agent"]) == {"A", "B"} and len(end) == 2, "expected exactly agents A and B"
e = end.set_index("agent")[["good1", "good2"]].astype(float)
if (e.values <= 0).any():
    raise ValueError("non-positive endowment: closed form not valid")
aA, aB = P["alpha_A"], P["alpha_B"]
for k, a in (("alpha_A", aA), ("alpha_B", aB)):
    if not (0 < a < 1):
        raise ValueError(f"{k} outside (0,1)")

E1 = float(e["good1"].sum()); E2 = float(e["good2"].sum())
S1 = P["aggregate_endowment_good1"]; S2 = P["aggregate_endowment_good2"]
if not np.isclose(E1, S1):
    issues.append({
        "issue": f"parameters.csv aggregate_endowment_good1 = {S1:g} but the sum of individual "
                 f"endowments of good 1 in endowments.csv is {E1:g} (A={e.loc['A','good1']:g}, B={e.loc['B','good1']:g}).",
        "handling": f"Used the sum of individual endowments ({E1:g}) as the aggregate endowment. In an exchange economy "
                    "aggregate supply is by definition the sum of individual endowments, and endowments.csv is declared the primary "
                    "data; the summary-sheet value is treated as a transcription error. Using it would make markets impossible to clear "
                    "(excess demand for good 1 at the equilibrium price would be "
                    f"{E1 - S1:g}, not 0); this is reported as a sensitivity figure only."})
if not np.isclose(E2, S2):
    issues.append({"issue": f"aggregate_endowment_good2 = {S2:g} differs from sum of individual endowments {E2:g}.",
                   "handling": f"Used the sum of individual endowments ({E2:g})."})

alpha = {"A": aA, "B": aB}

def income(p, i):
    return p * e.loc[i, "good1"] + e.loc[i, "good2"]

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

def z(p):
    d = [demand(p, i) for i in ("A", "B")]
    return sum(x[0] for x in d) - E1, sum(x[1] for x in d) - E2

# --- closed form: p*E1 = sum_i alpha_i (p e1_i + e2_i)  =>  p = sum alpha_i e2_i / sum (1-alpha_i) e1_i
p_cf = (aA * e.loc["A", "good2"] + aB * e.loc["B", "good2"]) / \
       ((1 - aA) * e.loc["A", "good1"] + (1 - aB) * e.loc["B", "good1"])
p_cf = float(p_cf)

# --- independent numerical root of z1(p) (bracketing; z1 -> +inf as p->0, -> negative as p->inf)
p_num = brentq(lambda p: z(p)[0], 1e-8, 1e8, xtol=1e-14, rtol=1e-14, maxiter=500)

xA = demand(p_cf, "A"); xB = demand(p_cf, "B")
z1, z2 = z(p_cf)
walras = p_cf * z1 + z2
# Walras' law at arbitrary off-equilibrium prices (identity check)
grid = np.logspace(-2, 2, 41)
walras_grid_max = float(max(abs(p * z(p)[0] + z(p)[1]) for p in grid))
z1_if_summary = float(xA[0] + xB[0] - S1)

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": float(xA[0]), "x2_A": float(xA[1]),
    "x1_B": float(xB[0]), "x2_B": float(xB[1]),
    "excess_demand_good1": float(z1),
    "excess_demand_good2": float(z2),
    "walras_law_residual": float(walras),
    "walras_law_max_abs_residual_offequilibrium_grid": walras_grid_max,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "price_ratio_numerical": float(p_num),
    "abs_diff_closed_form_vs_numerical": abs(p_cf - p_num),
    "sensitivity_excess_demand_good1_vs_summary_sheet_aggregate": z1_if_summary,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"endowments.csv": int(len(end)), "parameters.csv": int(len(par))},
    "n_analysed": 2,
    "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))