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

# For alpha=1/2: E_{1/2}(-z) = exp(z^2) erfc(z) = erfcx(z).
def ml_half(t, tau=1.0):
    z = np.sqrt(np.asarray(t, dtype=float) / tau)
    return erfcx(z)

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

def fit_bank(alpha, tau_target, K, max_lag=20000):
    # Fit the normalized finite-horizon causal kernel from the definition.
    lags = np.unique(np.r_[0, np.geomspace(1, max_lag, 600).astype(int)])
    raw = ml_half(lags + 1, tau_target)
    # Approximate the prefix normalizer using all integer lags.
    allraw = ml_half(np.arange(max_lag + 1) + 1, tau_target)
    target = raw / allraw.sum()
    taus = np.geomspace(1, max_lag, K)
    rho = np.exp(-1.0 / taus)
    A = ((1-rho)[:, None] * rho[:, None] ** lags[None, :]).T
    # Relative weighting prevents the long tail from being ignored.
    scale = 1.0 / np.maximum(target, 1e-12)
    try:
        w, _ = nnls(A * scale[:, None], target * scale, maxiter=100000)
    except RuntimeError:
        from scipy.optimize import lsq_linear
        w = lsq_linear(A * scale[:, None], target * scale, bounds=(0, np.inf), max_iter=10000).x
    w = w / max(w.sum(), 1e-30)
    test = np.unique(np.r_[np.arange(0, 1000), np.geomspace(1000, max_lag, 1000).astype(int)])
    yt = ml_half(test + 1, tau_target) / allraw.sum()
    yp = bank_kernel(test, taus, w)
    return dict(K=K, rel_l2=float(np.linalg.norm(yp-yt)/np.linalg.norm(yt)),
                max_rel_tail=float(np.max(np.abs(yp[-500:]-yt[-500:])/(yt[-500:]+1e-12))),
                taus=taus.tolist(), weights=w.tolist())

def normalized_filter(x, taus, weights, normalize=True):
    state=np.zeros(len(taus)); out=[]; mass=0.0
    rho=np.exp(-1/np.asarray(taus))
    for v in x:
        state=rho*state+(1-rho)*v
        mass += np.dot(weights, 1-rho) if normalize else 0.0
        out.append(np.dot(weights,state) / mass if normalize else np.dot(weights,state))
    return np.asarray(out)

def main():
    rng=np.random.default_rng(7)
    results={}
    # Prediction 1: stability and exact exponential half-life ln(.5)/ln(rho).
    stability=[]
    for rho in [0.5,0.8,0.95,0.99,0.999]:
        s=0.; peak=0.; vals=[]
        for t in range(10000):
            s=rho*s+(1-rho)*(1.0 if t==0 else 0.0)
            peak=max(peak,abs(s)); vals.append(s)
        arr=np.asarray(vals); ix=int(np.argmin(np.abs(arr-rho*0.5))) # after one? use first crossing half peak
        half=int(np.where(arr<=0.5*peak)[0][0])
        pred=math.log(.5)/math.log(rho)
        stability.append(dict(rho=rho, predicted_half_life=pred, observed_half_life=half,
                              max_abs=float(peak), final_abs=float(abs(arr[-1]))))
    results['stability_half_life']=stability
    # Prediction 2: fractional tail l^alpha h(l) -> 1/Gamma(1-alpha), alpha=.5.
    tail=[]
    for l in [100,1000,10000,100000,1000000]:
        h=float(ml_half(l+1,1.0)); tail.append(dict(l=l, scaled=float(h*math.sqrt(l)), predicted=1/math.sqrt(math.pi)))
    results['power_law_tail']=tail
    # Prediction 3: more logarithmic poles reduce approximation error.
    fits=[fit_bank(.5, 32., K) for K in [2,4,8,16,32]]
    results['bank_scaling']=fits
    # Secondary delayed-memory comparison: exact fractional normalized filter versus
    # the best one-pole normalized filter, on random sequences and a long horizon.
    T=16000; x=rng.normal(size=T); target=normalized_filter(x,[1], [1]) # placeholder overwritten
    frac=fit_bank(.5,32.,8); taus=np.asarray(frac['taus']); w=np.asarray(frac['weights'])
    y=normalized_filter(x,taus,w)
    # exact normalized causal kernel convolution, computed by direct recurrence-free convolution
    h=ml_half(np.arange(T)+1,32.)
    exact=np.convolve(x,h)[:T]/np.cumsum(h)
    one_tau=np.geomspace(1, T, 300)
    errs=[]
    for tt in one_tau:
        z=normalized_filter(x,[tt],[1.])
        errs.append(np.mean((z-exact)**2))
    best=float(min(errs)); fracerr=float(np.mean((y-exact)**2))
    results['filter_comparison']={'fractional_K8_mse':fracerr,'best_single_exponential_mse':best,
                                  'improvement_ratio':best/(fracerr+1e-30)}
    # delayed copy signal: correlation with a value 4000 steps back, highlighting retention.
    D=4000; probe=np.zeros(T); probe[D]=1.; response=normalized_filter(probe,taus,w)
    single=normalized_filter(probe,[32.],[1.])
    results['delayed_impulse']={'delay':D,'fractional_response':float(response[D]),
                                'single_tau32_response':float(single[D]),
                                'fractional_tail_at_delay':float(response[-1]),
                                'single_tail_at_delay':float(single[-1])}
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps({k:v for k,v in results.items() if k not in ('bank_scaling',)}, indent=2))
    print('bank_scaling', [(r['K'],r['rel_l2'],r['max_rel_tail']) for r in fits])

if __name__=='__main__': main()