import json, math
from pathlib import Path
import numpy as np

SEED = 1729
EPS = 0.10
PERSISTENCE = 3
RIDGE = 1e-4
CANDIDATES = [1, 2, 4, 8, 16, 32]


def closure_rho(z, max_lag=100):
    """Ensemble/time normalized autocorrelation, avoiding sequence concatenation."""
    z = np.asarray(z, float)
    z = z - np.mean(z)
    den = np.mean(z * z)
    return np.array([1.0 if k == 0 else np.mean(z[:, :-k] * z[:, k:]) / den
                     for k in range(max_lag)])


def first_persistent_crossing(rho, eps=EPS, persistence=PERSISTENCE):
    for k in range(1, len(rho) - persistence + 1):
        if np.all(np.abs(rho[k:k + persistence]) <= eps):
            return k
    return len(rho) - persistence


def make_data(a, nseq=240, length=240, noise=0.65, seed=SEED):
    rng = np.random.default_rng(seed)
    z = rng.normal(size=(nseq, length))
    for t in range(1, length):
        z[:, t] = a * z[:, t - 1] + math.sqrt(1 - a * a) * z[:, t]
    # x is a resolved signal with unresolved closure injection z observed noisily.
    x = z + noise * rng.normal(size=z.shape)
    return x, z


def fit_history_predictor(x, m, frac=.65):
    cut = int(len(x) * frac)

    def xy(arr):
        X, Y = [], []
        for row in arr:
            for t in range(m - 1, len(row) - 1):
                X.append(row[t - m + 1:t + 1][::-1])
                Y.append(row[t + 1])
        return np.asarray(X), np.asarray(Y)

    X, Y = xy(x[:cut])
    Xt, Yt = xy(x[cut:])
    w = np.linalg.solve(X.T @ X + RIDGE * np.eye(m), X.T @ Y)
    return w, float(np.mean((Xt @ w - Yt) ** 2))


def rollout_error(x, w, m, horizon=30):
    vals = []
    for row in x:
        t0 = len(row) - horizon - 1
        hist = list(row[t0 - m + 1:t0 + 1])
        pred = []
        for _ in range(horizon):
            y = float(np.dot(w, np.asarray(hist[-m:][::-1])))
            pred.append(y)
            hist.append(y)
        vals.append(np.mean((np.asarray(pred) - row[t0 + 1:t0 + 1 + horizon]) ** 2))
    return float(np.mean(vals))


def theoretical_crossing(a):
    return int(math.ceil(math.log(EPS) / math.log(abs(a))))


def theoretical_tau(a):
    # Discrete analogue of integral |rho| for rho(k)=a^k, including k=0.
    return 1.0 / (1.0 - abs(a))


def run():
    rows = []
    for a in (.2, .5, .7, .8, .9):
        x, z = make_data(a)
        rho = closure_rho(z)
        observed = first_persistent_crossing(rho)
        pred = theoretical_crossing(a)
        tau_obs = float(np.sum(np.abs(rho)))
        tau_pred = theoretical_tau(a)
        results = []
        for m in CANDIDATES:
            w, one = fit_history_predictor(x, m)
            results.append({'m': m, 'one_step': one, 'rollout': rollout_error(x, w, m)})
        scheduled = max(1, min(32, observed))
        ws, one = fit_history_predictor(x, scheduled)
        best = min(results, key=lambda q: q['one_step'])
        # Smallest window retaining at least 95% of the improvement from m=1 to m=32.
        e1, emax = results[0]['one_step'], results[-1]['one_step']
        target = e1 - .95 * (e1 - emax)
        elbow = next(q['m'] for q in results if q['one_step'] <= target)
        rows.append({'a': a, 'predicted_crossing': pred, 'observed_crossing': observed,
                     'predicted_tau': tau_pred, 'observed_tau': tau_obs,
                     'scheduled_m': scheduled, 'elbow_m_95pct': elbow, 'best_m': best['m'],
                     'results': results, 'scheduled_one_step': one,
                     'scheduled_rollout': rollout_error(x, ws, scheduled)})

    cross_err = [abs(q['observed_crossing'] - q['predicted_crossing']) for q in rows]
    tau_rel_err = [abs(q['observed_tau'] - q['predicted_tau']) / q['predicted_tau'] for q in rows]
    # Compare adaptive scheduler to Markov baseline and full 32-token model.
    baseline = float(np.mean([q['results'][0]['one_step'] for q in rows]))
    idea = float(np.mean([q['scheduled_one_step'] for q in rows]))
    full = float(np.mean([q['results'][-1]['one_step'] for q in rows]))
    summary = {
        'crossing_mae_steps': float(np.mean(cross_err)),
        'crossing_max_error_steps': int(max(cross_err)),
        'tau_relative_mae': float(np.mean(tau_rel_err)),
        'mean_baseline_m1_loss': baseline,
        'mean_scheduler_loss': idea,
        'mean_full_m32_loss': full,
        'scheduler_vs_m1_percent': 100 * (baseline - idea) / baseline,
        'scheduler_vs_m32_percent': 100 * (idea - full) / full,
        'mean_scheduler_m': float(np.mean([q['scheduled_m'] for q in rows])),
        'mean_elbow_m': float(np.mean([q['elbow_m_95pct'] for q in rows]))
    }
    out = {'config': {'epsilon': EPS, 'persistence': PERSISTENCE, 'seed': SEED,
                      'candidates': CANDIDATES}, 'summary': summary, 'rows': rows}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    run()