"""Walrasian equilibrium of a 2-agent, 2-good Cobb-Douglas exchange economy (task F5-conf-01).
Reads only data/; writes out/results.json deterministically. No randomness is 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 end["agent"].is_unique and set(end["agent"]) == {"A", "B"}, "expected exactly agents A and B"
assert not end.isna().any().any() and not par["value"].isna().any()
e = {r.agent: (float(r.good1), float(r.good2)) for r in end.itertuples()}
alpha = {"A": P["alpha_A"], "B": P["alpha_B"]}
for i in "AB":
    assert 0 < alpha[i] < 1, f"alpha_{i} outside (0,1)"
    assert e[i][0] > 0 and e[i][1] > 0, f"non-positive endowment for {i}"

E1 = sum(e[i][0] for i in "AB")
E2 = sum(e[i][1] for i in "AB")
E1_sheet, E2_sheet = P["aggregate_endowment_good1"], P["aggregate_endowment_good2"]


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


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


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


# diagnostic for the summary-sheet discrepancy (not used for estimates)
if abs(E1_sheet - E1) > 1e-12:
    p_sheet = brentq(lambda p: z1(p, E1_sheet), 1e-9, 1e9, xtol=1e-15, rtol=1e-15, maxiter=500)
    issues.append({
        "issue": f"parameters.csv aggregate_endowment_good1={E1_sheet} differs from the sum of individual endowments "
                 f"in endowments.csv ({E1}).",
        "handling": f"Used the sum of individual endowments ({E1}) since aggregate endowment in a pure exchange economy is "
                    f"by definition the sum of individual endowments and endowments.csv is the primary data. Diagnostic: "
                    f"clearing good 1 against {E1_sheet} would give p={p_sheet:.6f} with good-2 excess demand "
                    f"{z2(p_sheet):.6f} != 0, i.e. no consistent equilibrium; the summary value is treated as a recording error."})
if abs(E2_sheet - E2) > 1e-12:
    issues.append({"issue": f"aggregate_endowment_good2={E2_sheet} differs from sum {E2}",
                   "handling": "Used sum of individual endowments."})

# ---- baseline / independent numerical check: root of z1(p) -----------------
p_num = brentq(z1, 1e-9, 1e9, xtol=1e-15, rtol=1e-15, maxiter=500)

# ---- closed form -------------------------------------------------------------
# sum_i alpha_i (p e1_i + e2_i) = p E1  =>  p = sum alpha_i e2_i / sum (1-alpha_i) e1_i
p = sum(alpha[i] * e[i][1] for i in "AB") / sum((1 - alpha[i]) * e[i][0] for i in "AB")
x = {i: demand(i, p) for i in "AB"}
ed1 = z1(p)
walras = p * z1(p) + z2(p)

res = {
    "task_id": "F5-conf-01",
    "price_ratio_p1_over_p2": p,
    "x1_A": x["A"][0], "x2_A": x["A"][1], "x1_B": x["B"][0], "x2_B": x["B"][1],
    "excess_demand_good1": ed1,
    "excess_demand_good2": z2(p),
    "walras_law_residual": walras,
    "aggregate_endowment_good1": E1,
    "aggregate_endowment_good2": E2,
    "aggregate_endowment_good1_summary_sheet": E1_sheet,
    "aggregate_endowment_good2_summary_sheet": E2_sheet,
    "price_ratio_numerical": p_num,
    "price_abs_diff_closed_vs_numerical": abs(p - 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(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "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=1))