#!/usr/bin/env python3
"""nonlinear_check.py — ESM. Verifies the linearised system against the FULL
nonlinear equilibrium (E1)+(E2')+(E3') with the exactly deformed metric.
(A) forcing identity b*int K*Phi = Khat_n    (B) transfer function T_n
(C) welfare curvature dlogV/eps^2 = Lambda_1/4 at mu=0.7
(D) reversal-band boundary mu**_1 (nonlinear sign flip and crossing)"""
import numpy as np
from scipy.integrate import quad

SIG, B_, KB = 4.0, 15.0, 1.0

Khat = lambda n: 2*B_*(1-(-1)**n*np.exp(-B_/2))/(B_**2+4*np.pi**2*n**2)
T0 = 2/B_*(1-np.exp(-B_/2))
phi = lambda n: Khat(n)/T0
Om  = lambda n: phi(n)/(SIG+(SIG-1)*phi(n))
Ac  = lambda n: 1+(2*SIG-1)/(SIG-1)*Om(n)
Tn  = lambda n, mu: mu*(Ac(n)-1)/(KB*(1-mu*Ac(n)))
nuf = lambda m: ((m/(1-m))*(2*SIG-1)-(SIG-1)**2)/(SIG*(SIG-1))
Wc  = lambda m: m*(2*SIG-1)/(SIG-1)*(m*SIG-(SIG-1))/((1-m)*(SIG-1))
Hr  = lambda r: 2/B_**2*(np.exp(-B_*min(r,1-r))-np.exp(-B_/2))-2/B_*(0.5-min(r,1-r))*np.exp(-B_/2)
Cc  = (2-(2+B_)*np.exp(-B_/2))/B_**2
Hh  = lambda n: quad(lambda r: Hr(r)*np.cos(2*np.pi*n*r), 0, 1, limit=800)[0]
def Lam(n, m):
    p = phi(n); jen = -2*B_*Cc + B_**2*Hh(n)
    return (Wc(m)*p**2/(SIG**2*KB**2*(1-nuf(m)*p)**2)
            + m*jen/((SIG-1)*T0*KB**2) + 2*m*nuf(m)*p**2/((SIG-1)*KB**2*(1-nuf(m)*p)))

def equilibrium(mu, eta, eps, tol=1e-12, itmax=60000, ret="V"):
    N = len(eta)
    xs = np.arange(N)/N; dx = 1.0/N; i = np.arange(N)
    DIF = (i[None, :]-i[:, None]) % N
    right = DIF <= N//2
    dens = 1.0/(KB + eps*eta)
    c = np.concatenate(([0.0], np.cumsum(dens)*dx)); Ctot = c[-1]
    R = (c[i[None, :]] - c[i[:, None]]) % Ctot
    D = np.where(right, R, Ctot - R)
    K = np.exp(-B_*D)
    lam = np.ones(N); w = np.ones(N)
    for it in range(itmax):
        P1s = (lam*w**(1-SIG)) @ K * dx
        wn = ((lam*w/P1s) @ K.T * dx)**(1/SIG)
        wn /= np.exp(np.mean(np.log(wn)))
        w = 0.5*w + 0.5*wn
        P1s = (lam*w**(1-SIG)) @ K * dx
        V = w**mu * P1s**(mu/(SIG-1)) * lam**(-(1-mu))
        lamn = lam*(V/np.mean(V))**(0.5/(1-mu)); lamn /= np.mean(lamn)
        dl = np.max(np.abs(lamn-lam)); lam = lamn
        if dl < tol and it > 50:
            break
    P1s = (lam*w**(1-SIG)) @ K * dx
    V = w**mu * P1s**(mu/(SIG-1)) * lam**(-(1-mu))
    assert V.std()/V.mean() < 1e-8, ("utility not equalised", V.std()/V.mean())
    return (lam, w, float(np.mean(V))) if ret == "all" else float(np.mean(V))

if __name__ == "__main__":
    N = 768
    xs = np.arange(N)/N; dx = 1.0/N; i = np.arange(N)
    DIF = (i[None, :]-i[:, None]) % N
    dbase = np.minimum(DIF, N-DIF)/N; K0m = np.exp(-B_*dbase)
    for n in (1, 2, 3):
        r = DIF/N; r = np.where(r > 0.5, r-1.0, r)
        s = np.sign(r)*(np.exp(2j*np.pi*n*r)-1)/(2j*np.pi*n)
        PHI = np.exp(2j*np.pi*n*xs[:, None])*s
        ghat = ((B_*(K0m*PHI).sum(axis=0)*dx)*np.exp(-2j*np.pi*n*xs)).mean()
        print(f"(A) n={n}: b*int(K.Phi)/Khat = {abs(ghat)/Khat(n):.6f}   (identity: 1)")
    MU = 0.7
    V0 = equilibrium(MU, np.zeros(N), 0.0)
    print(f"(B) Vbar(0) = {V0:.6f}   theory T0^(mu/(sig-1)) = {T0**(MU/(SIG-1)):.6f}")
    for n in (1, 2):
        for eps in (0.02, 0.01):
            lam, w, V = equilibrium(MU, np.cos(2*np.pi*n*xs), eps, ret="all")
            a = 2*np.mean(lam*np.cos(2*np.pi*n*xs))
            print(f"    T_{n} eps={eps}: nonlinear {a/eps:.4f}   theory {Tn(n, MU):.4f}"
                  f"   rel.err {abs(a/eps-Tn(n, MU))/Tn(n, MU):.2%}")
    for eps in (0.02, 0.01, 0.005):
        V = equilibrium(MU, np.cos(2*np.pi*xs), eps)
        print(f"(C) eps={eps}: dlogV/eps^2 = {(np.log(V)-np.log(V0))/eps**2:+.5f}"
              f"   theory Lam1/4 = {Lam(1, MU)/4:+.5f}")
    N = 512; xs = np.arange(N)/N
    eta = np.cos(2*np.pi*xs); eps = 0.01; rows = []
    for mu in (0.62, 0.635, 0.640, 0.66):
        V0b = equilibrium(mu, np.zeros(N), 0.0)
        Vb = equilibrium(mu, eta, eps)
        cc = (np.log(Vb)-np.log(V0b))/eps**2; rows.append((mu, cc))
        print(f"(D) mu={mu}: nonlinear {cc:+.5f}   theory Lam1/4 = {Lam(1, mu)/4:+.5f}"
              f"   sign {'MATCH' if np.sign(cc) == np.sign(Lam(1, mu)) else 'MISMATCH'}")
    (m1, c1), (m2, c2) = rows[1], rows[2]
    print(f"(D) nonlinear crossing = {m1 - c1*(m2-m1)/(c2-c1):.4f}   (Hessian zero: 0.6377)")