"""Competitive equilibrium of a 2-agent, 2-good Cobb-Douglas exchange economy (F5-conf-01).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Deterministic; no randomness used."""
import json, 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 set(end["agent"]) == {"A", "B"} and len(end) == 2, "need exactly agents A and B"
assert not end.isna().any().any() and not end["agent"].duplicated().any()
assert (end[["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())
e = {r.agent: (float(r.good1), float(r.good2)) for r in end.itertuples()}
agents = ["A", "B"]

E1 = sum(e[i][0] for i in agents)
E2 = sum(e[i][1] for i in agents)
for g, Eg, key in [(1, E1, "aggregate_endowment_good1"), (2, E2, "aggregate_endowment_good2")]:
    rec = P[key]
    if abs(rec - Eg) > 1e-12:
        issues.append({
            "issue": f"parameters.csv {key}={rec} does not equal the sum of individual endowments in endowments.csv ({Eg}).",
            "handling": f"Used the sum of individual endowments ({Eg}) because endowments.csv is the primary data and "
                        f"market clearing/Walras' law require aggregates consistent with agents' budgets; "
                        f"the summary-sheet value {rec} was not used in the equilibrium computation."})

# --- demands and excess demand ---
def demand(i, p):
    m = p * e[i][0] + e[i][1]
    return alpha[i] * m / p, (1 - alpha[i]) * m

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

# --- closed form: sum_i alpha_i (p e1_i + e2_i) = p E1 ---
num = sum(alpha[i] * e[i][1] for i in agents)
den = E1 - sum(alpha[i] * e[i][0] for i in agents)
assert den > 0
p_cf = num / den

# --- independent numerical root of z1(p) = 0 (bracketing; z1 strictly decreasing for CD) ---
z1 = lambda p: z(p)[0]
lo, hi = 1e-6, 1.0
while z1(hi) > 0:
    hi *= 2.0
p_num = brentq(z1, lo, hi, xtol=1e-14, rtol=1e-14, maxiter=500)

xA, xB = demand("A", p_cf), demand("B", p_cf)
z1_cf, z2_cf = z(p_cf)
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_cf,
    "walras_law_residual": p_cf * z1_cf + z2_cf,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "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": len(agents),
    "random_seed": None,
}
# diagnostic only (not a results key): price if summary-sheet E1 were used
rec1 = P["aggregate_endowment_good1"]
p_alt = num / (rec1 - sum(alpha[i] * e[i][0] for i in agents))
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))
print(f"[diagnostic] |p_cf - p_num| = {abs(p_cf - p_num):.3e}; z2 at p_cf = {z2_cf:.3e}; "
      f"price if summary-sheet E1={rec1} were used: {p_alt:.6f} (not used)")