from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'code'))
import math
import numpy as np
import pytest
from scipy.optimize import minimize_scalar
from fold_model import *
from simulate_fold import simulate, summarize

@pytest.mark.parametrize('q,p',[(.2,-.3),(.2,.1),(.6,-.15),(.6,.2),(.85,-.05),(.85,.1),(.6,0)])
def test_fractional_integral(q,p):
    assert math.isclose(mellin_coefficient(p,q),mellin_quadrature(p,q),rel_tol=2e-8,abs_tol=2e-9)

@pytest.mark.parametrize('p',[-.1,.2,.5])
def test_killed_censored_decomposition(p):
    par=Parameters()
    h=np.geomspace(1e-7,10,100);v=np.linspace(-4,4,100)
    a=power_generator(h,v,p,par)
    b=power_generator(h,v,p,par,True)-jump_rate(h,v,par)*h**p
    np.testing.assert_allclose(a,b,rtol=2e-13)

@pytest.mark.parametrize('dimension,p',[(1.2,.2),(3.,-.1)])
def test_nonlocal_lyapunov(dimension,p):
    par=Parameters(delta0=dimension);ap=power_drift_constant(p,par)
    h=np.geomspace(1e-10,1e5,701)[:,None];v=np.linspace(-20,20,63)[None,:]
    ratio=power_generator(h,v,p,par)/h**(p-1)
    assert np.max(ratio)<=-ap+1e-11

@pytest.mark.parametrize('q',[.2,.6,.85])
def test_critical_barrier(q):
    par=Parameters(alpha=2*q,delta0=2.2)
    L,K=critical_barrier_constants(par)
    h=np.geomspace(L*1e-8,L,300)[:,None];v=np.linspace(-25,25,51)[None,:]
    assert np.max(critical_barrier_generator(h,v,par)*h**q)<=1e-8

@pytest.mark.parametrize('curvature',[.01,.25,2.])
def test_geometry(curvature):
    par=Parameters(kappa=curvature)
    rng=np.random.default_rng(721)
    v=rng.normal(size=100);h=rng.uniform(.01,2,100);z=rng.normal(size=100)
    u=h+par.kappa*v*v
    up,vp=post_jump(u,v,z,par)
    _,_,a=par.coefficients(v)
    np.testing.assert_allclose(up-par.kappa*vp*vp,h-par.kappa*a*a*z*z,rtol=1e-12,atol=1e-12)
    np.testing.assert_allclose((up-u)-2*par.kappa*v*(vp-v),0,atol=2e-12)

@pytest.mark.parametrize('curvature',[.001,.05,1.])
def test_optimized_bound(curvature):
    par=Parameters(kappa=curvature)
    p=.2;h0=1e-8;ap=power_drift_constant(p,par)
    bound,L=optimized_jump_bound(h0,p,par)
    A=par.c_max/(par.q*ap)
    f=lambda x: h0**p*(math.exp(-p*x)+A*math.exp((1-p-par.q)*x))
    res=minimize_scalar(f,bounds=(math.log(L)-5,math.log(L)+5),method='bounded')
    assert math.isclose(bound,res.fun,rel_tol=2e-10)
    assert math.isclose(math.log(L),res.x,abs_tol=2e-5)

@pytest.mark.parametrize('bad',[-1,0,2,3])
def test_invalid_alpha(bad):
    with pytest.raises(ValueError): Parameters(alpha=bad)

@pytest.mark.parametrize('dimension',[1.,1.6,2.,3.])
def test_bridge_probability(dimension):
    x=np.geomspace(1e-10,1000,100)
    q=bridge_survival(x,dimension)
    assert np.all((q>=0)&(q<=1))
    assert np.all(np.diff(q)>=-1e-14)
    if dimension>=2: np.testing.assert_array_equal(q,np.ones_like(q))

@pytest.mark.parametrize('dimension',[1.2,3.])
def test_simulator_records(dimension):
    par=Parameters(delta0=dimension)
    data,meta=simulate(300,par,seed=23,cutoff=.1,max_step=.05,horizon=.2)
    assert sum(summarize(data)[s]['count'] for s in ['contact','jump','survival'])==300
    assert np.all(data['time_lower']<=data['time_upper'])
    assert np.max(data['time_upper']-data['time_lower'])<=.05+1e-12
    if dimension==3.: assert not np.any(data['cause']==0)
    j=data['cause']==1
    # The fatal mark need only be tested with its own pre-state amplitude;
    # the stored final tangent state equals the pre-jump state on a fatal mark.
    _,_,a=par.coefficients(data['final_v'][j])
    assert np.all(data['pre_h'][j]-par.kappa*(a*data['terminal_mark'][j])**2<=0)

def test_curvature_probability_exponent():
    p=.2;h0=1e-10
    a,_=optimized_jump_bound(h0,p,Parameters(kappa=.01))
    b,_=optimized_jump_bound(h0,p,Parameters(kappa=.04))
    assert math.isclose(b/a,4**(.6*p/.4),rel_tol=1e-12)

def test_survival_tail_is_integrable():
    par=Parameters()
    T=np.geomspace(1e-3,1e16,500)
    b=survival_bound(T,.05,par)
    assert np.all(np.isfinite(b)) and np.all((b>=0)&(b<=1))
    assert b[-1]<1e-20