import json
import math
import random
from pathlib import Path
import numpy as np
from scipy.optimize import brentq, minimize
from scipy.special import logsumexp

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


def logcomb(n, k):
    return math.lgamma(n + 1) - math.lgamma(k + 1) - math.lgamma(n - k + 1)


def epsilon_exact(N, k, beta):
    """Theorem 1 epsilon: solve beta/N * sum(...) = 1."""
    if k >= N:
        return 1.0
    logs = np.array([logcomb(m, k) - logcomb(N, k) for m in range(k, N)])
    powers = np.arange(1, N - k + 1)  # N-m
    def f(e):
        z = logs - powers * math.log1p(-e)
        log_value = math.log(beta / N) + logsumexp(z)
        # Monotone residual in log-space; root is log_value == 0.
        return log_value
    lo = k / N
    # For sufficiently large k, the monotone equation reaches one at the
    # admissible endpoint; the interval bound is then the active certificate.
    a = lo + 1e-12
    if f(a) >= 0:
        return float(lo)
    return float(brentq(f, a, 1 - 1e-12, xtol=1e-12))


def calibrate(A, y, ids):
    """Minimum L2-norm theta satisfying selected linear constraints A theta >= y."""
    if len(ids) == 0:
        return np.zeros(A.shape[1])
    As, ys = A[ids], y[ids]
    fun = lambda t: 0.5 * float(t @ t)
    jac = lambda t: t
    cons = {'type': 'ineq', 'fun': lambda t: As @ t - ys,
            'jac': lambda t: As}
    out = minimize(fun, np.zeros(A.shape[1]), jac=jac, constraints=cons,
                   method='SLSQP', options={'maxiter': 300, 'ftol': 1e-11})
    if not out.success:
        raise RuntimeError(out.message)
    return out.x


def compress(A, y, k, mode, seed):
    local = np.random.default_rng(seed)
    chosen = []
    for _ in range(k):
        t = calibrate(A, y, chosen)
        violation = y - A @ t
        candidates = np.array([i for i in range(len(y)) if i not in chosen])
        if mode == 'greedy':
            i = int(candidates[np.argmax(violation[candidates])])
        else:
            i = int(local.choice(candidates))
        chosen.append(i)
    return np.array(chosen), calibrate(A, y, chosen)


def run_toy():
    # Scenario uncertainty is iid: each constraint is a random affine requirement.
    n, nt, k = 400, 50000, 2
    train = rng.uniform(-1, 1, size=n)
    test = rng.uniform(-1, 1, size=nt)
    # Convex envelope-like constraints; two extreme slopes are informative.
    A = np.column_stack([train, np.ones(n)])
    y = 0.34 + 0.42 * np.abs(train) + rng.normal(0, 0.018, n)
    At = np.column_stack([test, np.ones(nt)])
    yt = 0.34 + 0.42 * np.abs(test) + rng.normal(0, 0.018, nt)
    rows = []
    for kk in [2, 4, 8]:
        for mode in ['greedy', 'random']:
            chosen, theta = compress(A, y, kk, mode, SEED + kk)
            train_rate = float(np.mean(y - A @ theta > 1e-8))
            fresh_rate = float(np.mean(yt - At @ theta > 0))
            rows.append({'k': kk, 'method': mode, 'train_violation': train_rate,
                         'fresh_violation': fresh_rate, 'theta': theta.tolist(),
                         'selected': chosen.tolist()})
    # Non-iid stress test: shifted fresh distribution should not be called certified.
    shifted = rng.uniform(0.55, 1.0, size=nt)
    Ash = np.column_stack([shifted, np.ones(nt)])
    ysh = 0.34 + 0.42 * np.abs(shifted) + rng.normal(0, 0.018, nt)
    chosen, theta = compress(A, y, 2, 'greedy', SEED)
    shifted_rate = float(np.mean(ysh - Ash @ theta > 0))
    return rows, shifted_rate


def main():
    # Cheap math verification and explicit quantitative predictions.
    checks = []
    for N in [100, 200, 400, 800]:
        checks.append({'N': N, 'k': 2, 'beta': 1e-5, 'epsilon': epsilon_exact(N, 2, 1e-5)})
    beta_rows = [{'beta': b, 'epsilon': epsilon_exact(400, 2, b)}
                 for b in [1e-3, 1e-5, 1e-7]]
    k_rows = [{'k': k, 'epsilon': epsilon_exact(400, k, 1e-5)}
              for k in [0, 1, 2, 4, 8, 16]]
    toy, shifted = run_toy()
    result = {'seed': SEED, 'theorem_checks': checks, 'beta_sweep': beta_rows,
              'k_sweep': k_rows, 'toy': toy, 'shifted_fresh_violation_k2_greedy': shifted}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()