import json
import numpy as np

SEED=1051
rng=np.random.default_rng(SEED)


def monotone_quantile(h,qmin=0.,qmax=1.,eps=1e-8):
    h=np.asarray(h)
    delta=np.logaddexp(0.,h)+eps
    return qmin+(qmax-qmin)*np.cumsum(delta,axis=-1)/delta.sum(axis=-1,keepdims=True)


def pushforward_density(q,p,edges):
    # q and p include both endpoints; each p interval carries exactly its mass.
    mass=np.zeros(len(edges)-1)
    for a,b,m in zip(q[:-1],q[1:],np.diff(p)):
        if b<=a: continue
        lo,hi=max(a,edges[0]),min(b,edges[-1])
        if hi<=lo: continue
        j0=max(0,np.searchsorted(edges,lo,side='right')-1)
        j1=min(len(mass)-1,np.searchsorted(edges,hi,side='left'))
        for j in range(j0,j1+1):
            overlap=max(0.,min(hi,edges[j+1])-max(lo,edges[j]))
            mass[j]+=m*overlap/(b-a)
    return mass/np.diff(edges)


def invariant_checks():
    K=64; H=rng.normal(size=(2000,K)); out={}
    # Prediction 1: positivity of increments implies monotonicity for every logit scale.
    temps=[.05,.2,1.,5.,20.]; rows=[]
    for t in temps:
        q=monotone_quantile(t*H)
        rows.append({'temperature':t,'monotone_fraction':float(np.mean(np.all(np.diff(q,axis=1)>=0,axis=1))), 'min_increment':float(np.min(np.diff(q,axis=1)))})
    out['temperature_sweep']=rows
    # Prediction 2: normalization fixes the final span independent of logits.
    q=monotone_quantile(H,qmin=-2,qmax=3)
    out['span_prediction']={'predicted_final_endpoint':3.,'observed_final_endpoint':float(np.mean(q[:,-1])),'endpoint_error':float(np.max(np.abs(q[:,-1]-3)))}
    # Prediction 3: every increment is >= L*eps/sum(delta); increasing eps raises this bound.
    epses=[1e-10,1e-7,1e-4,1e-2]; pred=[]; obs=[]
    for e in epses:
        delta=np.logaddexp(0.,H)+e; q=monotone_quantile(H,eps=e)
        pred.append(float(np.min(e/delta.sum(axis=1))))
        obs.append(float(np.min(np.diff(q,axis=1))))
    out['epsilon_sweep']={'eps':epses,'predicted_bound':pred,'observed_min_increment':obs,'bound_holds':bool(all(o+1e-15>=p for o,p in zip(obs,pred)))}
    p=np.linspace(0,1,K); edges=np.linspace(0,1,257); u=pushforward_density(monotone_quantile(H[0]),p,edges)
    out['pushforward']={'negative_cells':int((u<0).sum()),'mass':float(np.sum(u*np.diff(edges))),'mass_error':float(abs(np.sum(u*np.diff(edges))-1.))}
    return out


def fields(n=800,nx=128):
    x=np.linspace(0,1,nx); ans=[]
    for _ in range(n):
        c=rng.uniform(.15,.85); w=rng.uniform(.025,.10)
        y=np.exp(-.5*((x-c)/w)**2)+.35*np.exp(-.5*((x-(c-.14))/(.6*w))**2)
        y=np.maximum(y,0); y/=np.trapz(y,x); ans.append(y)
    return np.asarray(ans)


def quantiles(U,K):
    x=np.linspace(0,1,U.shape[1]); dx=x[1]-x[0]; p=np.linspace(0,1,K); out=[]
    for u in U:
        c=np.cumsum(u)*dx; c=(c-c[0])/(c[-1]-c[0])
        out.append(np.interp(p,c,x))
    return np.asarray(out)


def experiment():
    U=fields(); split=600; trainU,testU=U[:split],U[split:]; Q=quantiles(U,64); trainQ,testQ=Q[:split],Q[split:]
    meanQ=trainQ.mean(0); _,_,VQ=np.linalg.svd(trainQ-meanQ,full_matrices=False)
    meanU=trainU.mean(0); _,_,VU=np.linalg.svd(trainU-meanU,full_matrices=False)
    edges=np.linspace(0,1,U.shape[1]+1); p=np.linspace(0,1,64); rows=[]
    for d in [2,4,8]:
        # Standard reduced decoder: unconstrained physical POD field.
        base=meanU+(testU-meanU)@VU[:d].T@VU[:d]
        # Proposed CDT/POD decoder, with cumulative maximum as a robust monotone projection.
        qraw=meanQ+(testQ-meanQ)@VQ[:d].T@VQ[:d]
        qidea=np.clip(np.maximum.accumulate(qraw,axis=1),0,1)
        rec=np.asarray([pushforward_density(q,p,edges) for q in qidea])
        rows.append({'d':d,'baseline_physical_L1':float(np.mean(np.abs(base-testU))), 'idea_physical_L1':float(np.mean(np.abs(rec-testU))), 'baseline_wasserstein_proxy':float(np.mean(np.abs(quantiles(np.maximum(base,0),64)-testQ))), 'idea_wasserstein':float(np.mean(np.abs(qidea-testQ))), 'baseline_negative_fraction':float(np.mean(base<0)), 'idea_negative_fraction':float(np.mean(rec<0)), 'baseline_mass_abs_error':float(np.mean(np.abs(np.trapz(base,dx=1/127,axis=1)-1))), 'idea_mass_abs_error':float(np.mean(np.abs(np.sum(rec,axis=1)*(1/128)-1))), 'idea_monotone_fraction':float(np.mean(np.all(np.diff(qidea,axis=1)>=-1e-12,axis=1)))})
    return rows

if __name__=='__main__':
    report={'seed':SEED,'checks':invariant_checks(),'experiment':experiment()}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))