#!/usr/bin/env python3
"""DLEGOS numerical-development companion. All experiments are constructed.
Run: python climate_math_lab.py --output results
Requires NumPy only. No climate observations or pretrained weights are used.
"""
from pathlib import Path
import argparse,csv,json,math,platform
import numpy as np


def graph_system(n=12):
    """Closed chain; B is -1 at edge sender, +1 at receiver."""
    B=np.zeros((n,n-1),dtype=np.float64)
    for e in range(n-1):B[e,e]=-1;B[e+1,e]=1
    C=1+0.4*np.linspace(0,1,n)
    x=np.linspace(-1,1,n-1)
    features=np.stack([x,x*x,np.cos(np.pi*x)],axis=1)
    truth=0.25+0.12*(x+1)+0.07*np.cos(np.pi*x)
    return B,C,features,truth


def rates(T,G,B,C):
    """T shape (batch,n); G (edges). Returns temperature tendencies."""
    diff=T@B
    q=-diff*G
    return (q@B.T)/C


def step(T,G,B,C,dt):
    """Positivity/maximum-principle bound for pure closed diffusion."""
    bound=float(np.min(C/(np.abs(B)@G)))
    if dt<0 or dt>bound*(1+1e-12):raise ValueError('Time step violates convex-combination bound')
    return T+dt*rates(T,G,B,C)


def softplus(x):return np.maximum(x,0)+np.log1p(np.exp(-np.abs(x)))
def sigmoid(x):return np.exp(-np.logaddexp(0,-x))


class EdgeMLP:
    """3 -> 8 tanh -> 1 positive conductance, shared across all edges."""
    def __init__(self,seed=23):
        rng=np.random.default_rng(seed)
        self.p={'W':rng.normal(0,.15,(3,8)),'b':np.zeros(8),'v':rng.normal(0,.15,8),'d':np.array(-.8)}
    def forward(self,X):
        H=np.tanh(X@self.p['W']+self.p['b']);z=H@self.p['v']+self.p['d']
        return softplus(z)+1e-6,(H,z)
    def loss_gradient(self,X,T,target,B,C):
        G,(H,z)=self.forward(X);diff=T@B;pred=((-diff*G)@B.T)/C
        error=pred-target;loss=float(np.mean(error**2));d_pred=2*error/error.size
        d_q=(d_pred/C)@B
        d_G=np.sum(d_q*(-diff),axis=0)
        d_z=d_G*sigmoid(z);d_H=d_z[:,None]*self.p['v'][None,:]
        d_A=d_H*(1-H*H)
        grad={'W':X.T@d_A,'b':d_A.sum(axis=0),'v':H.T@d_z,'d':np.array(d_z.sum())}
        return loss,grad


def gradient_check(model,X,T,target,B,C):
    _,grads=model.loss_gradient(X,T,target,B,C);errors=[]
    for key,idx in [('W',(0,0)),('W',(2,5)),('b',(4,)),('v',(3,)),('d',())]:
        old=float(model.p[key][idx]);eps=1e-6
        model.p[key][idx]=old+eps;plus=model.loss_gradient(X,T,target,B,C)[0]
        model.p[key][idx]=old-eps;minus=model.loss_gradient(X,T,target,B,C)[0]
        model.p[key][idx]=old
        finite=(plus-minus)/(2*eps);analytic=float(grads[key][idx]);errors.append(abs(finite-analytic)/max(1e-8,abs(finite),abs(analytic)))
    assert max(errors)<2e-5,errors
    return max(errors)


def train_constructed(seed=23,iterations=1200):
    B,C,X,trueG=graph_system();rng=np.random.default_rng(seed)
    T=285+rng.normal(0,3,(192,12));target=rates(T,trueG,B,C)
    Ttest=285+rng.normal(0,3,(96,12));test_target=rates(Ttest,trueG,B,C)
    model=EdgeMLP(seed);gc=gradient_check(model,X,T[:8],target[:8],B,C)
    first=model.loss_gradient(X,T,target,B,C)[0];m={k:np.zeros_like(v) for k,v in model.p.items()};v={k:np.zeros_like(v) for k,v in model.p.items()};history=[]
    # Adam update with bias correction, fixed full-batch constructed training set.
    for it in range(1,iterations+1):
        loss,grad=model.loss_gradient(X,T,target,B,C)
        for k in model.p:
            m[k]=.9*m[k]+.1*grad[k];v[k]=.999*v[k]+.001*grad[k]**2
            model.p[k]-=.015*(m[k]/(1-.9**it))/(np.sqrt(v[k]/(1-.999**it))+1e-8)
        if it==1 or it%20==0:history.append({'iteration':it,'training_tendency_mse':loss})
    G=model.forward(X)[0];final=model.loss_gradient(X,T,target,B,C)[0]
    baselineG=np.full_like(G,.4)
    rmse=lambda pred:float(np.sqrt(np.mean((pred-test_target)**2)))
    learned_rmse=rmse(rates(Ttest,G,B,C));fixed_rmse=rmse(rates(Ttest,baselineG,B,C))
    bound=float(np.min(C/(np.abs(B)@G)));dt=.8*bound
    state=Ttest[0:1].copy();E0=float((state@C).item());mean=E0/C.sum();v0=float(((state-mean)**2*C).sum()/2);roll=[]
    for k in range(101):
        E=float((state@C).item());variance=float(((state-mean)**2*C).sum()/2)
        roll.append({'step':k,'time_s':k*dt,'energy_J':E,'absolute_energy_drift_J':abs(E-E0),'weighted_variance_J_K':variance,'min_temperature_K':float(state.min()),'max_temperature_K':float(state.max())})
        if k<100:state=step(state,G,B,C,dt)
    drift=max(r['absolute_energy_drift_J'] for r in roll)
    assert final<first/100 and learned_rmse<fixed_rmse
    assert drift<1e-9
    assert all(roll[k+1]['weighted_variance_J_K']<=roll[k]['weighted_variance_J_K']+1e-9 for k in range(100))
    assert np.max(np.abs(B.sum(axis=0)))==0
    report={'seed':seed,'iterations':iterations,'n_cells':12,'n_edges':11,'train_fields':192,'test_fields':96,'test_scope':'unseen constructed temperature fields on the SAME graph, not unseen geography','initial_training_tendency_mse_K2_s2':first,'final_training_tendency_mse_K2_s2':final,'fixed_conductance_test_tendency_rmse_K_s':fixed_rmse,'learned_test_tendency_rmse_K_s':learned_rmse,'conductance_rmse_W_K':float(np.sqrt(np.mean((G-trueG)**2))),'finite_difference_gradient_max_relative_error':float(gc),'convexity_step_bound_s':bound,'rollout_dt_s':dt,'max_energy_drift_J':drift,'initial_weighted_variance_J_K':v0,'final_weighted_variance_J_K':roll[-1]['weighted_variance_J_K'],'evidence_class':'constructed code-verification experiment; not climate forecast validation'}
    edge_rows=[{'edge':i,'true_conductance_W_K':float(a),'learned_conductance_W_K':float(b)} for i,(a,b) in enumerate(zip(trueG,G))]
    return report,history,roll,edge_rows,model


def advection_check(n,velocity=.7,end_time=.4,cfl=.6):
    dx=1/n;steps=math.ceil(end_time*abs(velocity)/(cfl*dx));dt=end_time/steps;nu=velocity*dt/dx
    x=(np.arange(n)+.5)*dx
    # Exact cell averages for a smooth periodic solution.
    q0=2+np.sinc(1/n)*np.sin(2*np.pi*x);q=q0.copy()
    for _ in range(steps):
        if velocity>=0:q=q-nu*(q-np.roll(q,1))
        else:q=q-nu*(np.roll(q,-1)-q)
    exact=2+np.sinc(1/n)*np.sin(2*np.pi*(x-velocity*end_time))
    error=float(np.mean(np.abs(q-exact)));drift=float(abs(dx*q.sum()-dx*q0.sum()))
    assert abs(nu)<=1 and q.min()>=0 and drift<1e-12
    return {'cells':n,'dx':dx,'steps':steps,'dt_s':dt,'courant_number':nu,'L1_cell_average_error':error,'mass_drift':drift}


def weighted_projection(z,H,b,P):
    residual=b-H@z;gram=H@P@H.T
    # Solve, do not explicitly form a matrix inverse.
    return z+P@H.T@np.linalg.solve(gram,residual)


def crps_ensemble(samples,y,fair=False):
    x=np.asarray(samples,dtype=float);M=x.size
    if fair and M<2:raise ValueError('Fair CRPS needs at least two members')
    pair=np.abs(x[:,None]-x[None,:]).sum();den=2*M*(M-1) if fair else 2*M*M
    return float(np.mean(np.abs(x-y))-pair/den)


def supplementary_checks():
    H=np.array([[.25,.25,.25,.25]]);z=np.array([2.,4.,6.,8.]);P=np.eye(4);b=np.array([6.])
    projected=weighted_projection(z,H,b,P);assert np.max(abs(H@projected-b))<1e-12
    negative=weighted_projection(z,H,np.array([0.]),P)
    assert negative.min()<0 # Equality projection alone does not enforce nonnegativity.
    assert crps_ensemble([1,3],2)==.5 and crps_ensemble([1,3],2,True)==0
    conv=[advection_check(n) for n in [32,64,128,256]]
    for i in range(1,len(conv)):conv[i]['observed_order']=math.log(conv[i-1]['L1_cell_average_error']/conv[i]['L1_cell_average_error'],2)
    assert .85<conv[-1]['observed_order']<1.15
    # Constant-input exact preservation for the physical transport operator.
    B,C,_,G=graph_system();assert np.max(abs(rates(np.full((1,12),290.),G,B,C)))<1e-12
    return {'projection_example':{'prior':z.tolist(),'target_mean':6,'projected':projected.tolist(),'equality_residual':float(np.max(abs(H@projected-b))),'zero_mean_projection':negative.tolist(),'nonnegative_guarantee':False},'crps_example':{'members':[1,3],'observation':2,'empirical_crps':.5,'fair_crps':0.},'advection':conv,'checks':'all assertions passed'}


def write_csv(path,rows):
    with path.open('w',newline='') as f:
        writer=csv.DictWriter(f,fieldnames=list(dict.fromkeys(k for row in rows for k in row)));writer.writeheader();writer.writerows(rows)


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--output',default='results');args=parser.parse_args();out=Path(args.output);out.mkdir(parents=True,exist_ok=True)
    report,history,roll,edges,model=train_constructed();supp=supplementary_checks()
    summary={'status':'EXECUTED; all assertions passed','evidence_class':'constructed numerical-development tests only','python':platform.python_version(),'numpy':np.__version__,'training':report,'supplementary':supp}
    (out/'report.json').write_text(json.dumps(summary,indent=2)+'\n')
    write_csv(out/'training.csv',history);write_csv(out/'rollout.csv',roll);write_csv(out/'conductances.csv',edges);write_csv(out/'advection.csv',supp['advection'])
    (out/'weights.json').write_text(json.dumps({k:v.tolist() for k,v in model.p.items()},indent=2)+'\n')
    print(json.dumps(summary,indent=2))
if __name__=='__main__':main()