"""F5-conf-01: Walrasian equilibrium of a 2x2 Cobb-Douglas exchange economy.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
No randomness is used (random_seed = null).
"""
import json, os
import numpy as np
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")

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

# --- audit / validation ---
assert sorted(endw["agent"]) == ["A", "B"] and not endw["agent"].duplicated().any()
assert not endw.isna().any().any() and not pars["value"].isna().any()
assert (endw[["good1", "good2"]] > 0).all().all(), "endowments must be strictly positive"
alpha = {"A": P["alpha_A"], "B": P["alpha_B"]}
assert all(0 < a < 1 for a in alpha.values())
e1 = dict(zip(endw["agent"], endw["good1"].astype(float)))
e2 = dict(zip(endw["agent"], endw["good2"].astype(float)))
agents = ["A", "B"]
E1 = sum(e1.values()); E2 = sum(e2.values())
for E, key in [(E1, "aggregate_endowment_good1"), (E2, "aggregate_endowment_good2")]:
    if not np.isclose(E, P[key]):
        issues.append({
            "issue": f"parameters.csv {key}={P[key]:g} disagrees with the sum of individual "
                     f"endowments in endowments.csv ({E:g}).",
            "handling": f"Used the sum of individual endowments ({E:g}) because endowments.csv is the "
                        "primary data and market clearing must be consistent with the agents' budget "
                        "constraints (using the summary value would make the good-1 clearing condition "
                        "inconsistent with the budget constraints, so Walras' law would fail). Summary value not used in any computation except a reported sensitivity."})

# --- model ---
def demand(i, p):
    w = p * e1[i] + e2[i]
    return alpha[i] * w / p, (1 - alpha[i]) * w

def z(p):
    x1 = sum(demand(i, p)[0] for i in agents); x2 = sum(demand(i, p)[1] for i in agents)
    return x1 - E1, x2 - E2

# baseline / independent check: bracketing root of z1(p) (closed form not used)
lo, hi = 1e-6, 1.0
while z(hi)[0] > 0:
    hi *= 2
p_num = brentq(lambda q: z(q)[0], lo, hi, xtol=1e-14, rtol=1e-15, maxiter=500)

# closed form: p = sum alpha_i e2_i / sum (1-alpha_i) e1_i
p = sum(alpha[i] * e2[i] for i in agents) / sum((1 - alpha[i]) * e1[i] for i in agents)

z1, z2 = z(p)
alloc = {f"x{g}_{i}": float(demand(i, p)[g - 1]) for i in agents for g in (1, 2)}

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": float(p),
    **{k: alloc[k] for k in ["x1_A", "x2_A", "x1_B", "x2_B"]},
    "excess_demand_good1": float(z1),
    "excess_demand_good2": float(z2),
    "walras_law_residual": float(p * z1 + z2),
    "aggregate_endowment_good1": float(E1),
    "aggregate_endowment_good2": float(E2),
    "price_ratio_numerical": float(p_num),
    "abs_diff_closed_vs_numerical": float(abs(p - p_num)),
    "excess_demand_good1_at_numerical_price": float(z(p_num)[0]),
    "wealth_A": float(p * e1["A"] + e2["A"]),
    "wealth_B": float(p * e1["B"] + e2["B"]),
    "summary_sheet_aggregate_endowment_good1": float(P["aggregate_endowment_good1"]),
    "sensitivity_excess_demand_good1_vs_summary_E1": float(sum(alloc[f"x1_{i}"] for i in agents) - P["aggregate_endowment_good1"]),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": n_rows,
    "n_analysed": len(agents),
    "random_seed": None,
}
with open(os.path.join(O, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["price_ratio_p1_over_p2", "price_ratio_numerical",
                  "excess_demand_good1", "walras_law_residual"]}, indent=2))