import json, math
import numpy as np

SEED = 937
rng = np.random.default_rng(SEED)


def burnin_prediction(r, eps, E0=1.0):
    return int(math.ceil(math.log(eps / E0) / math.log(r)))


def scalar_verification():
    rs = [0.50, 0.70, 0.85, 0.95, 0.99]
    epses = [1e-2, 1e-4, 1e-6]
    rows = []
    for r in rs:
        for eps in epses:
            pred = burnin_prediction(r, eps)
            e, k = 1.0, 0
            while e > eps and k < 100000:
                e *= r
                k += 1
            rows.append({'r': r, 'epsilon': eps, 'predicted_k': pred,
                         'observed_k': k, 'abs_error': abs(pred-k)})
    scaling = []
    for r in rs:
        x = np.log(1 / np.asarray(epses))
        y = np.asarray([burnin_prediction(r, e) for e in epses])
        slope = float(np.polyfit(x, y, 1)[0])
        target = 1 / abs(math.log(r))
        scaling.append({'r': r, 'observed_slope': slope,
                        'predicted_slope': target,
                        'relative_error': abs(slope-target)/target})
    boundary = []
    for r in [0.98, 0.999, 1.0, 1.001, 1.02]:
        e = 1.0
        for _ in range(1000):
            e *= r
        boundary.append({'r': r, 'E_1000': float(e),
                         'stable_observed': bool(e < 1),
                         'stable_predicted': bool(r < 1)})
    # Cost crossover directly evaluates the two asserted asymptotic forms.
    crossover = []
    for r in [0.5, 0.8, 0.95]:
        for eps in [1e-2, 1e-4, 1e-6]:
            obs = math.log(1/eps) / (1-r*r)
            rec = eps**(-1.0)  # d=1 toy attractor
            crossover.append({'r':r, 'epsilon':eps,
                              'observer_cost':obs, 'retrieval_cost':rec,
                              'observer_cheaper':obs < rec})
    return {'burnin_rows': rows, 'log_scaling': scaling,
            'stability_boundary': boundary, 'cost_scaling': crossover}


def lorenz_step(s, dt=0.01, sigma=10., rho=28., beta=8/3):
    def f(v):
        x, y, z = v
        return np.array([sigma*(y-x), x*(rho-z)-y, x*y-beta*z])
    k1 = f(s); k2 = f(s + dt*k1/2); k3 = f(s + dt*k2/2); k4 = f(s + dt*k3)
    return s + dt*(k1 + 2*k2 + 2*k3 + k4)/6


def generate(n=2600, dt=0.01):
    x = np.array([1., 1., 1.])
    truth = []
    for _ in range(n):
        x = lorenz_step(x, dt)
        truth.append(x.copy())
    truth = np.asarray(truth)
    # Scalar observation with moderate noise and occasional mismatch bursts.
    noise = rng.normal(0, 0.7, n)
    y = truth[:, 0] + noise
    bad = np.zeros(n, dtype=bool)
    bad[800:900] = True; bad[1700:1780] = True
    y[bad] += rng.normal(0, 7.0, bad.sum())
    return truth, y, bad


def run_policy(truth, y, mode, threshold=2.2, rmax=.94, eps=1e-3):
    # Deliberately small linear observer: x prediction plus scalar x correction.
    # Retrieval is a nearest scalar-observation context from a bank of latent states.
    z = truth[0].copy() + np.array([5., -3., 2.])
    bank_y, bank_z = [], []
    sqerr, retrievals, residuals, costs = [], 0, [], 0.0
    bad_retrievals = 0
    for t in range(len(y)):
        # One cheap model rollout. Its local contraction proxy is fixed but
        # varies by phase, making the switch's decision rule testable.
        r = .78 if (t % 700) < 500 else .985
        pred = lorenz_step(z)
        q = abs(y[t] - pred[0])
        use_retrieval = mode == 'retrieve' or (mode == 'adaptive' and (r > rmax or q > threshold))
        if mode == 'observer' or (mode == 'adaptive' and not use_retrieval):
            z = pred.copy()
            # Stable scalar innovation update, a simple observer correction.
            z[0] += .18 * (y[t] - z[0])
            costs += 1.0
        else:
            if bank_y:
                j = int(np.argmin(np.abs(np.asarray(bank_y[-700:]) - y[t])))
                z = bank_z[-700:][j].copy()
            else:
                z = pred.copy()
            z[0] += .10 * (y[t] - z[0])
            retrievals += 1; bad_retrievals += int(abs(y[t] - truth[t, 0]) > 4.0); costs += 7.0
        bank_y.append(float(y[t])); bank_z.append(z.copy())
        residuals.append(q); sqerr.append(float(np.mean((z-truth[t])**2)))
    return {'mse':float(np.mean(sqerr)), 'rmse':float(np.sqrt(np.mean(sqerr))),
            'retrievals':retrievals, 'retrieval_rate':retrievals/len(y),
            'relative_cost':costs/len(y), 'mean_residual':float(np.mean(residuals)),
            'bad_period_retrievals': int(bad_retrievals)}


def main():
    scalar = scalar_verification()
    truth, y, bad = generate()
    results = {}
    for mode in ['observer', 'retrieve', 'adaptive']:
        results[mode] = run_policy(truth, y, mode)
    # Useful diagnostic specifically measures event sensitivity.
    z = truth[0] + np.array([5., -3., 2.])
    hits_good = hits_bad = 0
    for t in range(len(y)):
        r = .78 if (t % 700) < 500 else .985
        q = abs(y[t] - lorenz_step(z)[0])
        decision = r > .94 or q > 2.2
        if bad[t]: hits_bad += int(decision)
        else: hits_good += int(decision)
        z = lorenz_step(z); z[0] += .18*(y[t]-z[0])
    results['adaptive_event_diagnostics'] = {
        'bad_steps': int(bad.sum()), 'bad_switches': hits_bad,
        'bad_recall': hits_bad/max(1,int(bad.sum())),
        'good_switch_rate': hits_good/max(1,int((~bad).sum()))}
    out = {'seed': SEED, 'scalar_verification': scalar,
           'lorenz_results': results,
           'note': 'Toy observer/retrieval costs are normalized units; retrieval cost is 7x observer.'}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()