"""Tempered-stable volatility clock MVP and numerical verification."""
import json, math
from pathlib import Path
import numpy as np


def positive_stable_laplace_one(alpha, rng, size):
    """Kanter sampler: positive stable S with E exp(-u S)=exp(-u**alpha)."""
    u = rng.uniform(1e-9, math.pi - 1e-9, size)
    e = rng.exponential(1.0, size)
    return (np.sin(alpha*u) / np.sin(u)**(1.0/alpha) *
            (np.sin((1-alpha)*u) / e)**((1-alpha)/alpha))


def ts_sample(alpha, theta, delta, rng, size):
    """Exact rejection sampler for TS(alpha,theta,delta).

    If S has Laplace exp(-s^alpha), delta^(1/alpha)S exponentially tilted
    by exp(-theta*x) has the requested tempered-stable law.
    """
    scale = delta ** (1.0/alpha)
    out = np.empty(size, dtype=float)
    filled = 0
    while filled < size:
        # Acceptance is exp(-delta*theta**alpha), generally high for settings here.
        n = max(128, int((size-filled) / max(math.exp(-delta*theta**alpha), .05) * 1.15))
        x = scale * positive_stable_laplace_one(alpha, rng, n)
        keep = rng.random(n) < np.exp(-theta*x)
        got = x[keep]
        take = min(got.size, size-filled)
        if take:
            out[filled:filled+take] = got[:take]
            filled += take
    return out


def params(alpha, theta, phi):
    delta = (1-phi) * theta**(1-alpha) / alpha
    v = (1-alpha)/(theta*(1+phi))
    return delta, v


def simulate(alpha, theta, phi, n_paths, length, rng, burn=250):
    delta, _ = params(alpha, theta, phi)
    a = np.ones(n_paths)
    # Burn-in makes the finite recursion close to stationary; mean initialization
    # also avoids an unnecessarily long burn for phi near one.
    for _ in range(burn):
        a = phi*a + ts_sample(alpha, theta, delta, rng, n_paths)
    A = np.empty((n_paths, length))
    for i in range(length):
        a = phi*a + ts_sample(alpha, theta, delta, rng, n_paths)
        A[:, i] = a
    z = rng.standard_normal((n_paths, length))
    return A, np.sqrt(A)*z


def excess_kurtosis(x):
    x = x.ravel()
    m2 = np.mean(x*x)
    return np.mean(x**4)/(m2*m2)-3


def lag_sq_acf(x, lag):
    q = x*x
    q = q - q.mean()
    return np.mean(q[:, :-lag]*q[:, lag:]) / np.mean(q*q)


def laplace_check(alpha, theta, delta, rng, n=200000):
    x = ts_sample(alpha, theta, delta, rng, n)
    rows=[]
    for u in (0.2, 1.0, 3.0):
        empirical=np.mean(np.exp(-u*x))
        theory=math.exp(-delta*((theta+u)**alpha-theta**alpha))
        rows.append({'u':u, 'empirical':empirical, 'theory':theory,
                     'abs_error':abs(empirical-theory)})
    return rows


def run():
    rng=np.random.default_rng(1142)
    # Prediction 1: stationary mean is one across persistence values.
    mean_sweep=[]
    # Prediction 2: variance / excess kurtosis follows closed form across alpha.
    var_sweep=[]
    # Prediction 3: squared-return ACF equals phi^h*v/(2+3v), with geometric decay.
    acf_sweep=[]
    for phi in (0.0, 0.3, 0.7, 0.9):
        alpha,theta=.65,.8
        A,r=simulate(alpha,theta,phi,3500,160,rng)
        delta,v=params(alpha,theta,phi)
        mean_sweep.append({'phi':phi,'pred_mean':1.0,'obs_mean':float(A.mean()),
                           'pred_var':v,'obs_var':float(A.var())})
        acf_sweep.append({'phi':phi,'pred_lag1':phi*v/(2+3*v),
                          'obs_lag1':float(lag_sq_acf(r,1)),
                          'pred_lag5':phi**5*v/(2+3*v),
                          'obs_lag5':float(lag_sq_acf(r,5))})
    for alpha in (.35,.55,.75):
        theta=.8; phi=.7
        A,r=simulate(alpha,theta,phi,3500,160,rng)
        delta,v=params(alpha,theta,phi)
        var_sweep.append({'alpha':alpha,'pred_var':v,'obs_var':float(A.var()),
                          'pred_K':3*v,'obs_K':float(excess_kurtosis(r)),
                          'delta':delta})
    # Mechanism comparison at identical dimensions: iid Gaussian has no clustering.
    alpha,theta,phi=.65,.8,.7
    A,r=simulate(alpha,theta,phi,3500,160,rng)
    iid=rng.standard_normal(r.shape)
    delta,v=params(alpha,theta,phi)
    baseline={'K':float(excess_kurtosis(iid)), 'rho1':float(lag_sq_acf(iid,1))}
    idea={'K':float(excess_kurtosis(r)), 'rho1':float(lag_sq_acf(r,1)),
          'pred_K':3*v, 'pred_rho1':phi*v/(2+3*v)}
    result={'laplace':laplace_check(.65,.8,delta,rng),
            'mean_variance_sweep':mean_sweep,'kurtosis_sweep':var_sweep,
            'acf_sweep':acf_sweep,'baseline_iid_gaussian':baseline,
            'clock':idea}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': run()