import os, sys, json, random, time
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, make_report, sweep_baseline, evaluate

TRACK='vision'; MODEL='cnn_small'; EPOCHS=3; NTR=400; NTE=200; BATCH=128
LR_GRID=[1e-3, 3e-3, 1e-2]
SEEDS=tuple(range(8))


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


def interval_conv(conv, lo, hi):
    w=conv.weight.detach(); b=conv.bias.detach() if conv.bias is not None else 0
    wp=torch.clamp(w,min=0); wn=torch.clamp(w,max=0)
    return (nn.functional.conv2d(lo,wp,bias=None,padding=conv.padding)+
            nn.functional.conv2d(hi,wn,bias=None,padding=conv.padding)+b.view(1,-1,1,1),
            nn.functional.conv2d(hi,wp,bias=None,padding=conv.padding)+
            nn.functional.conv2d(lo,wn,bias=None,padding=conv.padding)+b.view(1,-1,1,1))


def prune_signature(net, x, eps=.02, tau=0.0):
    """Sound interval sign classification for the two CNN ReLU blocks.
    It evaluates the trained model on x and an independent perturbation sample.
    """
    m=net.net; lo=torch.clamp(x-eps,0,1); hi=torch.clamp(x+eps,0,1)
    fixed=unstable=viol=total=0; max_err=0.0; zlo,zhi=lo,hi
    for layer in m:
        if isinstance(layer,nn.Conv2d):
            L,U=interval_conv(layer,zlo,zhi)
            # retain spatial neuron instances; exact fixed-sign substitution
            active=L>=0; inactive=U<=0; unst=~(active|inactive)
            fixed += int(active.sum()+inactive.sum()); unstable += int(unst.sum()); total += int(L.numel())
            # independent point check, using the same certified bounds
            q=lo + (hi-lo)*torch.rand_like(lo)
            a=layer(q if zlo is lo else zlo)
            exact=torch.relu(a)
            substituted=torch.where(active,a,torch.where(inactive,torch.zeros_like(a),torch.relu(a)))
            max_err=max(max_err,float((exact-substituted).abs().max()))
            # propagate interval through the ReLU and pooling
            zlo,zhi=torch.relu(L),torch.relu(U)
        elif isinstance(layer,nn.ReLU):
            pass
        elif isinstance(layer,nn.MaxPool2d):
            zlo=layer(zlo); zhi=layer(zhi)
        elif isinstance(layer,nn.Flatten):
            zlo=layer(zlo); zhi=layer(zhi)
        elif isinstance(layer,nn.Linear):
            break
    # Behavioural test: exact substitutions must match ordinary predictions
    with torch.no_grad():
        normal=net(x).argmax(1)
    return {'fixed_sign_fraction': fixed/max(total,1), 'fixed_sign':fixed,
            'unstable':unstable, 'total_relu_instances':total,
            'max_fixed_substitution_error':max_err,
            'prediction_count':int(len(normal)), 'tau':tau,
            'predicted_reduction_fraction':fixed/max(total,1)}


def train_one(cfg, seed, collect=False):
    seed_all(seed); d=get_dataset(TRACK,seed,n_train=NTR,n_test=NTE)
    net=make_model(MODEL,d['input_shape'],d['out_dim'])
    net, metric, hist=train_model(net,d,epochs=EPOCHS,lr=float(cfg['lr']),batch=BATCH,log=lambda *_:None)
    if net is None: return float('inf')
    if collect:
        net = net.cpu()
        sig=prune_signature(net,d['xte'][:min(32,len(d['xte']))])
        return metric,sig
    return metric


def main():
    # Baseline sweep over exactly the union of all idea learning rates.
    base=sweep_baseline(lambda cfg: (lambda seed: train_one(cfg,seed)),
                        [{'lr':lr} for lr in LR_GRID], seeds=(0,1,2,3))
    # Re-evaluate the selected baseline and all three idea settings on paired seeds.
    base_full=evaluate(lambda s: train_one(base['best_cfg'],s), seeds=SEEDS)
    base['full']=base_full
    idea_results=[]; signatures=[]
    for cfg in [{'lr':lr} for lr in LR_GRID]:
        r=evaluate(lambda s,cfg=cfg: train_one(cfg,s), seeds=SEEDS)
        idea_results.append({'cfg':cfg,'result':r})
    best=min(idea_results,key=lambda z:z['result']['mean'])
    for s in SEEDS:
        _,sig=train_one(best['cfg'],s,collect=True); signatures.append(sig)
    sig={k:float(np.mean([x[k] for x in signatures])) for k in ('fixed_sign_fraction','max_fixed_substitution_error','predicted_reduction_fraction')}
    sig.update({'fixed_sign_mean':float(np.mean([x['fixed_sign'] for x in signatures])),
                'unstable_mean':float(np.mean([x['unstable'] for x in signatures])),
                'total_relu_instances_mean':float(np.mean([x['total_relu_instances'] for x in signatures])),
                'confirmed': bool(sig['max_fixed_substitution_error'] <= 1e-7 and sig['predicted_reduction_fraction'] > 0)})
    rep=make_report(TRACK,MODEL,base,best['result'],extra={'mechanism_signature':sig,
        'idea_sweep':idea_results,
        'protocol_note':'Verifier pruning is inference-time only; training and standard test metric are unchanged systems.'})
    rep['experiment_settings']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'batch':BATCH,'lr_union':LR_GRID,'seeds':list(SEEDS)}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()