"""Jump-adapted approximation of the *same* joint curved-fold model.

Retained marks have |z|>=cutoff and occur at their actual Poisson arrival times.
Between marks, the normal coefficients are frozen over intervals of length at
most max_step; a killed Bessel transition is then sampled. Tangential OU flow
is exact over that interval. The full variable-coefficient, infinite-activity
model is NOT sampled exactly. Small marks are omitted, not replaced by drift.
Contact times are saved as intervals; jump times are exact for the cutoff model.
"""
from __future__ import annotations
from dataclasses import asdict
from pathlib import Path
import argparse
import json
import numpy as np
from fold_model import Parameters, bridge_survival
from path_data import write_paths


def simulate(n: int, par: Parameters, *, seed: int=20261020,
             h0: float=.05, v0: float=0., horizon: float=1.,
             cutoff: float=.05, max_step: float=1/256):
    if n<1 or min(h0,horizon,cutoff,max_step)<=0 or par.d_min<1:
        raise ValueError('Require n>=1, positive scales and d_min>=1.')
    rng=np.random.default_rng(seed)
    rate=2*par.c/par.alpha*cutoff**(-par.alpha)
    h=np.full(n,h0); v=np.full(n,v0); t=np.zeros(n)
    next_jump=rng.exponential(1/rate,n)
    # -1 alive, 0 continuous contact, 1 exiting mark, 2 terminal-horizon survival.
    cause=np.full(n,-1,dtype=np.int8)
    lo=np.full(n,np.nan); hi=lo.copy(); pre_h=lo.copy(); terminal_mark=lo.copy()
    retained=np.zeros(n,dtype=np.int32)
    normal_occupation=np.zeros(n)
    rounds=0
    while True:
        idx=np.flatnonzero(cause==-1)
        if not idx.size:
            break
        rounds+=1
        if rounds>1000000:
            raise RuntimeError('Unexpectedly many time steps.')
        dt=np.minimum(max_step,np.minimum(next_jump[idx]-t[idx],horizon-t[idx]))
        if np.any(dt<=0):
            raise FloatingPointError('Nonpositive step encountered.')
        old_h=h[idx].copy(); old_v=v[idx].copy(); old_t=t[idx].copy()
        d,s,a=par.coefficients(old_v)
        sc=s*s*dt
        new_h=sc*rng.noncentral_chisquare(d,old_h/sc)
        z=np.sqrt(old_h*new_h)/sc
        bridge=bridge_survival(z,d)
        contact=rng.random(idx.size)>bridge
        # Diagnostic left-endpoint occupation only, no exact integral claim.
        normal_occupation[idx] += dt*(old_h<=.01)
        k=par.tangent_reversion
        if k>0:
            decay=np.exp(-k*dt)
            sd=par.tangent_diffusion*np.sqrt(-np.expm1(-2*k*dt)/(2*k))
            new_v=old_v*decay+sd*rng.standard_normal(idx.size)
        else:
            new_v=old_v+par.tangent_diffusion*np.sqrt(dt)*rng.standard_normal(idx.size)
        t[idx]=old_t+dt; h[idx]=new_h; v[idx]=new_v
        killed=idx[contact]
        cause[killed]=0; lo[killed]=old_t[contact]; hi[killed]=t[killed]
        pre_h[killed]=old_h[contact]; h[killed]=0.
        surv=idx[~contact]
        # next_jump is an absolute time; no rounded arrival placement is used.
        at_jump=surv[t[surv]>=next_jump[surv]-2e-14]
        if at_jump.size:
            _,_,aa=par.coefficients(v[at_jump])
            marks=cutoff*np.maximum(rng.random(at_jump.size),np.finfo(float).tiny)**(-1/par.alpha)
            marks*=np.where(rng.integers(0,2,at_jump.size)==0,-1.,1.)
            before=h[at_jump].copy()
            after=before-par.kappa*(aa*marks)**2
            retained[at_jump]+=1
            exits=after<=0
            ex=at_jump[exits]
            cause[ex]=1;lo[ex]=t[ex];hi[ex]=t[ex]
            pre_h[ex]=before[exits];terminal_mark[ex]=marks[exits]
            h[ex]=0.
            acc=at_jump[~exits]
            h[acc]=after[~exits]; v[acc]+=aa[~exits]*marks[~exits]
            next_jump[acc]=t[acc]+rng.exponential(1/rate,acc.size)
        end=surv[(cause[surv]==-1)&(t[surv]>=horizon-2e-14)]
        cause[end]=2;lo[end]=horizon;hi[end]=horizon;pre_h[end]=h[end]
    return dict(cause=cause,time_lower=lo,time_upper=hi,pre_h=pre_h,
                terminal_mark=terminal_mark,final_v=v,retained_marks=retained,
                occupation_left_001=normal_occupation),dict(
                parameters=asdict(par),n=n,seed=seed,h0=h0,v0=v0,horizon=horizon,
                cutoff=cutoff,max_step=max_step,rounds=rounds,
                cutoff_arrival_rate=rate,
                omitted_tangent_second_moment=2*par.c/(2-par.alpha)*cutoff**(2-par.alpha),
                omitted_normal_mean_upper=2*par.c/(2-par.alpha)*par.kappa*par.a_max**2*cutoff**(2-par.alpha),
                interpretation='Finite-cutoff, frozen-coefficient approximation; contact times are intervals.')


def summarize(data):
    n=len(data['cause']);out={}
    for j,name in [(0,'contact'),(1,'jump'),(2,'survival')]:
        count=int(np.sum(data['cause']==j));p=count/n
        out[name]={'count':count,'estimate':p,'standard_error':float(np.sqrt(p*(1-p)/n))}
    out['max_contact_interval']=float(np.max(data['time_upper'][data['cause']==0]-data['time_lower'][data['cause']==0],initial=0))
    return out

if __name__=='__main__':
    ap=argparse.ArgumentParser()
    ap.add_argument('--n',type=int,default=10000)
    ap.add_argument('--seed',type=int,default=20261020)
    ap.add_argument('--dimension',type=float,default=1.2)
    ap.add_argument('--cutoff',type=float,default=.05)
    ap.add_argument('--step',type=float,default=1/256)
    ap.add_argument('--output',type=Path,
                    default=Path(__file__).resolve().parents[1]/'data'/'paths')
    args=ap.parse_args()
    par=Parameters(delta0=args.dimension)
    data,metadata=simulate(args.n,par,seed=args.seed,cutoff=args.cutoff,max_step=args.step)
    args.output.parent.mkdir(parents=True,exist_ok=True)
    write_paths(args.output.with_suffix('.csv'),data)
    metadata['summary']=summarize(data)
    args.output.with_suffix('.json').write_text(json.dumps(metadata,indent=2))
    print(json.dumps(metadata['summary'],indent=2))