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

import sys
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report

TRACK='dynamics'; MODEL='rnn_small'
# Equal-budget settings; the union is used on both sides.
LR_GRID=[1e-3, 3e-3, 6e-3]
EPOCHS=12
N_TRAIN=400
N_POOL=1600
SEEDS=tuple(range(8))


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


def state_view(x):
    # Last observed pendulum state in each 8-step input window.
    return np.asarray(x, dtype=np.float64).reshape(len(x), 8, 3)[:, -1, :]


def residual(x):
    # Conservative, task-independent safety margin: theta and angular velocity.
    z=state_view(x)
    return np.minimum(1.25-np.abs(z[:,0]), 2.0-0.35*np.abs(z[:,1]))


def cover_radius(query, samples):
    # Small pool: exact vectorized nearest-neighbor radius.
    q=np.asarray(query); s=np.asarray(samples)
    best=np.full(len(q), np.inf)
    for j in range(0,len(s),128):
        d=((q[:,None,:]-s[None,j:j+128,:])**2).sum(2)
        best=np.minimum(best, np.sqrt(d.min(1)))
    return float(best.max()), best


def adaptive_indices(pool_x, n, seed, gamma=0.35):
    z=state_view(pool_x)
    rng=np.random.default_rng(seed)
    # Four seed states, then q = normalized distance + gamma*low-margin.
    chosen=list(rng.choice(len(pool_x), size=8, replace=False))
    chosen_set=set(chosen)
    scale=np.std(z,axis=0)+1e-6
    zn=z/scale
    r=residual(pool_x)
    eps=0.15
    while len(chosen)<n:
        sel=np.asarray(chosen)
        best=np.full(len(pool_x), np.inf)
        for j in range(0,len(sel),128):
            d=((zn[:,None,:]-zn[sel[j:j+128]][None,:,:])**2).sum(2)
            best=np.minimum(best, np.sqrt(d.min(1)))
        q=best + gamma*np.maximum(0.,eps-r)
        q[list(chosen_set)]=-np.inf
        k=int(np.argmax(q)); chosen.append(k); chosen_set.add(k)
    return np.asarray(chosen,dtype=int)


def prepare(seed, mode):
    # Pool and test set are generated from the same benchmark dynamics task.
    base=get_dataset(TRACK, seed, n_train=N_TRAIN, n_test=1000)
    pool=get_dataset(TRACK, seed+10000, n_train=N_POOL, n_test=1000)
    rng=np.random.default_rng(seed+991)
    if mode=='baseline': idx=rng.choice(N_POOL,size=N_TRAIN,replace=False)
    else: idx=adaptive_indices(pool['xtr'].numpy(),N_TRAIN,seed)
    ds=dict(base)
    ds['xtr']=pool['xtr'][idx].clone(); ds['ytr']=pool['ytr'][idx].clone()
    return ds, pool, idx


def run_one(seed, mode, cfg, return_net=False):
    seed_all(seed + (0 if mode=='baseline' else 100000))
    ds,pool,idx=prepare(seed,mode)
    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
    net,metric,hist=train_model(net,ds,epochs=EPOCHS,lr=float(cfg['lr']),batch=128,log=lambda *_:None)
    if return_net:
        return float(metric),net,ds,pool,idx
    return float(metric)


def math_check():
    rng=np.random.default_rng(4); cloud=rng.uniform(-1,1,(500,3)); samples=cloud[rng.choice(500,40,False)]
    # r(x)=0.2+L||x-c|| is L-Lipschitz; verify the stated lower bound.
    L=0.7; c=np.array([0.1,-0.2,0.3]); r=lambda a:0.2+L*np.linalg.norm(a-c,axis=1)
    d=cover_radius(cloud,samples)[0]; assert r(cloud).min() >= r(samples).min()-L*d-1e-10
    return {'checked':True,'L':L,'delta':d,'sample_min':float(r(samples).min()),'cloud_min':float(r(cloud).min()),'bound':float(r(samples).min()-L*d)}


def main():
    check=math_check()
    grid=[{'lr':x} for x in LR_GRID]
    base=sweep_baseline(lambda cfg: lambda s: run_one(s,'baseline',cfg),grid,seeds=(0,1,2,3))
    # Explicitly evaluate idea at all three settings; best is selected on the same tuning seeds.
    idea_cfg_results=[]
    for cfg in grid:
        rr=evaluate(lambda s,cfg=cfg: run_one(s,'idea',cfg),seeds=(0,1,2,3))
        idea_cfg_results.append((rr['mean'],cfg))
    idea_cfg=min(idea_cfg_results,key=lambda x:x[0])[1]
    idea=evaluate(lambda s: run_one(s,'idea',idea_cfg),seeds=SEEDS)
    # Signature uses trained models, not a toy identity: compare pool coverage and
    # boundary/general-region prediction errors for one paired seed.
    bm,bnet,bds,bpool,bidx=run_one(0,'baseline',base['best_cfg'],True)
    im,inet,ids,ipool,iidx=run_one(0,'idea',idea_cfg,True)
    with torch.no_grad():
        bx=bpool['xtr']; by=bpool['ytr']; ix=ipool['xtr']; iy=ipool['ytr']
        bdev=next(bnet.parameters()).device; idev=next(inet.parameters()).device
        bp=bnet(bx.to(bdev)).detach().cpu().numpy().reshape(-1); ip=inet(ix.to(idev)).detach().cpu().numpy().reshape(-1)
    bz=state_view(bpool['xtr'].numpy()); iz=state_view(ipool['xtr'].numpy())
    bdelta=cover_radius(bz,state_view(bpool['xtr'][bidx].numpy()))[0]
    idelta=cover_radius(iz,state_view(ipool['xtr'][iidx].numpy()))[0]
    br=residual(bpool['xtr'].numpy()); ir=residual(ipool['xtr'].numpy())
    sig={'prediction':'adaptive reachable-state selection should reduce empirical coverage radius at equal simulator budget',
         'observed':{'baseline_delta':bdelta,'idea_delta':idelta,'baseline_pool_mse':float(np.mean((bp-by.numpy())**2)),'idea_pool_mse':float(np.mean((ip-iy.numpy())**2)), 'baseline_min_sample_residual':float(br[bidx].min()),'idea_min_sample_residual':float(ir[iidx].min())},
         'relative_delta_reduction_pct':float(100*(bdelta-idelta)/bdelta),
         'confirmed':bool(idelta < bdelta)}
    report=make_report(TRACK,MODEL,base,idea,extra={'math_check':check,'mechanism_signature':sig,'idea_grid_results':idea_cfg_results,'selected_idea_cfg':idea_cfg,'budget':{'epochs':EPOCHS,'train_samples':N_TRAIN,'pool_samples':N_POOL}})
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()