"""MVP for safe receding-horizon topology switching.

The toy has correlated features x2=x1+noise.  The deployed model uses both
features with weights (1,1), while the target topology drops feature 2 and
uses target weight (2,0).  The endpoint is accurate, but immediately after
changing the mask, retaining the old active weight gives a large error.
"""
import json
import math
from dataclasses import dataclass
import numpy as np

SEED = 217

@dataclass
class Toy:
    x1: np.ndarray
    x2: np.ndarray
    y: np.ndarray

    @property
    def a(self):
        # residual after topology switch and interpolation alpha:
        # (1+alpha)x1 - (x1+x2) = alpha*x1-x2
        return float(np.mean(self.x1 * self.x1))
    @property
    def b(self):
        return float(np.mean(self.x1 * self.x2))
    @property
    def c(self):
        return float(np.mean(self.x2 * self.x2))
    def loss_alpha(self, alpha):
        return self.a * alpha * alpha - 2*self.b*alpha + self.c
    def old_loss(self):
        return float(np.mean((self.x1 + self.x2 - self.y)**2))
    def target_loss(self):
        return self.loss_alpha(1.0)
    def predicted_root(self, eps):
        # smallest alpha in [0,1] with loss_alpha <= eps
        # positive root of a alpha^2 - 2b alpha + c-eps=0
        d = self.b*self.b - self.a*(self.c-eps)
        if d < 0: return None
        roots = [(self.b-math.sqrt(max(0,d)))/self.a,
                 (self.b+math.sqrt(max(0,d)))/self.a]
        feasible = [r for r in roots if 0 <= r <= 1]
        if self.loss_alpha(0) <= eps: return 0.0
        return min(feasible) if feasible else float('inf')

def make_toy(noise, n=20000):
    rng = np.random.default_rng(SEED + int(round(noise*1000)))
    x1 = rng.normal(size=n)
    x2 = x1 + noise*rng.normal(size=n)
    return Toy(x1, x2, x1+x2)

def certify(toy, alpha, eps):
    return toy.loss_alpha(alpha) - toy.old_loss() <= eps + 1e-12

def filtered_switch(toy, eps, candidates):
    rejected = []
    for alpha in candidates:
        if certify(toy, alpha, eps):
            return alpha, rejected
        rejected.append(alpha)
    return None, rejected

def run():
    # Core math verification: sweep noise and tolerance, compare quadratic root
    # to the observed first accepted planner candidate.
    rows = []
    candidate_grid = np.linspace(0.0, 1.0, 101)
    for noise in [0.05, 0.10, 0.20, 0.35]:
        toy = make_toy(noise)
        for eps in [0.10, 0.25, 0.50]:
            root = toy.predicted_root(eps)
            observed, rej = filtered_switch(toy, eps, candidate_grid)
            rows.append({
                'noise': noise, 'epsilon': eps,
                'predicted_alpha_min': root,
                'observed_alpha_min_grid': observed,
                'grid_error': None if observed is None or root is None else observed-root,
                'rejected_candidates': len(rej),
                'endpoint_loss': toy.target_loss(),
                'initial_new_mask_loss': toy.loss_alpha(0),
            })

    # Prediction 1: residual follows the derived quadratic to numerical precision.
    toy = make_toy(0.20)
    alphas = np.linspace(0,1,51)
    empirical = np.array([toy.loss_alpha(a) for a in alphas])
    formula = toy.a*alphas**2 - 2*toy.b*alphas + toy.c
    quadratic_max_abs_error = float(np.max(np.abs(empirical-formula)))

    # Prediction 2: larger tolerance monotonically lowers the required transition.
    roots = [toy.predicted_root(e) for e in [0.10,0.25,0.50]]
    tolerance_monotone = bool(all(roots[i] >= roots[i+1] for i in range(2)))

    # Baselines at the same candidate schedule. One-shot has an unsafe
    # intermediate (alpha=0 under new mask); unfiltered gradual schedule also
    # violates until its first safe alpha. The filter executes only safe alpha.
    eps = 0.25
    one_shot = make_toy(0.20)
    gradual = np.linspace(0,1,11)
    unfiltered_violations = int(sum(not certify(one_shot,a,eps) for a in gradual))
    accepted, rejected = filtered_switch(one_shot, eps, gradual)
    filtered_violations = 0 if accepted is not None and certify(one_shot, accepted, eps) else 1

    # Cut reuse: without cuts, each replan retries all unsafe candidates;
    # with cuts, rejected alpha values are remembered and not evaluated again.
    unsafe = [float(a) for a in gradual if not certify(one_shot,a,eps)]
    replans = 5
    no_cut_evals = replans * len(unsafe)
    with_cut_evals = len(unsafe)

    out = {
        'seed': SEED,
        'quadratic_max_abs_error': quadratic_max_abs_error,
        'predicted_vs_observed': rows,
        'prediction_checks': {
            'quadratic_formula_exact': quadratic_max_abs_error < 1e-12,
            'tolerance_lowers_required_alpha': tolerance_monotone,
            'cut_reuse_reduces_rejected_evaluations': with_cut_evals < no_cut_evals,
        },
        'comparison': {
            'epsilon': eps,
            'one_shot_endpoint_loss': one_shot.target_loss(),
            'one_shot_intermediate_violations': 1 if not certify(one_shot,0,eps) else 0,
            'unfiltered_gradual_violations_of_11': unfiltered_violations,
            'filtered_first_accepted_alpha': accepted,
            'filtered_rejected_before_accept': len(rejected),
            'filtered_accepted_step_violations': filtered_violations,
            'cut_reuse_evaluations_over_5_replans': {'without_cuts': no_cut_evals, 'with_cuts': with_cut_evals},
        },
    }
    print(json.dumps(out, indent=2, sort_keys=True))

if __name__ == '__main__':
    run()