import json
import numpy as np


def opnorm(a):
    return np.linalg.norm(a, 2)


def gram(stack, p=None):
    if p is None:
        p = stack[0].shape[1]
    return sum((w.T @ w for w in stack), np.zeros((p, p)))


def theorem_bound(true_stack, errors):
    return sum(2 * opnorm(w) * e + e * e for w, e in zip(true_stack, errors))


def computable_bound(estimated_stack, errors):
    # Since ||Omega|| <= ||Omega_hat|| + ||Delta||, this is a valid
    # estimated-state-only inflation of the theorem bound.
    return sum(2 * (opnorm(h) + e) * e + e * e
               for h, e in zip(estimated_stack, errors))


def greedy_stack(candidates, max_n=8):
    selected, selected_ids, current = [], [], -np.inf
    p = candidates[0].shape[1]
    for i, w in enumerate(candidates):
        trial = selected + [w]
        score = np.linalg.eigvalsh(gram(trial, p))[0]
        if len(selected) < p or score > current + 1e-10:
            selected, selected_ids, current = trial, selected_ids + [i], score
        if len(selected) >= max_n:
            break
    return selected, selected_ids


def mechanism_checks(rng):
    p, m = 2, 18
    true = [rng.normal(size=(3, p)) for _ in range(m)]
    true = [w / opnorm(w) for w in true]

    rows = []
    violations = 0
    conservative_violations = 0
    for e in np.geomspace(1e-5, 0.8, 12):
        estimated, errors = [], []
        for w in true:
            d = rng.normal(size=w.shape); d *= e / opnorm(d)
            estimated.append(w + d); errors.append(e)
        diff = opnorm(gram(estimated) - gram(true))
        b = theorem_bound(true, errors)
        bc = computable_bound(estimated, errors)
        violations += int(diff > b * (1 + 1e-10))
        conservative_violations += int(diff > bc * (1 + 1e-10))
        rows.append((e, diff, b, bc, diff / b))

    data = np.array(rows)
    # Prediction 1: exact perturbation is bounded for every mismatch size.
    # Prediction 2: theorem bound has O(e) scaling at small e.
    small_slope = np.median(data[:5, 1] / data[:5, 0])
    # Prediction 3: the quadratic term becomes non-negligible at large e;
    # report observed quadratic/theorem-bound contribution fraction.
    large_fraction = np.median((m * data[-4:, 0] ** 2) / data[-4:, 2])

    # Certificate threshold with one fixed perturbation direction.  This
    # isolates the predicted monotone loss of the conservative certificate.
    gamma = 0.18
    directions = []
    for w in true:
        d = rng.normal(size=w.shape)
        directions.append(d / opnorm(d))
    threshold_rows = []
    for e in np.linspace(0, 0.55, 56):
        estimated = [w + e*d for w, d in zip(true, directions)]
        gh, gt = gram(estimated), gram(true)
        rho = theorem_bound(true, [e] * m)
        q = np.linalg.eigvalsh(gh)[0] - rho
        actual = np.linalg.eigvalsh(gt)[0]
        threshold_rows.append((e, q, actual))
    cert = np.array(threshold_rows)
    positive = cert[cert[:, 1] > gamma]
    observed_e = positive[-1, 0] if len(positive) else np.nan
    pred_e = np.nan
    for a, b in zip(cert[:-1], cert[1:]):
        if (a[1] - gamma) * (b[1] - gamma) <= 0:
            pred_e = a[0] + (gamma-a[1]) * (b[0]-a[0]) / (b[1]-a[1])
            break
    # Directly verify Weyl at every sweep point, not only q-positive points.
    cert_margin = float(np.min(cert[:, 2] - cert[:, 1]))
    return {
        'bound_violations': int(violations),
        'computable_bound_violations': int(conservative_violations),
        'bound_max_ratio': float(max(r[4] for r in rows)),
        'small_epsilon_slope_median': float(small_slope),
        'large_epsilon_quadratic_fraction': float(large_fraction),
        'gamma': gamma,
        'certificate_crossing_epsilon_observed': float(observed_e),
        'certificate_crossing_epsilon_interpolated': float(pred_e),
        'minimum_actual_minus_q': cert_margin,
        'bound_table': [[float(x) for x in r] for r in rows],
    }


def identification(rng):
    p, n_candidates = 2, 45
    theta = np.array([1.25, -0.8])
    candidates = []
    for i in range(n_candidates):
        angle = 0.12 * i
        base = np.array([np.cos(angle), np.sin(angle)])
        w = np.tile(base, (4, 1)) + 0.08 * rng.normal(size=(4, p))
        candidates.append(w / opnorm(w))
    noise = 0.035
    y = [w @ theta + noise * rng.normal(size=4) for w in candidates]
    greedy, idx = greedy_stack(candidates, max_n=8)
    g = gram(greedy)

    def estimate(ids):
        H = gram([candidates[i] for i in ids], p)
        b = sum((candidates[i].T @ y[i] for i in ids), np.zeros(p))
        return np.linalg.solve(H + 1e-5*np.eye(p), b)

    errors_gated, errors_all, activated = [], [], False
    for t in range(1, n_candidates + 1):
        errors_all.append(float(np.linalg.norm(estimate(range(t)) - theta)))
        gs = [i for i in idx if i < t]
        if len(gs) >= p and np.linalg.eigvalsh(gram([candidates[i] for i in gs], p))[0] > 0.05:
            activated = True
            est = estimate(gs)
        else:
            est = estimate(gs) if activated else np.zeros(p)
        errors_gated.append(float(np.linalg.norm(est-theta)))
    return {
        'greedy_stack_size': len(greedy),
        'greedy_lambda_min': float(np.linalg.eigvalsh(g)[0]),
        'baseline_final_parameter_error': errors_all[-1],
        'idea_final_parameter_error': errors_gated[-1],
        'baseline_best_error': min(errors_all),
        'idea_best_error': min(errors_gated),
        'activation_step': next((t for t in range(1, n_candidates + 1) if len([i for i in idx if i < t]) >= p and np.linalg.eigvalsh(gram([candidates[i] for i in idx if i < t], p))[0] > 0.05), None),
        'errors_baseline': errors_all,
        'errors_idea': errors_gated,
    }


def main():
    rng = np.random.default_rng(2914)
    result = {'seed': 2914, 'mechanism': mechanism_checks(rng), 'identification': identification(rng)}
    with open('results.json', 'w') as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()