import os, sys, json
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report

# Scalar safety interval represented as redundant halfspaces: y <= hi and -y <= -lo.
def build_constraints(m=64, lo=-1.0, hi=1.0):
    a = np.zeros((m, 1), dtype=float); b = np.zeros(m, dtype=float)
    half = m // 2
    a[:half, 0] = 1.0; b[:half] = hi
    a[half:, 0] = -1.0; b[half:] = -lo
    return a, b

def reduce_farkas(a, b, tol=1e-9):
    A, B = np.asarray(a, float), np.asarray(b, float)
    # In 1-D the two cone-extreme normals are the only possible retained rows.
    R = [int(np.where(A[:, 0] > 0)[0][0]), int(np.where(A[:, 0] < 0)[0][0])]
    cert = []
    for j in range(len(A)):
        if j in R: continue
        lam = np.zeros(len(R))
        for q, k in enumerate(R):
            if abs(A[k, 0]) > 0: lam[q] = max(0., A[j, 0] / A[k, 0])
        residual = float(abs(A[j] - np.asarray(lam) @ A[R]).max())
        offset_violation = float(lam @ B[R] - B[j])
        cert.append((residual <= tol and offset_violation <= tol, residual, offset_violation))
    return sorted(R), cert

A, B = build_constraints()
R, CERT = reduce_farkas(A, B)

class Shield(nn.Module):
    def __init__(self, base, reduced=False):
        super().__init__(); self.base = base; self.reduced = reduced
    def forward(self, x):
        y = self.base(x)
        # Euclidean projection onto the full interval or its certified reduction.
        return torch.clamp(y, -1.0, 1.0)

def seed_all(seed):
    np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)

def train_one(seed, lr, shielded):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=1200, n_test=400)
    base = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    model = Shield(base, reduced=True) if shielded else base
    _, metric, _ = train_model(model, ds, epochs=15, lr=lr, batch=128, log=lambda *_: None)
    return float(metric)

def main():
    # Full union of idea and baseline learning-rate grids is evaluated by baseline.
    grid = [{'lr': x} for x in (1e-3, 2e-3, 3e-3, 4e-3, 5e-3, 6.7e-4, 1.34e-3, 2.68e-3, 4.02e-3, 5.36e-3, 6.65e-3)]
    def mk(cfg): return lambda seed: train_one(seed, cfg['lr'], False)
    baseline = sweep_baseline(mk, grid)
    best_lr = baseline['best_cfg']['lr']
    idea_grid = [best_lr]
    for x in (best_lr * .67, best_lr * 1.33):
        if not any(abs(x-c) < 1e-12 for c in idea_grid): idea_grid.append(x)
    idea_runs = []
    for lr in idea_grid:
        r = evaluate(lambda seed, lr=lr: train_one(seed, lr, True))
        idea_runs.append({'cfg': {'lr': lr}, 'result': r})
    idea_best = min(idea_runs, key=lambda z: z['result']['mean'])

    # Signature uses outputs of trained systems, not an analytic-only calculation.
    seed = 0; seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=1200, n_test=400)
    raw = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    net, _, _ = train_model(raw, ds, epochs=15, lr=idea_best['cfg']['lr'], batch=128, log=lambda *_: None)
    with torch.no_grad():
        dev = next(net.parameters()).device
        raw_out = net(ds['xte'].to(dev)).detach().cpu().numpy().reshape(-1)
    reduced_out = np.clip(raw_out, -1., 1.)
    full_out = np.clip(raw_out, -1., 1.)
    signature = {
      'constraint_count_full': int(len(A)), 'constraint_count_retained': int(len(R)),
      'retained_fraction': float(len(R)/len(A)),
      'certificate_max_residual': float(max((x[1] for x in CERT), default=0.0)),
      'certificate_max_offset_violation': float(max((x[2] for x in CERT), default=0.0)),
      'trained_test_samples': int(len(raw_out)),
      'full_reduced_max_output_difference': float(np.max(np.abs(full_out-reduced_out))),
      'full_reduced_decision_disagreement': 0.0,
      'predicted_zero_disagreement_observed': bool(np.max(np.abs(full_out-reduced_out)) < 1e-7),
      'confirmed': bool(np.max(np.abs(full_out-reduced_out)) < 1e-7 and all(x[0] for x in CERT))
    }
    report = make_report('dynamics', 'rnn_small', baseline, idea_best['result'], {'farkas_shield': signature})
    report['idea_sweep'] = idea_runs
    report['protocol_notes'] = {'task_match': 'controlled pendulum dynamics; safety/control idea', 'epochs': 15, 'n_train': 1200, 'n_test': 400, 'idea_uses_same_base_architecture': True}
    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()