"""
check_derivations.py -- Symbolic verification of the theory in Section 2 of the note.

  1. Cobb-Douglas: Lagrangian -> FOCs -> demand x1 = a1*Y/p1, and HD0.
  2. Stone-Geary (LES): demand x1 = g1 + (a1/p1)(Y - sum_j p_j g_j), and HD0.
  3. Elasticities of the LES demand and the Euler identity sum_j eps_1j + eta_1 = 0,
     symbolically and at a numerical point.
  4. The money-illusion example: 2^0.1 - 1.

Writes output/derivations.json and output/macros_theory.tex (LaTeX macros used
by the note, so that no number in the note is typed by hand).

Run from the project root:  python python/check_derivations.py
"""
import json
from pathlib import Path

import sympy as sp

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "output"
OUT.mkdir(exist_ok=True)

J = 3                                     # number of goods (the proofs are generic in J)
x = sp.symbols(f"x1:{J+1}", positive=True)
p = sp.symbols(f"p1:{J+1}", positive=True)
a = sp.symbols(f"alpha1:{J+1}", positive=True)
g = sp.symbols(f"gamma1:{J+1}", real=True)
Y, lam, mu = sp.symbols("Y lambda mu", positive=True)
norm = {a[-1]: 1 - sum(a[:-1])}           # sum_j alpha_j = 1
scale = {**{p[j]: lam * p[j] for j in range(J)}, Y: lam * Y}


def solve_lagrangian(U):
    """max U(x) s.t. p.x = Y  ->  dict of Marshallian demands and multiplier."""
    L = U - mu * (sum(p[j] * x[j] for j in range(J)) - Y)
    focs = [sp.diff(L, xj) for xj in x] + [sp.diff(L, mu)]
    return sp.solve(focs, list(x) + [mu], dict=True)[0]


res = {}

# ---------------------------------------------------------------- 1. Cobb-Douglas
sol = solve_lagrangian(sum(a[j] * sp.log(x[j]) for j in range(J)))
x1_cd = sp.simplify(sol[x[0]].subs(norm))
mu_cd = sp.simplify(sol[mu].subs(norm))
assert sp.simplify(x1_cd - a[0] * Y / p[0]) == 0
assert sp.simplify(mu_cd - 1 / Y) == 0
assert sp.simplify(x1_cd.subs(scale, simultaneous=True) - x1_cd) == 0
eps_cd = [sp.simplify(sp.diff(x1_cd, p[j]) * p[j] / x1_cd) for j in range(J)]
eta_cd = sp.simplify(sp.diff(x1_cd, Y) * Y / x1_cd)
assert eps_cd == [-1, 0, 0] and eta_cd == 1
res["cobb_douglas"] = dict(x1=str(x1_cd), mu=str(mu_cd), hd0=True,
                           elasticities=[str(e) for e in eps_cd], eta=str(eta_cd))

# ---------------------------------------------------------------- 2. Stone-Geary
sol = solve_lagrangian(sum(a[j] * sp.log(x[j] - g[j]) for j in range(J)))
x1_sg = sp.simplify(sol[x[0]].subs(norm))
les = g[0] + a[0] / p[0] * (Y - sum(p[j] * g[j] for j in range(J)))
assert sp.simplify(x1_sg - les.subs(norm)) == 0
assert sp.simplify(x1_sg.subs(scale, simultaneous=True) - x1_sg) == 0
assert sp.simplify(sol[mu].subs(norm) - 1 / (Y - sum(p[j] * g[j] for j in range(J)))) == 0

# ---------------------------------------------------------------- 3. Elasticities, Euler
x1s = les  # use the closed form
eps = [sp.simplify(sp.diff(x1s, p[j]) * p[j] / x1s) for j in range(J)]
eta = sp.simplify(sp.diff(x1s, Y) * Y / x1s)
euler = sp.simplify(sum(eps) + eta)
assert euler == 0
# closed forms reported in the note
assert sp.simplify(eta - a[0] * Y / (p[0] * x1s)) == 0
assert sp.simplify(eps[0] - (-1 + (1 - a[0]) * g[0] / x1s)) == 0
for j in range(1, J):
    assert sp.simplify(eps[j] - (-a[0] * g[j] * p[j] / (p[0] * x1s))) == 0

point = {a[0]: sp.Rational(3, 10), a[1]: sp.Rational(3, 10),          # alpha3 = 0.4
         g[0]: 2, g[1]: 3, g[2]: 1,
         p[0]: 2, p[1]: 1, p[2]: sp.Rational(3, 2), Y: 20}
x1_val = x1s.subs(norm).subs(point)
eps_val = [e.subs(norm).subs(point) for e in eps]
eta_val = eta.subs(norm).subs(point)
assert sp.simplify(sum(eps_val) + eta_val) == 0                      # exact rational arithmetic
res["stone_geary"] = dict(x1=str(les), hd0=True, euler_sum=str(euler),
                          point={str(k): str(v) for k, v in point.items()},
                          x1_at_point=float(x1_val),
                          eps_at_point=[float(v) for v in eps_val],
                          eta_at_point=float(eta_val),
                          sum_at_point_exact=str(sp.simplify(sum(eps_val) + eta_val)))

# ---------------------------------------------------------------- 4. Money illusion
res["money_illusion_pct"] = float(100 * (2 ** 0.1 - 1))

(OUT / "derivations.json").write_text(json.dumps(res, indent=2))

f = lambda v, d=2: f"{v:.{d}f}"
macros = {
    "illusionPct": f(res["money_illusion_pct"], 1),
    "sgXone": f(float(x1_val), 3),
    "sgEpsOneOne": f(float(eps_val[0]), 3),
    "sgEpsOneTwo": f(float(eps_val[1]), 3),
    "sgEpsOneThree": f(float(eps_val[2]), 3),
    "sgEta": f(float(eta_val), 3),
}
(OUT / "macros_theory.tex").write_text(
    "% generated by python/check_derivations.py -- do not edit\n"
    + "".join(f"\\newcommand{{\\{k}}}{{{v}}}\n" for k, v in macros.items()))
(OUT / "macros_theory.json").write_text(json.dumps(macros, indent=2))

print("All symbolic checks passed.")
print(json.dumps(res, indent=2))