import json, math, random
from pathlib import Path
import numpy as np
from scipy.optimize import nnls
from scipy.special import gamma


def fit_bank(p, smin=1.0, smax=2048.0, J=24):
    lam = np.logspace(-math.log10(smax), -math.log10(smin), J)
    grid = np.geomspace(smin, smax, 500)
    target = grid ** (p - 1.0)
    A = np.exp(-grid[:, None] * lam[None, :])
    w, _ = nnls(A, target)
    pred = A @ w
    scale = np.exp(np.mean(np.log(target + 1e-30) - np.log(pred + 1e-30)))
    return lam, w * scale


def bank_impulse(lam, w, n=4096, dt=1.0):
    return np.exp(-np.outer(np.arange(n) * dt, lam)) @ w


def bank_transfer(lam, w, omega):
    return np.sum(w[None, :] / (lam[None, :] + 1j * omega[:, None]), axis=1)


def mechanism_checks():
    rows = []
    for p in [0.2, 0.5, 0.8]:
        lam, w = fit_bank(p)
        g = bank_impulse(lam, w)
        lags = np.arange(8, 512)
        slope = np.polyfit(np.log(lags), np.log(g[lags]), 1)[0]
        rows.append({"check": "power_law_slope", "p": p,
                     "predicted": p - 1.0, "observed": float(slope),
                     "abs_error": float(abs(slope - (p - 1.0))),
                     "positive_weights": bool(np.all(w >= 0))})

    # Exact discretization predicts a spectral radius exp(-lambda_min*dt),
    # strictly below one for positive rates, for every tested time step.
    lam, w = fit_bank(0.5)
    for dt in [0.1, 1.0, 4.0]:
        rho = float(np.max(np.exp(-lam * dt)))
        predicted = float(np.exp(-np.min(lam) * dt))
        rows.append({"check": "discrete_stability", "dt": dt,
                     "predicted_rho": predicted, "observed_rho": rho,
                     "stable": rho < 1.0})

    # For Q(w)=integral exp(-lambda s)x(t-s) ds, a fractional kernel has
    # Q phase approximately -pi*p/2 and magnitude slope -p. This is the
    # Fourier-sign counterpart of the paper's C_p phase +pi*p/2.
    for p in [0.2, 0.5, 0.8]:
        lam, w = fit_bank(p)
        omega = np.logspace(-2.3, -0.7, 300)
        H = bank_transfer(lam, w, omega)
        sel = (omega > 0.004) & (omega < 0.2)
        observed_phase = float(np.median(np.unwrap(np.angle(H))[sel]))
        observed_slope = float(np.polyfit(np.log(omega[sel]),
                                           np.log(np.abs(H[sel])), 1)[0])
        phase_pred = -math.pi * p / 2.0
        rows.append({"check": "fractional_transfer", "p": p,
                     "predicted_phase_rad": phase_pred,
                     "observed_phase_rad": observed_phase,
                     "phase_error_rad": float(abs(observed_phase-phase_pred)),
                     "predicted_log_slope": -p,
                     "observed_log_slope": observed_slope,
                     "slope_error": float(abs(observed_slope+p))})
    return rows

def fractional_filter(x, p=0.5, smax=128, J=16):
    lam, w = fit_bank(p, 1.0, float(smax), J)
    q = np.zeros(J)
    out = np.zeros(len(x))
    decay = np.exp(-lam)
    gain = (1.0 - decay) / lam
    a = np.sum(w / lam)
    for t, xt in enumerate(x):
        q = decay * q + gain * xt
        out[t] = a * xt - np.dot(w, q)
    return out


def delayed_retrieval(seed=7, ntrain=6000, ntest=3000, delay=48):
    rng = np.random.default_rng(seed)
    # Isolated random pulses make the target unambiguously depend on a past
    # event. Evaluate the memory feature exactly at the requested delay.
    gap = delay + 1
    def make(n):
        x = np.zeros(n); y = np.zeros(n)
        for t in range(0, n-gap, gap):
            bit = rng.integers(0, 2)
            x[t] = bit
            y[t+delay] = bit
        return x, y
    train_x, train_y = make(ntrain)
    test_x, test_y = make(ntest)
    result = {}
    for kind in ["fractional", "exp", "raw"]:
        if kind == "fractional":
            z = fractional_filter(train_x, 0.5, max(128, delay * 3), 16)
            zt = fractional_filter(test_x, 0.5, max(128, delay * 3), 16)
        elif kind == "exp":
            lam = math.log(2) / delay
            def filt(x):
                q = 0.; out = np.zeros(len(x))
                e = math.exp(-lam)
                for t, xt in enumerate(x):
                    q = e*q + (1-e)/lam*xt
                    out[t] = q
                return out
            z, zt = filt(train_x), filt(test_x)
        else:
            z, zt = train_x, test_x
        X = np.column_stack([np.ones(len(z)), z])
        beta = np.linalg.lstsq(X, train_y, rcond=None)[0]
        pred = np.column_stack([np.ones(len(zt)), zt]) @ beta
        mask = test_y != 0
        result[kind] = float(np.mean((pred[mask] - test_y[mask]) ** 2))
    return result

def main():
    random.seed(0); np.random.seed(0)
    checks = mechanism_checks()
    retrieval = delayed_retrieval()
    report = {"checks": checks, "retrieval_mse": retrieval}
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()