#!/usr/bin/env python3
"""Exact algebraic calculations for the accompanying portfolio-choice article.

The multidimensional calculations use three traded risks and two factors with
symbolically unrestricted covariance entries. Additional calculations check
commutator-scaled orthogonal markets, opportunity clocks, and moment boundaries. All gradients, Hessians, and
third derivatives are formal independent jets, subject only to symmetry.
No data files, network access, random samples, or absolute paths are used.
"""
from __future__ import annotations

from itertools import combinations_with_replacement
import sys

try:
    import sympy as sp
except ImportError as exc:
    raise SystemExit(
        "SymPy is required. Install supplement/requirements.txt from the package root."
    ) from exc


class Calculations:
    """Collect exact residuals and raise an error for any nonzero expression."""

    def __init__(self) -> None:
        self.count = 0

    def zero(self, name: str, residual: sp.Expr | sp.MatrixBase) -> None:
        entries = list(residual) if isinstance(residual, sp.MatrixBase) else [residual]
        for position, entry in enumerate(entries):
            numerator = sp.together(entry).as_numer_denom()[0]
            result = sp.expand(numerator)
            if result != 0:
                result = sp.simplify(result)
            if result != 0:
                raise AssertionError(f"{name}, entry {position}: {result}")
        self.count += 1
        print(f"{self.count:02d}  {name}: zero residual", flush=True)


def multidimensional(check: Calculations) -> None:
    """Check the marginal, gradient, divergence, and density formulas."""
    n, d = 3, 2
    m = d + 1
    p = sp.Symbol("p", nonzero=True)
    s = sp.Matrix(sp.symbols("s1:3"))
    grad = sp.Matrix([p, *s])
    C = sp.Matrix(d, n, lambda i, j: sp.Symbol(f"c{i+1}{j+1}"))
    R = sp.eye(d) - C * C.T
    lam = sp.Matrix(sp.symbols("l1:4"))
    beta = sp.Matrix(sp.symbols("b1:3"))
    lam_y = sp.Matrix(n, d, lambda i, j: sp.Symbol(f"l{i+1}_{j+1}"))
    beta_y = sp.Matrix(d, d, lambda i, j: sp.Symbol(f"b{i+1}_{j+1}"))
    lam_yy = {
        (a, i, j): sp.Symbol(f"l{a+1}_{i+1}{j+1}")
        for a in range(n) for i, j in combinations_with_replacement(range(d), 2)
    }
    hess_entries = {
        ij: sp.Symbol("h" + "".join(map(str, ij)))
        for ij in combinations_with_replacement(range(m), 2)
    }
    third = {
        ijk: sp.Symbol("t" + "".join(map(str, ijk)))
        for ijk in combinations_with_replacement(range(m), 3)
    }
    H = sp.Matrix(m, m, lambda i, j: hess_entries[tuple(sorted((i, j)))])

    def dz(expr: sp.Expr, coordinate: int) -> sp.Expr:
        """Total spatial derivative on formal jets of order at most two."""
        ans = sum(sp.diff(expr, grad[a]) * H[a, coordinate] for a in range(m))
        ans += sum(
            sp.diff(expr, symbol) * third[tuple(sorted((*ij, coordinate)))]
            for ij, symbol in hess_entries.items()
        )
        if coordinate:
            q = coordinate - 1
            ans += sum(sp.diff(expr, lam[a]) * lam_y[a, q] for a in range(n))
            ans += sum(sp.diff(expr, beta[a]) * beta_y[a, q] for a in range(d))
            ans += sum(
                sp.diff(expr, lam_y[a, j]) * lam_yy[(a, *sorted((j, q)))]
                for a in range(n) for j in range(d)
            )
        return ans

    k = lam + C.T * s
    ck = C * k
    k2 = (k.T * k)[0]
    s2 = (s.T * s)[0]
    A = sp.zeros(m)
    A[0, 0] = k2 / (2 * p**2)
    for i in range(d):
        A[0, i + 1] = A[i + 1, 0] = -ck[i] / (2 * p)
        A[i + 1, i + 1] = sp.Rational(1, 2)
    F = beta.dot(s) + s2 / 2 - k2 * (1 + 1 / p) / 2
    h_rhs = sum(A[i, j] * H[i, j] for i in range(m) for j in range(m)) + F
    raw_marginal = (
        sum(H[i + 1, i + 1] for i in range(d)) / 2
        + s2 / 2 + beta.dot(s)
        - (1 + p) * k2 / (2 * p) - dz(k2 / (2 * p), 0)
    )
    check.zero("multi-factor logarithmic marginal equation", raw_marginal - h_rhs)

    drift = sp.Matrix([sp.diff(h_rhs, g) for g in grad])
    hp = H[1:, 0]
    explicit_b0 = -k2 * H[0, 0] / p**3 + ck.dot(hp) / p**2 + k2 / (2 * p**2)
    explicit_by = ck * H[0, 0] / p**2 - (C * C.T) * hp / p + beta + s - ck * (1 + 1 / p)
    check.zero("linearized wealth drift", drift[0] - explicit_b0)
    check.zero("linearized factor drift", drift[1:, :] - explicit_by)

    Qp = k * H[0, 0] / p - C.T * hp
    sources: list[sp.Expr] = [sp.Integer(0)]
    for i in range(d):
        li = lam_y[:, i]
        sources.append(li.dot(Qp) / p + beta_y[:, i].dot(s) - (1 + 1 / p) * k.dot(li))
    time_grad = [dz(h_rhs, i) for i in range(m)]
    for a in range(m):
        linearized = sum(
            A[i, j] * third[tuple(sorted((a, i, j)))]
            for i in range(m) for j in range(m)
        ) + drift.dot(H[:, a])
        check.zero(f"gradient equation {a}", time_grad[a] - linearized - sources[a])

    J = sp.Matrix([-k2 / (2 * p), *[v / 2 for v in s]])
    B = J.jacobian(grad)
    B_stated = sp.zeros(m)
    B_stated[0, 0] = k2 / (2 * p**2)
    for i in range(d):
        B_stated[0, i + 1] = -ck[i] / p
        B_stated[i + 1, i + 1] = sp.Rational(1, 2)
    check.zero("full nonsymmetric flux Jacobian", B - B_stated)
    check.zero("symmetric part of the flux Jacobian", (B + B.T) / 2 - A)
    check.zero("marginal equation in divergence form", h_rhs - sum(dz(J[i], i) for i in range(m)) - F)
    for a in range(m):
        flux = B * H[:, a]
        if a == 0:
            src = sp.Matrix([F, 0, 0])
        else:
            src = sp.zeros(m, 1)
            src[0] = -k.dot(lam_y[:, a - 1]) / p
            src[a] = F
        div_rhs = sum(dz(flux[i] + src[i], i) for i in range(m))
        check.zero(f"scalar divergence equation {a}", time_grad[a] - div_rhs)

    xi = sp.Matrix(sp.symbols("x0:3"))
    qxi = k * xi[0] / p - C.T * xi[1:, :]
    check.zero("principal quadratic form", 2 * (xi.T * A * xi)[0] - qxi.dot(qxi) - (xi[1:, :].T * R * xi[1:, :])[0])
    check.zero("primal principal determinant", A.det() - (k.T * (sp.eye(n) - C.T * C) * k)[0] / (2**m * p**2))

    r_time = 2 * sum(grad[a] * time_grad[a] for a in range(m))
    r_gradient = sp.Matrix([2 * grad.dot(H[:, i]) for i in range(m)])
    r_hessian = sp.Matrix(m, m, lambda i, j: 2 * sum(
        H[a, i] * H[a, j] + grad[a] * third[tuple(sorted((a, i, j)))] for a in range(m)
    ))
    r_linear = sum(A[i, j] * r_hessian[i, j] for i in range(m) for j in range(m)) + drift.dot(r_gradient)
    dissipation = sp.Integer(0)
    for a in range(m):
        spatial_y = H[1:, a]
        Qa = k * H[0, a] / p - C.T * spatial_y
        dissipation -= Qa.dot(Qa) + (spatial_y.T * R * spatial_y)[0]
    bernstein_source = (
        2 * (lam_y * s).dot(Qp) / p + 2 * (s.T * beta_y.T * s)[0]
        - 2 * (1 + 1 / p) * k.dot(lam_y * s)
    )
    check.zero("joint Bernstein identity", r_time - r_linear - dissipation - bernstein_source)

    kap = -k / p
    trade_noise = p * kap + C.T * s
    check.zero("marginal-utility traded diffusion", trade_noise + lam)
    h_drift = (
        -h_rhs + p * (kap.dot(lam) - kap.dot(kap) / 2) + beta.dot(s)
        + kap.dot(kap) * H[0, 0] / 2 + (C * kap).dot(hp)
        + sum(H[i + 1, i + 1] for i in range(d)) / 2
    )
    check.zero("marginal-utility logarithmic Ito drift", h_drift + (lam.dot(lam) + (s.T * R * s)[0]) / 2)
    check.zero("optimal wealth-density budget drift", kap.dot(lam) + kap.dot(trade_noise))

    G = sp.Matrix(sp.symbols("G1:3"))
    kd = lam - C.T * G
    dual_cross = C * lam + R * G
    Ad = sp.zeros(m)
    Ad[0, 0] = (lam.dot(lam) + (G.T * R * G)[0]) / 2
    for i in range(d):
        Ad[0, i + 1] = Ad[i + 1, 0] = -dual_cross[i] / 2
        Ad[i + 1, i + 1] = sp.Rational(1, 2)
    check.zero("dual principal determinant", Ad.det() - (kd.T * (sp.eye(n) - C.T * C) * kd)[0] / 2**m)
    check.zero("Legendre reduction of nonlinear terms", (kd.dot(kd) - G.dot(G)) / 2 - (lam.dot(lam) / 2 - (C * lam).dot(G) - (G.T * R * G)[0] / 2))


def scalar_and_utilities(check: Calculations) -> None:
    rho, p = sp.symbols("rho p", nonzero=True)
    lam, ly, lyy, lt, beta, by, s, hll, hly = sp.symbols("lam ly lyy lt beta by s hll hly")
    k = lam + rho * s
    drift_y = rho * k * hll / p**2 - rho**2 * hly / p + beta + s - rho * k * (1 + 1 / p)
    source_s = ly / p * (k * hll / p - rho * hly) + by * s - k * ly * (1 + 1 / p)
    raw_k = rho * source_s - lyy / 2 - drift_y * ly
    c = by - ly / rho
    g = lam * ly / rho - by * lam - beta * ly - lyy / 2
    check.zero("single-factor effective-premium equation", raw_k - c * k - g)
    check.zero("calendar-time source correction", raw_k - lt - c * k - (g - lt))

    a, bmag, rmag, omega = sp.symbols("a bmag rmag omega", positive=True)
    t = sp.Symbol("t")
    for sign in (1, -1):
        b, corr = sign * bmag, sign * rmag
        price = a + b * t
        first = b * omega * (1 - t**2)
        second = -2 * b * omega**2 * t * (1 - t**2)
        cr = -first / corr
        gr = price * first / corr - second / 2
        lower, upper = a - bmag - rmag * omega, a + bmag + rmag * omega
        factor = first / corr * (bmag + rmag * omega)
        check.zero(f"tanh lower barrier, orientation {sign}", cr * lower + gr - factor * (1 + sign * t))
        check.zero(f"tanh upper barrier, orientation {sign}", cr * upper + gr - factor * (-1 + sign * t))

    r, f0, f1, f2 = sp.symbols("r f0 f1 f2")
    g0, g1 = r * f0 - f1, r * f1 - f2
    qshape = r * (r + 1) * f0 - (2 * r + 1) * f1 + f2
    check.zero("curvature of a logarithmically modulated power", qshape - ((r + 1) * g0 - g1))
    e, q = sp.symbols("epsilon q")
    g_sine = 1 + e * (sp.sin(q) - sp.cos(q))
    j_sine = 2 + e * (sp.sin(q) - 3 * sp.cos(q))
    check.zero("oscillatory terminal curvature", j_sine - (2 * g_sine - sp.diff(g_sine, q)))

    phi = sp.Symbol("phi", positive=True)
    phiy, phiyy = sp.symbols("phiy phiyy")
    exponent = 1 / (1 + r * rho**2)
    # Divide the transformed factor equation by F, so all residuals are rational.
    logarithmic_y = exponent * phiy / phi
    Fyy_over_F = exponent * phiyy / phi + exponent * (exponent - 1) * phiy**2 / phi**2
    F_rhs_over_F = Fyy_over_F / 2 + (beta + r * rho * lam) * logarithmic_y + r * lam**2 / 2 + r * rho**2 * logarithmic_y**2 / 2
    phi_rhs = phiyy / 2 + (beta + r * rho * lam) * phiy + r * (1 + r * rho**2) * lam**2 * phi / 2
    check.zero("power-utility linear transformation", F_rhs_over_F - exponent * phi_rhs / phi)

    ell, y = sp.symbols("ell y")
    H = sp.Function("H")
    price_log = sp.Function("g")(y)
    h0 = H(ell - price_log) - price_log
    h0ell = sp.diff(h0, ell)
    check.zero("purchasing-power factor sensitivity", sp.diff(h0, y) + (h0ell + 1) * sp.diff(price_log, y))


def purchasing_power_demand(check: Calculations) -> None:
    """Verify the terminal hedge and the risk-aversion derivative of a two-power mixture."""
    aversion = sp.Symbol("aversion", positive=True)
    lam, vg = sp.symbols("lambda vg")
    s_projected = (aversion - 1) * vg
    kappa = (lam + s_projected) / aversion
    check.zero(
        "terminal purchasing-power demand",
        kappa - lam / aversion - (1 - 1 / aversion) * vg,
    )
    ell = sp.Symbol("ell", real=True)
    c1, c2, g1, g2 = sp.symbols("c1 c2 gamma1 gamma2", positive=True)
    w1, w2 = c1 * sp.exp(-g1 * ell), c2 * sp.exp(-g2 * ell)
    A = (g1 * w1 + g2 * w2) / (w1 + w2)
    variance = w1 * w2 * (g1 - g2)**2 / (w1 + w2)**2
    check.zero("two-power wealth derivative of relative risk aversion", sp.diff(A, ell) + variance)


def classification_identities(check: Calculations) -> None:
    """Check the new horizontal fluxes, clock jets, and Gaussian boundaries."""
    l, y = sp.symbols("ell y", real=True)
    h = sp.Function("h")(l, y)
    p, s = sp.diff(h, l), sp.diff(h, y)
    beta = sp.Function("beta")(y)
    v = [sp.Function(f"v{j}")(y) for j in range(3)]
    v2 = sum(a**2 for a in v)
    ap = 1 / (2 * p**2)
    F = beta * s + s**2 / 2 - v2 * (1 + 1 / p) / 2
    rhs = v2 * sp.diff(h, l, 2) / (2 * p**2) + sp.diff(h, y, 2) / 2 + F

    def X(j: int, expr: sp.Expr) -> sp.Expr:
        return v[j] * sp.diff(expr, l)

    def horizontal_spatial(w: sp.Expr) -> sp.Expr:
        return (
            sum(X(j, ap * X(j, w)) for j in range(3))
            + sp.diff(w, y, 2) / 2
            + sum(v[j] * X(j, w) / (2 * p**2) for j in range(3))
            + (beta + s) * sp.diff(w, y)
        )

    check.zero("horizontal wealth-gradient equation", sp.diff(rhs, l) - horizontal_spatial(p))
    source_s = (
        sum(X(j, -sp.diff(v[j], y) / p) for j in range(3))
        + sp.diff(beta, y) * s
        - sum(v[j] * sp.diff(v[j], y) * (1 + 1 / p) for j in range(3))
    )
    check.zero("horizontal factor-gradient equation", sp.diff(rhs, y) - horizontal_spatial(s) - source_s)
    for j in range(3):
        w = sp.Function("w")(l, y)
        check.zero(f"commutator with factor field {j}", sp.diff(X(j, w), y) - X(j, sp.diff(w, y)) - sp.diff(v[j], y) * sp.diff(w, l))

    av, u, x, lam, ux = sp.symbols("A u x lambda ux", nonzero=True)
    # av stands for the negative logarithmic wealth derivative p in this identity.
    Vxx = u * av / x
    check.zero("factor-source HJB reconstruction", -lam**2 * x * u / (2 * av) + lam**2 * u**2 / (2 * Vxx))
    check.zero("limiting orthogonal traded marginal diffusion", av * (-lam / av) + lam)

    r, cc, rem, zz = sp.symbols("r c remaining z", positive=True)
    density = sp.Symbol("Z", positive=True)
    jr = r * zz**(-r-1) * sp.exp(r*(r+1)*rem/2) * density**(-r-1)
    # d remaining = -c^2 dt; d Z = -c Z dW.
    drift = -cc**2 * sp.diff(jr, rem) + cc**2 * density**2 * sp.diff(jr, density, 2) / 2
    noise = -cc * density * sp.diff(jr, density)
    check.zero("compact power-mixture component drift", drift - (r+1)*cc**2*jr)
    check.zero("compact power-mixture component diffusion", noise - (r+1)*cc*jr)

    eps, tau = sp.symbols("epsilon tau", positive=True)
    yy = sp.Symbol("factor")
    for order in range(1, 9):
        Aplus = tau + 2*eps*yy**order*(sp.exp(order*tau)-1)/order + eps**2*yy**(2*order)*(sp.exp(2*order*tau)-1)/(2*order)
        Aminus = tau + 2*eps*yy**order*(1-sp.exp(-order*tau))/order + eps**2*yy**(2*order)*(1-sp.exp(-2*order*tau))/(2*order)
        jump = sp.diff(Aplus-Aminus, yy, order).subs(yy, 0)
        expected = 2*eps*sp.factorial(order-1)*(sp.exp(order*tau)+sp.exp(-order*tau)-2)
        check.zero(f"order-{order} opportunity-clock derivative jump", jump - expected)

    z = sp.Symbol("z", positive=True)
    # alpha and b here correspond to zeta and b_* in the manuscript.
    alpha, b = sp.symbols("alpha b", positive=True)
    logI = -b * sp.log(z) + alpha * sp.log(1 + 1/z)**2
    reciprocal_rra = b + 2*alpha*sp.log(1+1/z)/(1+z)
    check.zero("inverse-marginal risk-aversion formula", -z*sp.diff(logI, z) - reciprocal_rra)

    A, q = sp.symbols("A q", positive=True)
    N, logz = sp.symbols("N logz", real=True)
    c0 = A/2 - logz
    u = sp.sqrt(A)*N + c0
    expanded = sp.expand(q*(b*u + alpha*u**2) - N**2/2)
    expected = (q*alpha*A-sp.Rational(1,2))*N**2 + q*sp.sqrt(A)*(b+2*alpha*c0)*N + q*(b*c0+alpha*c0**2)
    check.zero("Gaussian moment asymptotic polynomial", expanded - expected)
    logzc = A/2 + b/(2*alpha)
    check.zero("critical linear coefficient", (q*sp.sqrt(A)*(b+2*alpha*c0)).subs(logz, logzc))
    check.zero("L2 critical variance", (2*alpha*A-sp.Rational(1,2)).subs(A, 1/(4*alpha)))
    check.zero("finite-value quadratic coefficient at L2 boundary", (alpha*A-sp.Rational(1,2)).subs(A, 1/(4*alpha)) + sp.Rational(1,4))
    check.zero("explicit critical multiplier logarithm", logzc.subs({A:sp.Rational(2,3), b:2, alpha:sp.Rational(3,8)}) - 3)

    py, rho, bet, by, l0, ly, lyy, kconst = sp.symbols("p rho beta beta_y lambda lambda_y lambda_yy k", nonzero=True)
    sy = (kconst-l0)/rho
    primitive = rho*by*sy - (bet+sy)*ly - lyy/2
    cr = by - ly/rho
    gr = l0*ly/rho - by*l0 - bet*ly - lyy/2
    check.zero("constant-boundary effective-premium reaction", primitive-cr*kconst-gr)


def exposure_energy_identity(check: Calculations) -> None:
    """Verify the completion of squares in the replication energy estimate."""
    theta, wealth, premium, rate = sp.symbols(
        "theta wealth premium rate", real=True
    )
    drift = theta**2 + 2 * premium * wealth * theta + rate * wealth**2
    squares = (
        theta**2 / 2
        + (theta + 2 * premium * wealth)**2 / 2
        + (rate - 2 * premium**2) * wealth**2
    )
    check.zero("replication energy completion of squares", drift - squares)


def main() -> int:
    checks = Calculations()
    multidimensional(checks)
    scalar_and_utilities(checks)
    purchasing_power_demand(checks)
    classification_identities(checks)
    exposure_energy_identity(checks)
    print(f"All {checks.count} exact algebraic calculations returned zero residual.")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except AssertionError as error:
        print(f"Calculation failed: {error}", file=sys.stderr)
        sys.exit(1)