import json, math, random
import numpy as np
import torch

SEED=2543
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DT=0.1; K=30; N=3000

def features(x):
    # Smooth characteristic features plus moments; sufficient to expose law mismatch.
    fs=[x, x*x, x**3]
    for w in [0.35,0.7,1.2,2.0,3.0]:
        fs += [torch.cos(w*x), torch.sin(w*x)]
    return torch.stack(fs,-1)

def np_features(x):
    x=torch.as_tensor(x,dtype=torch.float32)
    return features(x)

def mixture_stats(paths, hazards):
    # States x_0,...,x_K: stop at k<K; all residual mass emits x_K.
    q=1-torch.exp(-torch.nn.functional.softplus(hazards)*DT)
    surv=torch.cumprod(torch.cat([torch.ones(1),1-q[:-1]]),0)
    w=torch.zeros_like(q)
    w[:-1]=surv[:-1]*q[:-1]
    w[-1]=surv[-1]
    z=(features(paths)*w[None,:,None]).sum(1)
    expected_steps=(surv[:-1]*DT).sum()
    return z.mean(0), w, expected_steps

def mmd_feature(a,b):
    return ((a-b)**2).mean()

def exact_checks():
    rows=[]
    for h in [0.2,0.5,1.0,2.0]:
        q=1-math.exp(-h*DT)
        s=np.array([(1-q)**k for k in range(K+1)])
        # fit log survival slope, excluding endpoint numerical issues
        slope=np.polyfit(np.arange(K+1)*DT,np.log(s),1)[0]
        # discrete expected number of active intervals, matching sum survival
        expected=float(s[:-1].sum()*DT)
        pred=(1-math.exp(-h*K*DT))/h
        rows.append({'h':h,'predicted_log_slope':-h,'observed_log_slope':float(slope),
                     'slope_abs_err':float(abs(slope+h)),
                     'predicted_expected_time':pred,'observed_expected_time':expected,
                     'time_abs_err':abs(expected-pred)})
    # survival-mixture mass for arbitrary positive hazards
    hz=np.array([.1,.4,1.1,2.3,.7]*7)[:K+1]
    q=1-np.exp(-hz*DT); surv=np.cumprod(np.r_[1.,1-q[:-1]])
    mass=float(np.sum(surv*q)+surv[-1]*(1-q[-1]))
    # Prediction 3: discrete expected stopping time converges to the
    # continuous integral as dt decreases (first-order Riemann error).
    conv=[]
    h=.8; horizon=3.0
    for dt in [.2,.1,.05,.025]:
        n=round(horizon/dt); r=math.exp(-h*dt)
        discrete=dt*(1-r**n)/(1-r)
        continuous=(1-math.exp(-h*horizon))/h
        conv.append({'dt':dt,'discrete_expected_time':discrete,
                     'continuous_prediction':continuous,
                     'abs_error':abs(discrete-continuous)})
    return rows, mass, conv

def train_policy(paths, target_feat, steps=350):
    p=torch.nn.Parameter(torch.full((K+1,),-1.0))
    opt=torch.optim.Adam([p],lr=.08)
    history=[]
    for it in range(steps):
        opt.zero_grad()
        pred,w,etime=mixture_stats(paths,p)
        loss=mmd_feature(pred,target_feat)+1e-4*etime
        loss.backward(); opt.step()
        if it in [0,49,149,349]: history.append(float(loss.detach()))
    with torch.no_grad():
        pred,w,etime=mixture_stats(paths,p)
        q=1-torch.exp(-torch.nn.functional.softplus(p)*DT)
        # empirical stopped samples use weighted mixture CDF proxy; law metric is feature MSE
        return float(mmd_feature(pred,target_feat)),float(etime),q.numpy(),history

def run():
    # Common Brownian paths; X_0=0 and X_k=sqrt(t_k) Gaussian marginals.
    inc=np.random.randn(N,K).astype('float32')*math.sqrt(DT)
    paths=np.concatenate([np.zeros((N,1),dtype='float32'),np.cumsum(inc,axis=1)],axis=1)
    pt=torch.from_numpy(paths)
    # Target law is a known constant-hazard stopped Brownian law (independent samples).
    htrue=.8
    u=np.random.rand(N); tau=-np.log(u)/htrue
    kt=np.minimum((tau/DT).astype(int),K)
    target=np.zeros(N,dtype='float32')
    for k in range(K):
        mask=kt==k
        target[mask]=paths[mask,k] # couple only for reproducibility; marginal is Brownian
    mask=kt==K; target[mask]=paths[mask,K]
    target_feat=np_features(target).mean(0)
    # Standard fixed-budget baseline: emit X_K.
    baseline=float(mmd_feature(np_features(paths[:,:][:,-1]).mean(0),target_feat))
    idea,etime,q,history=train_policy(pt,target_feat)
    # Also report learned hazard averaged over its (time-only) head and fit survival.
    surv=np.cumprod(np.r_[1.,1-q[:-1]])
    slope=float(np.polyfit(np.arange(K+1)*DT,np.log(np.maximum(surv,1e-12)),1)[0])
    return {'checks':{'constant_hazard_sweep':exact_checks()[0],
                      'mixture_mass_error':abs(exact_checks()[1]-1.0),
                      'dt_convergence_sweep':exact_checks()[2]},
            'mini_experiment':{'target_true_hazard':htrue,'baseline_feature_mse':baseline,
                               'idea_feature_mse':idea,'idea_expected_compute_time':etime,
                               'fixed_budget_compute_time':K*DT,
                               'learned_survival_log_slope':slope,
                               'training_loss_checkpoints':history,
                               'learned_q_first_last':[float(q[0]),float(q[-1])]}}

if __name__=='__main__':
    try:
        # CUDA is allowed, but this tiny verification is deliberately CPU-safe.
        out=run()
    except Exception as e:
        out={'error':repr(e)}
    print(json.dumps(out,indent=2))