import json, math
import numpy as np
from scipy.special import erfcx
from scipy.optimize import nnls


def ml_half(t, tau):
    return erfcx(np.sqrt(np.asarray(t, dtype=float) / tau))


def bank(lags, taus, weights):
    rho = np.exp(-1 / np.asarray(taus))
    return ((weights * (1-rho))[:, None] * rho[:, None]**np.asarray(lags)[None, :]).sum(0)


def fit(K, target_tau=32., horizon=16000):
    lags = np.arange(horizon + 1)
    target = ml_half(lags + 1, target_tau)
    # The causal filter is normalized, and each exponential component has unit
    # infinite-horizon mass, so fit the same unit-mass target.
    target = target / target.sum()
    # Fit with simplex weights to preserve unit DC mass.
    taus = np.geomspace(1, horizon, K)
    A = ((1-np.exp(-1/taus))[:, None] * np.exp(-1/taus)[:, None]**lags[None, :]).T
    # Relative error across logarithmic lags, with mild absolute stabilization.
    ix = np.unique(np.r_[np.arange(100), np.geomspace(100, horizon, 500).astype(int)])
    scale = 1 / np.maximum(target[ix], 1e-8)
    w, _ = nnls(A[ix] * scale[:, None], target[ix] * scale)
    w /= w.sum()
    pred = bank(lags, taus, w)
    rel = np.linalg.norm(pred[ix]-target[ix]) / np.linalg.norm(target[ix])
    return taus, w, float(rel)


def normalized(x, taus, w):
    rho = np.exp(-1/np.asarray(taus)); s=np.zeros(len(rho)); out=[]; mass=0
    for v in x:
        s = rho*s + (1-rho)*v
        mass += np.dot(w, 1-rho)
        out.append(np.dot(w,s)/mass)
    return np.asarray(out)


def main():
    rng=np.random.default_rng(7); out={}
    # Prediction 1: rho^m decay and half-life ln(2)/(-ln rho).
    rows=[]
    for rho in [.8,.95,.99,.999]:
        impulse=(1-rho)*rho**np.arange(20000)
        obs=int(np.ceil(math.log(.5)/math.log(rho)))
        rows.append({'rho':rho,'predicted_half_life':math.log(.5)/math.log(rho),
                     'observed_half_life':obs,'max_impulse':float(impulse.max()),
                     'stable_max_state':float(np.max(np.cumsum(impulse)))})
    out['prediction_stability_half_life']=rows
    # Prediction 2: Mittag-Leffler alpha=.5 has h(l)*sqrt(l) -> 1/sqrt(pi).
    rows=[]; constant=1/math.sqrt(math.pi)
    for l in [100,1000,10000,100000,1000000]:
        scaled=float(ml_half(l,1)*math.sqrt(l))
        rows.append({'lag':l,'observed':scaled,'predicted':constant,
                     'relative_error':abs(scaled-constant)/constant})
    out['prediction_powerlaw']=rows
    # Prediction 3: increasing logarithmic poles improves fit, at least in the useful range.
    fits=[]
    for k in [1,2,4,8,16,32]:
        taus,w,e=fit(k); fits.append({'K':k,'relative_fit_error':e})
    out['prediction_bank_scaling']=fits
    # Secondary controlled signal: fractional bank versus best single exponential.
    T=16000; x=rng.normal(size=T); h=ml_half(np.arange(T)+1,32.)
    exact=np.convolve(x,h)[:T]/np.cumsum(h)
    taus,w,e=fit(8); frac=normalized(x,taus,w)
    candidates=np.geomspace(1,T,250)
    singles=[np.mean((normalized(x,[q],[1.])-exact)**2) for q in candidates]
    frac_mse=float(np.mean((frac-exact)**2)); best=float(min(singles))
    out['secondary_filter_comparison']={'fractional_K8_mse':frac_mse,
      'best_single_exponential_mse':best,'single_over_fractional':best/frac_mse}
    # Delayed retention: compare tails after a pulse at delay from the start.
    D=4000; pulse=np.zeros(T); pulse[0]=1
    fr=normalized(pulse,taus,w); one=normalized(pulse,[32],[1])
    out['secondary_delayed_impulse']={'delay':D,'fractional_at_delay':float(fr[D]),
      'single_at_delay':float(one[D]),'ratio':float(fr[D]/one[D])}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()