"""Curved-fold jump model: exact generator coefficients and analytical bounds.

The Lévy density is c*|z|**(-1-alpha) on the entire punctured real line.
No finite mark cutoff is used in these generator and probability calculations.
"""
from __future__ import annotations
from dataclasses import dataclass, asdict
import math
import numpy as np
from numpy.typing import ArrayLike, NDArray
from scipy.special import gammaln, digamma, ive, kve
from scipy.integrate import quad

@dataclass(frozen=True)
class Parameters:
    alpha: float = 1.2
    c: float = 0.12
    kappa: float = 0.25
    delta0: float = 1.2
    delta_modulation: float = 0.2
    sigma_modulation: float = 0.15
    amplitude_modulation: float = 0.1
    tangent_diffusion: float = 0.4
    tangent_reversion: float = 1.0

    def __post_init__(self) -> None:
        if not all(math.isfinite(x) for x in asdict(self).values()):
            raise ValueError('All parameters must be finite.')
        if not 0 < self.alpha < 2 or self.c <= 0 or self.kappa <= 0:
            raise ValueError('Require 0 < alpha < 2 and positive c, kappa.')
        if self.delta0 < abs(self.delta_modulation):
            raise ValueError('Dimension coefficient must be nonnegative.')
        if abs(self.sigma_modulation) >= 1 or abs(self.amplitude_modulation) >= 1:
            raise ValueError('Noise amplitudes must be uniformly positive.')
        if self.tangent_diffusion < 0 or self.tangent_reversion < 0:
            raise ValueError('Nonnegative tangent coefficients are required.')

    @property
    def q(self) -> float:
        return self.alpha / 2

    @property
    def d_min(self) -> float:
        return self.delta0 - abs(self.delta_modulation)

    @property
    def d_max(self) -> float:
        return self.delta0 + abs(self.delta_modulation)

    @property
    def s_min(self) -> float:
        return 1 - abs(self.sigma_modulation)

    @property
    def s_max(self) -> float:
        return 1 + abs(self.sigma_modulation)

    @property
    def a_min(self) -> float:
        return 1 - abs(self.amplitude_modulation)

    @property
    def a_max(self) -> float:
        return 1 + abs(self.amplitude_modulation)

    @property
    def c_min(self) -> float:
        return self.c * self.kappa**self.q * self.a_min**self.alpha

    @property
    def c_max(self) -> float:
        return self.c * self.kappa**self.q * self.a_max**self.alpha

    def coefficients(self, v: ArrayLike):
        t = np.tanh(np.asarray(v, dtype=float))
        return (self.delta0 + self.delta_modulation*t,
                1+self.sigma_modulation*t,
                1+self.amplitude_modulation*t)


def mellin_coefficient(p: float, q: float) -> float:
    """m_q(p), restricted to its positive range p > q-1."""
    if not 0 < q < 1 or not p > q-1 or not math.isfinite(p):
        raise ValueError('Require 0 < q < 1 and p > q-1.')
    return float(np.exp(gammaln(1-q)+gammaln(1+p)-gammaln(1+p-q))/q)


def mellin_quadrature(p: float, q: float) -> float:
    """Independent real integral, with endpoint subtraction at zero."""
    if not 0 < q < 1 or not p > q-1:
        raise ValueError('Require 0 < q < 1 and p > q-1.')
    # Integrate after removing the linear term: cancellation at zero is avoided.
    def remainder(u: float) -> float:
        if u == 0:
            return 0.0
        if u < 1e-5:
            r = (p*(p-1)/2*u*u - p*(p-1)*(p-2)/6*u**3
                 +p*(p-1)*(p-2)*(p-3)/24*u**4)
        else:
            r = math.expm1(p*math.log1p(-u)) + p*u
        return r*u**(-1-q)
    val = quad(remainder, 0, .5, epsabs=2e-11, epsrel=2e-11)[0]
    # Split upper endpoint so p<0 is handled as an integrable power.
    val += quad(remainder, .5, 1, epsabs=2e-10, epsrel=2e-10, limit=300)[0]
    return float(1/q+p/(1-q)-val)


def power_drift_constant(p: float, par: Parameters) -> float:
    """A positive uniform a_p, or raise if the proposed power is inadmissible."""
    if not par.q-1 < p < 1 or p == 0:
        raise ValueError('p must be nonzero and lie in (q-1,1).')
    if p > 0:
        result = par.s_min**2*p*(2-par.d_max-2*p)
    else:
        theta = -p
        result = par.s_min**2*theta*(par.d_min-2-2*theta)
    if result <= 0:
        raise ValueError('The requested exponent does not give a negative continuous drift.')
    return result


def power_generator(h: ArrayLike, v: ArrayLike, p: float, par: Parameters,
                    censored: bool = False) -> NDArray[np.float64]:
    h, v = np.broadcast_arrays(np.asarray(h, float), np.asarray(v, float))
    if np.any(h <= 0) or np.any(~np.isfinite(h)):
        raise ValueError('h must be finite and positive.')
    d, s, a = par.coefficients(v)
    c_v = par.c*par.kappa**par.q*a**par.alpha
    m = mellin_coefficient(p, par.q) - (1/par.q if censored else 0)
    return s*s*p*(d+2*p-2)*h**(p-1) - c_v*m*h**(p-par.q)


def jump_rate(h: ArrayLike, v: ArrayLike, par: Parameters) -> NDArray[np.float64]:
    h, v = np.broadcast_arrays(np.asarray(h, float), np.asarray(v, float))
    if np.any(h <= 0):
        raise ValueError('Positive discriminants are required.')
    _, _, a = par.coefficients(v)
    return par.c*par.kappa**par.q*a**par.alpha/par.q*h**(-par.q)


def optimized_jump_bound(h0: float, p: float, par: Parameters) -> tuple[float, float]:
    """Upper bound on eventual jump termination; returns (bound, optimal collar)."""
    if h0 <= 0 or not 0 < p < 1-par.q:
        raise ValueError('Require h0>0 and 0<p<1-q.')
    ap = power_drift_constant(p, par)
    b = 1-p-par.q
    A = par.c_max/(par.q*ap)
    Lstar = (p/(A*b))**(1/(1-par.q))
    raw = h0**p*(1-par.q)/b*(A*b/p)**(p/(1-par.q))
    return min(1.0, raw), Lstar


def collar_contact_bound(h0: float, L: float, p: float, par: Parameters) -> float:
    if not 0 < h0 < L or not 0 < p < 1-par.q:
        raise ValueError('Require 0<h0<L and 0<p<1-q.')
    ap = power_drift_constant(p, par)
    return max(0., 1-(h0/L)**p-par.c_max*h0**p*L**(1-p-par.q)/(par.q*ap))


def occupation_bound(h0: float, radius: ArrayLike, p: float, par: Parameters):
    if h0 <= 0 or np.any(np.asarray(radius) <= 0):
        raise ValueError('Positive initial state and radii are required.')
    ap = power_drift_constant(p, par)
    return h0**p/ap*np.asarray(radius)**(2*(1-p))


def near_jump_bound(h0: float, radius: ArrayLike, p: float, par: Parameters):
    if p >= 1-par.q:
        raise ValueError('Require p<1-q.')
    return np.minimum(1., par.c_max/(par.q)*
                      occupation_bound(h0, radius, p, par)*np.asarray(radius)**(-par.alpha))


def critical_barrier_constants(par: Parameters) -> tuple[float, float]:
    if par.d_min < 2:
        raise ValueError('The critical barrier requires d_min>=2.')
    r = 1-par.q
    Aq = (digamma(1)-digamma(1-par.q))/par.q
    Dq = mellin_coefficient(r, par.q)-1/par.q
    L = min(1., (par.s_min**2*r*r/(par.c_max*Dq))**(1/r))
    K = par.c_max*Aq/(par.s_min**2*r*r)
    return float(L), float(K)


def critical_barrier_generator(h: ArrayLike, v: ArrayLike, par: Parameters):
    L, K = critical_barrier_constants(par)
    h, v = np.broadcast_arrays(np.asarray(h, float), np.asarray(v, float))
    if np.any(h <= 0) or np.any(h > L*(1+1e-12)):
        raise ValueError('h must lie in the proven collar.')
    d,s,a = par.coefficients(v)
    Cv = par.c*par.kappa**par.q*a**par.alpha
    Aq = (digamma(1)-digamma(1-par.q))/par.q
    return s*s*(2-d)/h+Cv*Aq*h**(-par.q)-K*power_generator(h,v,1-par.q,par,True)


def survival_bound(T: ArrayLike, h0: float, par: Parameters):
    T = np.asarray(T, float)
    if np.any(T<=0) or h0<=0:
        raise ValueError('Require T>0 and h0>0.')
    gamma = (1-par.q)/(1+par.q)
    a0 = par.c_min/par.q
    c0 = a0**(1/(1+par.q))*(4*par.s_max**2)**(-par.q/(1+par.q))
    # logaddexp avoids overflow in a bound that may be larger than one.
    logpref = np.logaddexp(0., par.d_max/2*np.log(2)+h0/(2*par.s_max**2*T))
    return np.exp(np.minimum(0.,logpref-c0*T**gamma))


def post_jump(u: ArrayLike, v: ArrayLike, z: ArrayLike, par: Parameters):
    u,v,z=np.broadcast_arrays(np.asarray(u,float),np.asarray(v,float),np.asarray(z,float))
    _,_,a=par.coefficients(v)
    up=u+2*par.kappa*v*a*z
    vp=v+a*z
    return up,vp


def bridge_survival(z: ArrayLike, dimension: ArrayLike):
    """Exact Bessel bridge probability for frozen coefficients, dimension>=1."""
    z,d=np.broadcast_arrays(np.asarray(z,float),np.asarray(dimension,float))
    if np.any(z<0) or np.any(d<1) or np.any(~np.isfinite(z)):
        raise ValueError('Require finite z>=0 and dimension>=1.')
    out=np.ones_like(z)
    sub=d<2
    out[sub & (z==0)]=0
    mask=sub & (z>0) & (z<350)
    zz=z[mask]; eta=1-d[mask]/2
    if zz.size:
        lr=(np.log(2/np.pi*np.sin(np.pi*eta))+np.log(kve(eta,zz))
            -np.log(ive(eta,zz))-2*zz)
        out[mask]=np.exp(-np.logaddexp(0.,lr))
    if np.any(~np.isfinite(out)):
        raise FloatingPointError('Bessel special-function evaluation failed.')
    return out