"""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 pandas as pd
from scipy.optimize import brentq

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
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 sorted(end["agent"]) == ["A", "B"], "expected exactly agents A and B"
e = {r.agent: (float(r.good1), float(r.good2)) for r in end.itertuples()}
assert all(v > 0 for t in e.values() for v in t), "endowments must be strictly positive"
a = {"A": P["alpha_A"], "B": P["alpha_B"]}
assert all(0 < v < 1 for v in a.values()), "alpha must lie in (0,1)"

E1 = e["A"][0] + e["B"][0]
E2 = e["A"][1] + e["B"][1]
E1_sheet = P.get("aggregate_endowment_good1")
E2_sheet = P.get("aggregate_endowment_good2")
for g, Esum, Esh in [(1, E1, E1_sheet), (2, E2, E2_sheet)]:
    if Esh is not None and abs(Esum - Esh) > 1e-9:
        issues.append({
            "issue": (f"parameters.csv aggregate_endowment_good{g} = {Esh:g} is inconsistent with the "
                      f"sum of individual endowments in endowments.csv ({Esum:g})."),
            "handling": (f"Used the sum of individual endowments ({Esum:g}) as the aggregate endowment, because "
                         "endowments.csv is the primary data and in an exchange economy the aggregate endowment "
                         "is by definition the sum of individual endowments (otherwise Walras' law plus budget "
                         "balance would make market clearing impossible). The summary-sheet value was not used.")})

# --- demands ---
def demand(i, p):
    w = p * e[i][0] + e[i][1]
    return a[i] * w / p, (1 - a[i]) * w

def z1(p, E1agg=E1):
    return sum(demand(i, p)[0] for i in "AB") - E1agg

def z2(p, E2agg=E2):
    return sum(demand(i, p)[1] for i in "AB") - E2agg

# --- closed form: p = sum_i alpha_i e2_i / sum_i (1-alpha_i) e1_i ---
p_cf = sum(a[i] * e[i][1] for i in "AB") / sum((1 - a[i]) * e[i][0] for i in "AB")

# --- independent numerical root of z1(p)=0 (bracketing, no closed form used) ---
lo, hi = 1e-6, 1.0
while z1(hi) > 0:
    hi *= 2
p_num = brentq(z1, lo, hi, xtol=1e-14, rtol=1e-15, maxiter=500)

alloc = {i: demand(i, p_cf) for i in "AB"}
ez1 = z1(p_cf); ez2 = z2(p_cf)
walras = p_cf * ez1 + ez2

# Sensitivity: what if the summary-sheet aggregate were (wrongly) imposed for good 1?
sens = None
if E1_sheet is not None and abs(E1_sheet - E1) > 1e-9:
    f = lambda p: z1(p, E1_sheet)
    hi2 = 1.0
    while f(hi2) > 0:
        hi2 *= 2
    p_bad = brentq(f, 1e-6, hi2, xtol=1e-14)
    sens = {"p_if_sheet_E1_imposed": p_bad,
            "z2_at_that_p_using_true_E2": z2(p_bad),
            "note": "Good-2 market does not clear: the sheet value is inconsistent with budget-balanced demands."}

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p_cf,
    "x1_A": alloc["A"][0], "x2_A": alloc["A"][1],
    "x1_B": alloc["B"][0], "x2_B": alloc["B"][1],
    "excess_demand_good1": ez1,
    "excess_demand_good2": ez2,
    "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),
    "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,
    "sensitivity_sheet_aggregate": sens,
}
with open(os.path.join(HERE, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2))