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

# Local custom track: observations are Fourier measurements of two separated sources.
# Target is source coordinates (fixed two-source instance per sample); this is the
# closest valid supervised neural benchmark for the continuous localization mechanism.
META={'name':'fourier_source_localization','domain':'continuous_localization','description':'Noisy Fourier measurements of two separated planar sources; regress their coordinates.'}
def get_dataset(seed,n_train,n_test):
    def make(n, ss):
        rng=np.random.default_rng(ss); k=np.array([2,3,5,7,11,13,17,19],float)
        def one():
            x=np.sort(rng.uniform(.12,.88,2)); y=rng.uniform(.12,.88,2)
            # fixed Fourier feature bank, complex coefficients encoded real/imag
            ph=np.outer(k,x)+np.outer(k*0.73,y)
            a=(np.exp(1j*ph[:,0])+np.exp(1j*ph[:,1]))/np.sqrt(2)
            z=np.concatenate([a.real,a.imag])+rng.normal(0,.035,16)
            return z, np.array([x[0],y[0],x[1],y[1]])
        q=[one() for _ in range(n)]; return np.float32([a for a,b in q]),np.float32([b for a,b in q])
    xtr,ytr=make(n_train,seed); xte,yte=make(n_test,seed+5000)
    return {'xtr':torch.from_numpy(xtr),'ytr':torch.from_numpy(ytr),'xte':torch.from_numpy(xte),'yte':torch.from_numpy(yte),'task':'regression','metric':'mse','out_dim':4,'input_shape':(xtr.shape[1],)}

def train_eval(seed, lr, idea, epochs=12):
    d=get_dataset(seed,400,200); torch.manual_seed(seed+991)
    net=make_model('mlp_tiny',d['input_shape'],4)
    # Same architecture and training path. Idea is a differentiable soft coarse-grid
    # proposal/refinement layer applied to the model's coordinate hypotheses.
    if idea:
        class Refined(torch.nn.Module):
            def __init__(self, base): super().__init__(); self.base=base
            def forward(self,x):
                raw=self.base(x).view(-1,2,2).sigmoid()
                # calibrated fixed-step contraction toward the nearest coarse proposal;
                # this is the coarse-to-fine inductive bias, not an oracle.
                grid=torch.tensor([.25,.5,.75],dtype=x.dtype,device=x.device)
                g=torch.stack(torch.meshgrid(grid,grid,indexing='ij'),-1).reshape(-1,2)
                dist=((raw[:,:,None,:]-g[None,None,:,:])**2).sum(-1)
                w=torch.softmax(-30*dist,-1)
                prop=(w[:,:,:,None]*g[None,None,:,:]).sum(2)
                return (raw + 0.35*(prop-raw)).reshape(-1,4)
        net=Refined(net)
    net, metric, hist=train_model(net,d,epochs=epochs,lr=lr)
    dev=next(net.parameters()).device
    with torch.no_grad():
        pred=net(d['xte'].to(dev)).cpu().numpy(); obs=d['yte'].numpy()
    mse=float(np.mean((pred-obs)**2))
    # trained-model mechanism signature: measured proposal displacement and residual.
    raw=net.base(d['xte'].to(dev)).view(-1,2,2).sigmoid().detach().cpu().numpy() if idea else pred.reshape(-1,2,2)
    final=pred.reshape(-1,2,2); displacement=float(np.mean(np.linalg.norm(final-raw,axis=-1)))
    return mse, {'mse':mse,'proposal_displacement':displacement,'n':len(obs)}

def main():
    # shared union of baseline/idea learning rates; baseline sweep includes all.
    grid=[{'lr':1e-3,'epochs':12},{'lr':3e-3,'epochs':12},{'lr':1e-2,'epochs':12}]
    seeds=list(range(8)); base_sweep=[]
    for cfg in grid:
        vals=[train_eval(s,cfg['lr'],False,cfg['epochs'])[0] for s in seeds[:4]]
        base_sweep.append({'config':cfg,'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':4})
    best=min(base_sweep,key=lambda z:z['mean'])['config']
    base=[]; idea=[]; sig=[]
    for s in seeds:
        b,_=train_eval(s,best['lr'],False,best['epochs']); i,sg=train_eval(s,best['lr'],True,best['epochs'])
        base.append(b); idea.append(i); sig.append(sg)
    # nearby idea settings are run, but only shared/best setting is reported to keep
    # the paired comparison transparent and budget equal.
    nearby={str(lr):[train_eval(s,lr,True,12)[0] for s in seeds[:4]] for lr in [1e-3,1e-2]}
    delta=np.asarray(idea)-np.asarray(base)
    rng=np.random.default_rng(0); signs=rng.choice([-1,1],(20000,8)); p=float(np.mean(np.abs((signs*delta).mean(1))>=abs(delta.mean())))
    report={'track':'custom_fourier_source_localization','model':'mlp_tiny','baseline':{'best_cfg':best,'sweep':base_sweep,'full':{'mean':float(np.mean(base)),'std':float(np.std(base)),'per_seed':base,'n':8}},'idea':{'mean':float(np.mean(idea)),'std':float(np.std(idea)),'per_seed':idea,'n':8,'nearby':nearby},'paired_delta_mean':float(delta.mean()),'permutation_pvalue':p,'mechanism_signature':{'predicted':'coarse proposal contracts raw coordinates with fixed bandwidth-independent shrinkage','observed_mean_proposal_displacement':float(np.mean([x['proposal_displacement'] for x in sig])),'observed_mse':float(np.mean(idea)),'confirmed':False},'custom_track':{'name':'fourier_source_localization','file':'bench_experiment.py','domain':'continuous_localization'}}
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()