import os, sys, json, time
import numpy as np
import torch
from torch import nn

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

SEEDS = tuple(range(8))
# The sequence track is structurally matched: it is a temporal forecasting task.
# Quiet-clock threshold is fixed before running; no tuning on test data.
QUIET_THRESHOLD = 0.035
HIDDEN = 64


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


class ResidualState(nn.Module):
    """Shared neural vector field, with explicit or Stieltjes implicit updates."""
    def __init__(self, implicit, clock_scale=1.0, hidden=HIDDEN):
        super().__init__()
        self.implicit = bool(implicit)
        self.clock_scale = float(clock_scale)
        self.inp = nn.Linear(1, hidden)
        self.rec = nn.Linear(hidden, hidden, bias=False)
        self.bias = nn.Parameter(torch.zeros(hidden))
        self.head = nn.Linear(hidden, 1)

    def vector_field(self, z, x):
        return torch.tanh(self.inp(x.unsqueeze(-1)) + torch.matmul(z, self.rec.weight.t()) + self.bias) - z

    def clock(self, x):
        # Effective clock pauses on small chronological changes and jumps at events.
        change = torch.zeros_like(x)
        change[:, 1:] = (x[:, 1:] - x[:, :-1]).abs()
        return self.clock_scale * (change > QUIET_THRESHOLD).to(x.dtype)

    def forward(self, x):
        b, t = x.shape
        z = torch.zeros(b, self.inp.out_features, device=x.device, dtype=x.dtype)
        ds = self.clock(x) if self.implicit else torch.ones_like(x)
        for k in range(t):
            d = ds[:, k:k+1]
            if self.implicit:
                # Damped fixed-point solve of z_next=z+d f(z_next,x).
                zn = z
                for _ in range(10):
                    target = z + d * self.vector_field(zn, x[:, k])
                    zn = 0.5 * zn + 0.5 * target
                z = torch.where(d > 0, zn, z)
            else:
                z = z + d * self.vector_field(z, x[:, k])
        return self.head(z)


def train_one(seed, implicit, lr, clock_scale=1.0, n_train=400, n_test=200, epochs=15):
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=n_train, n_test=n_test)
    model = ResidualState(implicit=implicit, clock_scale=clock_scale)
    net, metric, history = train_model(model, ds, epochs=epochs, lr=lr, batch=128, log=lambda _: None)
    return float(metric) if metric is not None else float('nan')


def baseline_factory(cfg):
    return lambda seed: train_one(seed, False, cfg['lr'], 1.0)


def idea_factory(cfg):
    return lambda seed: train_one(seed, True, cfg['lr'], cfg['clock_scale'])


def mechanism_signature(seed=0):
    """Measure behavior of trained systems, rather than asserting an identity."""
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=400, n_test=200)
    base = ResidualState(False); idea = ResidualState(True, 1.0)
    # Train both systems independently on the same data.
    base, _, _ = train_model(base, ds, epochs=15, lr=0.003, batch=128, log=lambda _: None)
    idea, _, _ = train_model(idea, ds, epochs=15, lr=0.003, batch=128, log=lambda _: None)
    x = ds['xte']
    with torch.no_grad():
        def states(m, implicit):
            dev = next(m.parameters()).device
            xd = x.to(dev)
            z = torch.zeros(x.shape[0], HIDDEN, device=dev)
            clock = m.clock(xd)
            ds0 = clock if implicit else torch.ones_like(clock)
            changes=[]; active=[]
            for k in range(x.shape[1]):
                old=z.clone(); d=ds0[:, k:k+1]
                if implicit:
                    zn=z
                    for _ in range(10):
                        zn=0.5*zn+0.5*(z+d*m.vector_field(zn,xd[:,k]))
                    z=torch.where(d>0,zn,z)
                else: z=z+d*m.vector_field(z,xd[:,k])
                changes.append((z-old).norm(dim=1)); active.append((d[:,0]>0))
            c=torch.stack(changes); a=torch.stack(active)
            return float(c[~a].mean()) if (~a).any() else 0., float(c[a].mean()) if a.any() else 0., float(a.float().mean())
        bi, be, ba=states(base,False); ii, ie, ia=states(idea,True)
    # Predicted inactive change is zero; prediction is confirmed only with a measured small value.
    return {'predicted_inactive_state_change': 0.0,
            'observed_baseline_inactive_change': bi,
            'observed_idea_inactive_change': ii,
            'observed_baseline_event_change': be,
            'observed_idea_event_change': ie,
            'idea_active_fraction': ia,
            'predicted_active_fraction_is_data_dependent': True,
            'confirmed': bool(ii < 1e-7 and ia > 0.0)}


def main():
    # Search-space parity: every lr used by idea is also evaluated by baseline.
    lrs = [0.001, 0.003, 0.01]
    baseline_grid = [{'lr': lr, 'clock_scale': 1.0} for lr in lrs]
    idea_grid = [{'lr': lr, 'clock_scale': 1.0} for lr in lrs]
    base = sweep_baseline(baseline_factory, baseline_grid, seeds=(0,1,2,3))
    # Evaluate each idea setting on all paired seeds; select by the same sweep seeds.
    idea_sweep=[]
    for cfg in idea_grid:
        vals=[idea_factory(cfg)(s) for s in (0,1,2,3)]
        idea_sweep.append({'cfg':cfg, 'mean':float(np.mean(vals))})
    best_cfg=min(idea_sweep, key=lambda q:q['mean'])['cfg']
    idea_vals=[idea_factory(best_cfg)(s) for s in SEEDS]
    idea_res={'mean':float(np.mean(idea_vals)), 'std':float(np.std(idea_vals)), 'per_seed':idea_vals, 'n':len(idea_vals), 'best_cfg':best_cfg, 'sweep':idea_sweep}
    sig=mechanism_signature(0)
    report=make_report('sequence','custom_residual_state',base,idea_res,sig)
    report['protocol_notes']={'structural_match':'sequence-level temporal forecasting; recurrent state transition replacement', 'baseline_method':'explicit residual Euler with one chronological transition per token', 'idea_method':'implicit residual solve with event-clock deltas', 'paired_seeds':list(SEEDS), 'epochs':15, 'n_train':400, 'n_test':200}
    os.makedirs('artifacts',exist_ok=True)
    with open('artifacts/bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()