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

SEED = 1217
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

# One-step stochastic point-mass abstraction. z is a severity/initial-state feature.
# The policy is u=theta*z; task optimum theta=1, while safety requires theta*z <= c.
# Thus g(theta,z)=theta*z-c and v=[g]_+ (scale s=1).
C = 0.62
B = 32
ROUNDS = 18
UPDATES = 12
BETA = 10.0
LR = 0.045
T = 0.035

def soft_penalty(g):
    return torch.nn.functional.softplus(g / T).pow(2) * T*T

def update(theta, z, picked=False):
    # z is a tensor of trajectories in the current uniform batch or constraint buffer.
    theta = theta.detach().clone().requires_grad_(True)
    g = theta * z - C
    task = (theta - 1.0).pow(2)
    loss = task + BETA * soft_penalty(g).mean()
    loss.backward()
    with torch.no_grad():
        theta -= LR * theta.grad
    return theta.detach(), float(loss.detach())

def run(method, seed=SEED, record=False):
    rng = np.random.default_rng(seed)
    theta = torch.tensor(1.0)
    buffer = []
    history = []
    probe = torch.linspace(0, 1, 4001)
    for m in range(ROUNDS):
        z_np = rng.random(B)
        z = torch.tensor(z_np, dtype=torch.float32)
        v_np = np.maximum(theta.item() * z_np - C, 0)
        if method == 'pick':
            # The exact argmax trajectory is retained, as prescribed by C_{m+1}.
            j = int(np.argmax(v_np))
            buffer.append(float(z_np[j]))
            train_z = torch.tensor(buffer, dtype=torch.float32)
        else:
            train_z = z
        for _ in range(UPDATES):
            theta, _ = update(theta, train_z)
        probe_v = torch.relu(theta * probe - C)
        history.append({'round': m+1, 'theta': float(theta),
                        'batch_max_v': float(v_np.max()),
                        'probe_max_v': float(probe_v.max())})
    return theta.item(), history

def certify(theta, seed=991):
    rng = np.random.default_rng(seed)
    out = []
    delta = .05
    for n in [50, 100, 200, 400, 600, 1200, 2400]:
        z = rng.random(n)
        v = np.maximum(theta*z-C, 0)
        phat = float(np.mean(v > 0))
        eps = math.sqrt(math.log(1/delta)/(2*n))
        out.append({'N': n, 'phat': phat, 'upper': min(1., phat+eps), 'epsilon': eps})
    # The correction itself is the claimed O(N^-1/2) object.
    x = np.log([r['N'] for r in out]); y = np.log([r['epsilon'] for r in out])
    slope = float(np.polyfit(x, y, 1)[0])
    return out, slope

def main():
    # Prediction 1: for z~U[0,1], E[max z]=B/(B+1), hence selected max
    # should approach 1 and exceed the random-batch mean ~1/2.
    order_rows = []
    for b in [4, 8, 16, 32, 64, 128]:
        rng = np.random.default_rng(700+b)
        reps = 20000
        maxima = rng.random((reps,b)).max(axis=1)
        means = rng.random((reps,b)).mean(axis=1)
        order_rows.append({'B': b, 'observed_max': float(maxima.mean()),
                           'predicted_max': b/(b+1),
                           'observed_batch_mean': float(means.mean()),
                           'predicted_batch_mean': .5,
                           'max_minus_mean': float((maxima-means).mean())})
    # Prediction 2: the adaptive penalty should lower the worst-case probe margin
    # more than uniform empirical-risk updates at equal rollout budget.
    pick_theta, pick_hist = run('pick')
    uni_theta, uni_hist = run('uniform')
    ztest = np.linspace(0,1,200001)
    def eval_theta(th):
        v = np.maximum(th*ztest-C,0)
        return {'theta': th, 'violation_rate': float(np.mean(v>0)),
                'max_violation': float(v.max()), 'mean_violation': float(v.mean())}
    # Prediction 3: fixed-policy certificate correction has log-log slope -1/2.
    cert, cert_slope = certify(pick_theta)
    # Independent held-out certification requested by the idea.
    rng = np.random.default_rng(4242)
    zh = rng.random(600); vh = np.maximum(pick_theta*zh-C,0)
    heldout = {'N': 600, 'violations': int(np.sum(vh>0)),
               'phat': float(np.mean(vh>0)),
               'upper_95': float(min(1., np.mean(vh>0)+math.sqrt(math.log(20)/(1200))))}
    result = {
      'config': {'C': C, 'batch': B, 'rounds': ROUNDS, 'updates': UPDATES,
                 'beta': BETA, 'lr': LR, 'temperature': T, 'seed': SEED},
      'prediction_1_order_statistics': order_rows,
      'prediction_2_equal_budget': {'pick': eval_theta(pick_theta),
                                     'uniform': eval_theta(uni_theta),
                                     'pick_history': pick_hist,
                                     'uniform_history': uni_hist,
                                     'rollouts_each': B*ROUNDS},
      'prediction_3_certificate': {'rows': cert, 'observed_loglog_slope': cert_slope,
                                   'predicted_slope': -0.5, 'heldout_600': heldout},
      'notes': 'Violation is g=[theta*z-C], with U[0,1] severity and fixed policy during certification.'
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()