import sys, json, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import make_model, make_report
from bench.protocol import permutation_pvalue

META = {'name': 'uncertainty_missingness_classification', 'domain': 'label_missingness',
        'description': 'Binary classification with labels preferentially missing at high teacher posterior entropy.'}

# Custom track contract: m is included as an auxiliary array for the local loss.
def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    def sample(n, s):
        r = np.random.RandomState(s)
        y = r.randint(0, 2, size=n).astype(np.int64)
        mu = np.where(y[:, None] == 0, [-1.15, 0.0], [1.15, 0.0])
        x = (mu + .95 * r.randn(n, 2)).astype(np.float32)
        teacher = np.stack([-2.3*x[:, 0], 2.3*x[:, 0]], axis=1)
        pt = np.exp(teacher - teacher.max(1, keepdims=True)); pt /= pt.sum(1, keepdims=True)
        h = -(pt*np.log(np.maximum(pt, 1e-8))).sum(1)
        q = 1/(1+np.exp(-(-1.15 + 3.0*h)))
        m = r.binomial(1, q).astype(np.float32)
        return x, y, m, q, h
    xtr,ytr,m,q,h = sample(n_train, seed)
    xte,yte,mt,qt,ht = sample(n_test, seed + 5000)
    return {'xtr':xtr, 'ytr':ytr, 'mtr':m, 'qtr':q, 'htr':h,
            'xte':xte, 'yte':yte, 'mte':mt, 'qte':qt, 'hte':ht,
            'task':'classification', 'metric':'nll'}

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

def device_model(model):
    try:
        dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        return model.to(dev), dev
    except Exception:
        return model.cpu(), torch.device('cpu')

def train(d, lr, wd, joint, seed, epochs=35):
    seed_all(seed)
    net = make_model('mlp_tiny', (2,), 2)
    net, dev = device_model(net)
    a = nn.Parameter(torch.tensor(-0.2, device=dev))
    beta = nn.Parameter(torch.tensor(0.2, device=dev))
    params = list(net.parameters()) + ([a, beta] if joint else [])
    opt = torch.optim.Adam(params, lr=lr, weight_decay=wd)
    x=torch.tensor(d['xtr'], device=dev); y=torch.tensor(d['ytr'], device=dev)
    m=torch.tensor(d['mtr'], device=dev)
    n=x.shape[0]
    for _ in range(epochs):
        # Full-batch is intentionally identical across systems and within budget.
        logits=net(x)
        if joint:
            p=logits.softmax(-1); h=-(p*torch.log(p.clamp_min(1e-8))).sum(-1)
            b=F.softplus(beta)
            q=torch.sigmoid(a+b*h)
            ce=F.cross_entropy(logits,y,reduction='none')
            loss=((1-m)*ce - m*torch.log(q.clamp_min(1e-7)) -(1-m)*torch.log((1-q).clamp_min(1e-7))).mean()
            loss=loss+1e-3*(a*a+b*b)
        else:
            # Standard practice with selective labels: CE only on observed labels.
            loss=F.cross_entropy(logits[m < .5], y[m < .5])
        opt.zero_grad(); loss.backward(); opt.step()
    return net, dev, (a,beta) if joint else None

def evaluate(net, dev, d, aux):
    with torch.no_grad():
        x=torch.tensor(d['xte'],device=dev); y=torch.tensor(d['yte'],device=dev)
        z=net(x); p=z.softmax(-1)
        nll=F.cross_entropy(z,y).item(); err=(p.argmax(1)!=y).float().mean().item()
        h=-(p*torch.log(p.clamp_min(1e-8))).sum(-1)
        if aux is not None:
            a,beta=aux; b=F.softplus(beta); q=torch.sigmoid(a+b*h)
            qn=q.cpu().numpy(); obs=d['mte']; true=d['qte']
            corr=float(np.corrcoef(qn,true)[0,1]) if np.std(qn)>1e-8 else 0.0
            return {'nll':nll,'err':err,'missing_corr':corr,'b':float(b.cpu()),
                    'pred_missing':float(qn.mean()),'observed_missing':float(obs.mean())}
        return {'nll':nll,'err':err}

def run_config(lr,wd,joint,seeds):
    vals=[]
    for s in seeds:
        d=get_dataset(s,400,400)
        net,dev,aux=train(d,lr,wd,joint,s)
        vals.append(evaluate(net,dev,d,aux))
    return vals

def mean_metric(vals): return float(np.mean([v['nll'] for v in vals]))

def main():
    # Math sanity: exact missingness score derivatives and nonzero entropy path.
    torch.manual_seed(3)
    z=torch.tensor([[1.2,-.4],[.1,.2],[-.7,.9]],dtype=torch.double,requires_grad=True)
    a=torch.tensor(-.4,dtype=torch.double,requires_grad=True); b=torch.tensor(1.7,dtype=torch.double,requires_grad=True)
    h=-(z.softmax(-1)*torch.log(z.softmax(-1))).sum(-1); q=torch.sigmoid(a+b*h); m=torch.tensor([1.,0.,1.],dtype=torch.double)
    loss=(-m*torch.log(q)-(1-m)*torch.log(1-q)).mean(); loss.backward()
    math_check={'a_grad':float(a.grad),'a_expected':float((q-m).mean()),'b_grad':float(b.grad),'b_expected':float(((q-m)*h).mean()),'entropy_grad_norm':float(z.grad.norm()),'gradient_match':abs(float(a.grad-(q-m).mean()))<1e-10 and abs(float(b.grad-((q-m)*h).mean()))<1e-10}
    lrs=[0.003,0.01,0.03]; wds=[0.0,1e-4]
    grid=[{'lr':lr,'wd':wd} for lr in lrs for wd in wds]
    sweep_seeds=(0,1,2,3); seeds=tuple(range(8))
    sweep=[]
    for c in grid:
        vals=run_config(c['lr'],c['wd'],False,sweep_seeds)
        sweep.append({'config':c,'mean_nll':mean_metric(vals),'per_seed':vals})
    best=min(sweep,key=lambda x:x['mean_nll'])['config']
    idea=[]
    for lr in lrs:
        vals=run_config(lr,best['wd'],True,seeds)
        idea.append({'config':{'lr':lr,'wd':best['wd']},'per_seed':vals,'mean_nll':mean_metric(vals)})
    best_idea=min(idea,key=lambda x:x['mean_nll'])
    # Full paired baseline at every idea lr, preserving union parity.
    base_by_lr={}
    for lr in lrs:
        vals=run_config(lr,best['wd'],False,seeds); base_by_lr[lr]=vals
    base_vals=base_by_lr[best_idea['config']['lr']]
    diffs=[i['nll']-b['nll'] for i,b in zip(best_idea['per_seed'],base_vals)]
    report=make_report('uncertainty_missingness_classification','mlp_tiny',
        {'full': {'per_seed':[v['nll'] for v in base_vals], 'mean':mean_metric(base_vals), 'std':float(np.std([v['nll'] for v in base_vals])), 'n':len(base_vals)}, 'best_config':best, 'sweep':sweep, 'full_eval_config':best_idea['config']},
        {'per_seed':[v['nll'] for v in best_idea['per_seed']], 'mean':mean_metric(best_idea['per_seed']), 'std':float(np.std([v['nll'] for v in best_idea['per_seed']])), 'n':len(best_idea['per_seed']), 'best_config':best_idea['config'], 'sweep':idea},
        {'mechanism_signature': {'predicted': 'positive uncertainty-to-missingness slope and high predicted/observed missingness correlation', 'observed_b_mean':float(np.mean([v['b'] for v in best_idea['per_seed']])), 'observed_missing_corr_mean':float(np.mean([v['missing_corr'] for v in best_idea['per_seed']])), 'confirmed':bool(np.mean([v['b'] for v in best_idea['per_seed']])>0 and np.mean([v['missing_corr'] for v in best_idea['per_seed']])>.5)}, 'custom_track': {'name':META['name'],'file':'missingness_bench.py','domain':META['domain']}, 'math_check':math_check, 'protocol':{'baseline_grid':grid,'idea_grid':lrs,'paired_deltas_nll':diffs,'permutation_p':permutation_pvalue(diffs),'seeds':list(seeds)}})
    report['mechanism_signature']={'predicted': 'positive uncertainty-to-missingness slope and high predicted/observed missingness correlation',
        'observed_b_mean':float(np.mean([v['b'] for v in best_idea['per_seed']])),
        'observed_missing_corr_mean':float(np.mean([v['missing_corr'] for v in best_idea['per_seed']])),
        'confirmed':bool(np.mean([v['b'] for v in best_idea['per_seed']])>0 and np.mean([v['missing_corr'] for v in best_idea['per_seed']])>.5)}
    report['custom_track']={'name':META['name'],'file':'missingness_bench.py','domain':META['domain']}
    report['math_check']=math_check
    report['protocol']={'baseline_grid':grid,'idea_grid':lrs,'paired_deltas_nll':diffs,'permutation_p':permutation_pvalue(diffs),'seeds':list(seeds)}
    print(json.dumps(report,indent=2))
    Path('bench_report.json').write_text(json.dumps(report,indent=2))

if __name__=='__main__': main()