"""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, 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 ---
assert set(end["agent"]) == {"A", "B"} and len(end) == 2
E = end.set_index("agent")[["good1", "good2"]].astype(float)
if (E.values <= 0).any():
    issues.append({"issue": "non-positive endowment", "handling": "closed form requires positive wealth; flagged"})
aA, aB = P["alpha_A"], P["alpha_B"]
for k, a in (("alpha_A", aA), ("alpha_B", aB)):
    if not (0 < a < 1):
        issues.append({"issue": f"{k} outside (0,1)", "handling": "flagged"})
sum1, sum2 = E["good1"].sum(), E["good2"].sum()
sheet1, sheet2 = P["aggregate_endowment_good1"], P["aggregate_endowment_good2"]
if not np.isclose(sum1, sheet1):
    issues.append({"issue": f"parameters.csv aggregate_endowment_good1 = {sheet1:g} but individual endowments in endowments.csv sum to {sum1:g} (A {E.loc['A','good1']:g} + B {E.loc['B','good1']:g}).",
                   "handling": f"Used the sum of individual endowments ({sum1:g}) as the aggregate, because in an exchange economy aggregate supply is by definition the sum of agents' endowments and endowments.csv is designated primary data; the summary-sheet value {sheet1:g} is treated as a transcription error. Using {sheet1:g} would put {sheet1-sum1:g} units of good 1 that no agent owns into the market-clearing condition (inconsistent with Walras' law); a sensitivity check is reported in report.md."})
if not np.isclose(sum2, sheet2):
    issues.append({"issue": f"aggregate_endowment_good2 sheet {sheet2:g} != sum {sum2:g}", "handling": "used sum of individual endowments"})

e1A, e2A, e1B, e2B = E.loc["A", "good1"], E.loc["A", "good2"], E.loc["B", "good1"], E.loc["B", "good2"]
W1, W2 = sum1, sum2

def demands(p):
    mA, mB = p * e1A + e2A, p * e1B + e2B
    return aA * mA / p, (1 - aA) * mA, aB * mB / p, (1 - aB) * mB

def z(p, W1=W1, W2=W2):
    x1A, x2A, x1B, x2B = demands(p)
    return x1A + x1B - W1, x2A + x2B - W2

# --- closed form: p*W1 = aA*(p e1A+e2A) + aB*(p e1B+e2B) ---
p_cf = (aA * e2A + aB * e2B) / ((1 - aA) * e1A + (1 - aB) * e1B)
x1A, x2A, x1B, x2B = demands(p_cf)
z1, z2 = z(p_cf)
walras = p_cf * z1 + z2

# --- independent numerical root of z1(p) ---
p_num = brentq(lambda p: z(p)[0], 1e-6, 1e6, xtol=1e-14, rtol=1e-15, maxiter=500)

# Walras' law holds identically: check on a grid
grid = np.logspace(-3, 3, 61)
walras_grid_max = float(max(abs(p * z(p)[0] + z(p)[1]) for p in grid))

# sensitivity: summary-sheet aggregate (inconsistent)
p_sheet = brentq(lambda p: z(p, sheet1, sheet2)[0], 1e-6, 1e6, xtol=1e-14)
z_sheet = z(p_sheet, sheet1, sheet2)
sens = {"p_root_using_sheet_aggregate": p_sheet,
        "z2_at_that_price_using_sheet_aggregate": z_sheet[1],
        "walras_residual_using_sheet_aggregate": p_sheet * z_sheet[0] + z_sheet[1],
        "walras_grid_max_abs_residual": walras_grid_max,
        "abs_diff_closed_form_vs_numerical": abs(p_cf - p_num),
        "excess_demand_good2": z2}

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": x1A, "x2_A": x2A, "x1_B": x1B, "x2_B": x2B,
    "excess_demand_good1": z1,
    "walras_law_residual": walras,
    "aggregate_endowment_good1": W1,
    "aggregate_endowment_good2": W2,
    "price_ratio_numerical": p_num,
    "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)
json.dump(res, open("out/results.json", "w"), indent=2)
json.dump(sens, open("out/diagnostics.json", "w"), indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(sens, indent=2))