import json
import sys
import numpy as np
import torch

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

TRACK = 'two_view_gauge_localization_v2'
MODEL = 'mlp_tiny'
EPOCHS = 12
NTRAIN, NTEST = 400, 200

class CoarseToFine(torch.nn.Module):
    """Differentiable coarse proposal followed by one fixed refinement step.

    The scalar target is the localized x-coordinate in the registered localization
    track. The coarse grid is fixed in the known bounded target domain [-2, 2].
    """
    def __init__(self, base, temperature=18.0, refinement=0.35):
        super().__init__()
        self.base = base
        self.temperature = temperature
        self.refinement = refinement
        self.register_buffer('grid', torch.linspace(-2.0, 2.0,  nine if False else 9))

    def forward(self, x):
        raw = self.base(x).reshape(-1, 1)
        dist = (raw - self.grid.reshape(1, -1)) ** 2
        weights = torch.softmax(-self.temperature * dist, dim=1)
        proposal = (weights * self.grid.reshape(1, -1)).sum(dim=1, keepdim=True)
        return raw + self.refinement * (proposal - raw)


def run_one(seed, lr, idea):
    d = get_dataset(TRACK, int(seed), n_train=NTRAIN, n_test=NTEST)
    torch.manual_seed(10000 + int(seed))
    base = make_model(MODEL, d['input_shape'], d['out_dim'])
    net = CoarseToFine(base) if idea else base
    net, _, _ = train_model(net, d, epochs=EPOCHS, lr=float(lr), log=lambda *_: None)
    if net is None:
        return float('nan'), {'raw_mae': float('nan'), 'refined_mae': float('nan'), 'proposal_shift': float('nan')}
    device = next(net.parameters()).device
    xte, yte = d['xte'].to(device), d['yte'].to(device)
    with torch.no_grad():
        pred = net(xte)
        if idea:
            raw = net.base(xte)
            refined_mae = float(torch.mean(torch.abs(pred - yte)).cpu())
            raw_mae = float(torch.mean(torch.abs(raw - yte)).cpu())
            shift = float(torch.mean(torch.abs(pred - raw)).cpu())
        else:
            refined_mae = float(torch.mean(torch.abs(pred - yte)).cpu())
            raw_mae, shift = refined_mae, 0.0
        mse = float(torch.mean((pred - yte) ** 2).cpu())
    # The benchmark metric remains the task MSE; signature numbers are measured
    # from predictions of the trained systems, not from an analytic toy identity.
    return mse, {'raw_mae': raw_mae, 'refined_mae': refined_mae, 'proposal_shift': shift}


def main():
    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
    cache = {}
    def baseline_factory(cfg):
        return lambda seed: run_one(seed, cfg['lr'], False)[0]
    base_block = sweep_baseline(baseline_factory, grid)
    best_lr = float(base_block['best_cfg']['lr'])
    idea_res = evaluate(lambda seed: run_one(seed, best_lr, True)[0])
    nearby = {}
    for lr in (1e-3, 1e-2):
        r = evaluate(lambda seed, lr=lr: run_one(seed, lr, True)[0], seeds=(0,1,2,3))
        nearby[str(lr)] = r
    sig = evaluate(lambda seed: run_one(seed, best_lr, True)[1]['proposal_shift'])
    raw = evaluate(lambda seed: run_one(seed, best_lr, True)[1]['raw_mae'])
    refined = evaluate(lambda seed: run_one(seed, best_lr, True)[1]['refined_mae'])
    extra = {
        'mechanism_signature': {
            'prediction': 'coarse proposal plus fixed refinement should reduce raw localization error',
            'raw_mae_mean': raw['mean'],
            'refined_mae_mean': refined['mean'],
            'mean_proposal_shift': sig['mean'],
            'confirmed': bool(np.isfinite(raw['mean']) and refined['mean'] < raw['mean'])
        },
        'idea_nearby_sweep': nearby,
        'custom_track': {'name': TRACK, 'registered': True, 'domain': 'geometry/localization'}
    }
    report = make_report(TRACK, MODEL, base_block, idea_res, extra)
    with open('registered_bench_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()