import json, math
import numpy as np
from scipy.linalg import expm

SEED=1094
np.random.seed(SEED)

# ---------- Core geometric calculation ----------
def oblique_sweep():
    # At a vertical face n=(1,0), G=[[1,c],[c,1]].
    # r=Gn/(n^T G n)=(1,c), so r.grad=0 iff grad_x+c grad_y=0.
    cs=np.array([0., .2, .5, .8, .95])
    # A deliberately oblique tangent-compatible gradient: grad=(-c,1).
    oblique=np.abs((-cs)+cs*1.)
    # The normal Neumann condition grad_x=0 leaves residual c in the oblique BC.
    normal_res=np.abs(cs)
    # Normal diffusive flux n^T G grad has the same residual for this test.
    flux_res=np.abs(cs+(-cs)) # exact zero under oblique condition
    # Use a second gradient to show the normal condition fails flux cancellation.
    wrong_flux=np.abs(cs)
    return {"correlations":cs.tolist(), "oblique_residual":oblique.tolist(),
            "normal_bc_residual":normal_res.tolist(), "oblique_flux_residual":flux_res.tolist(),
            "normal_bc_flux_residual":wrong_flux.tolist()}

# ---------- Finite-volume reflecting diffusion ----------
def reflected_generator(n, D):
    # Cell-centered 1D finite-volume generator for q_t=D q_xx, zero face flux.
    # Column vector q, Q has column sums zero and nonnegative off diagonal.
    h=1.0/n; Q=np.zeros((n,n))
    rate=D/(h*h)
    for j in range(n):
        if j>0: Q[j-1,j]+=rate; Q[j,j]-=rate
        if j<n-1: Q[j+1,j]+=rate; Q[j,j]-=rate
    return Q

def conservation_sweep():
    n=80; x=(np.arange(n)+.5)/n
    q=np.exp(-((x-.23)/.07)**2); q/=q.sum()
    rows=[]
    for D in [0., 1e-3, .03, .3, 1.0]:
        Q=reflected_generator(n,D)
        q1=expm(Q*.37)@q
        # Boundary face current is imposed to exactly zero by construction.
        mass_err=abs(q1.sum()-1.)
        minq=q1.min()
        rows.append({"D":D,"mass_error":float(mass_err),"min_density":float(minq),
                     "boundary_flux_left":0.,"boundary_flux_right":0.})
    return rows

def stability_sweep():
    n=80; D=1.; h=1/n
    Q=reflected_generator(n,D)
    lam=np.linalg.eigvalsh(Q).min()
    dt_spec=2/abs(lam)
    # For forward Euler, positivity requires dt <= 1/max_j(-Q_jj),
    # while l2 stability only requires dt <= 2/|lambda_min|.
    dt_pos=1/np.max(-np.diag(Q))
    # Positive alternating data strongly excites the highest-frequency mode.
    z0=1.+.9*(-1.)**np.arange(n); z0/=z0.sum()
    ratios=[.5,.9,1.0,1.05,1.5]
    out=[]
    for ratio in ratios:
        dt=ratio*dt_pos; z=z0.copy()
        for _ in range(80): z=z+dt*(Q@z)
        out.append({"ratio_to_positivity_bound":ratio,"dt":dt,
                    "min_after_80":float(z.min()),"l1_norm":float(np.linalg.norm(z,1)),
                    "nonnegative_observed":bool(z.min()>=-1e-12)})
    return {"lambda_min":float(lam),"predicted_positivity_dt":float(dt_pos),
            "predicted_spectral_dt":float(dt_spec),"sweep":out}

# ---------- Small endpoint scaling experiment ----------
def kl(p,q):
    return float(np.sum(p*np.log((p+1e-15)/(q+1e-15))))

def sinkhorn(P, p0, p1, iters=100):
    # P maps initial mass to terminal mass: terminal=P @ initial.
    a=np.ones(len(p0)); b=np.ones(len(p1))
    # joint diag(a) P diag(b), alternating endpoint marginal fitting
    for k in range(iters):
        a=p0/(P.T@b+1e-15)
        b=p1/(P@a+1e-15)
    return a,b

def sinkhorn_demo():
    n=60; x=(np.arange(n)+.5)/n; h=1/n
    p0=np.exp(-.5*((x-.22)/.07)**2); p0/=p0.sum()
    p1=np.exp(-.5*((x-.78)/.07)**2); p1/=p1.sum()
    D=.015; T=.45
    # PDE/reflected kernel from finite-volume Fokker-Planck operator.
    Pref=expm(reflected_generator(n,D)*T)
    # Unconstrained Gaussian transition kernel, row-normalized only for comparison.
    K=np.exp(-(x[:,None]-x[None,:])**2/(4*D*T)); K/=K.sum(axis=0,keepdims=True)
    results=[]
    for name,P in [("reflected_PDE",Pref),("unconstrained_kernel",K)]:
        a,b=sinkhorn(P,p0,p1,80)
        left=a*(P.T@b); right=b*(P@a)
        results.append({"method":name,"initial_KL":kl(left/left.sum(),p0),
                        "terminal_KL":kl(right/right.sum(),p1),
                        "boundary_mass_initial":float(left[0]+left[-1]),
                        "boundary_mass_terminal":float(right[0]+right[-1])})
    return results

def main():
    report={"seed":SEED,"oblique_geometry":oblique_sweep(),
            "finite_volume_no_flux":conservation_sweep(),"explicit_stability":stability_sweep(),
            "sinkhorn_demo":sinkhorn_demo()}
    with open("results.json","w") as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=="__main__": main()