import json
import numpy as np
import torch
from torch import nn
from volatility_clock import simulate, params


def make_batch(kind, rng, batch=96, length=32):
    # A deliberately small synthetic diffusion-like task. x0 is an AR(1)
    # signal; the denoiser is blind to realized A and sees only xt.
    x0 = np.zeros((batch, length), dtype=np.float32)
    for j in range(1, length):
        x0[:, j] = .85*x0[:, j-1] + rng.standard_normal(batch).astype(np.float32)
    if kind == 'clock':
        _, eps = simulate(.65, .8, .7, batch, length, rng, burn=120)
    else:
        eps = rng.standard_normal((batch, length))
    ab = .55
    xt = (np.sqrt(ab)*x0 + np.sqrt(1-ab)*eps).astype(np.float32)
    return torch.tensor(xt[:, None, :], dtype=torch.float32), torch.tensor(eps[:, None, :], dtype=torch.float32)


class TinyDenoiser(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Conv1d(1, 16, 5, padding=2), nn.GELU(),
                                 nn.Conv1d(16, 16, 5, padding=2), nn.GELU(),
                                 nn.Conv1d(16, 1, 5, padding=2))
    def forward(self, x): return self.net(x)


def run_one(kind):
    torch.manual_seed(1142)
    rng = np.random.default_rng(9000 + (kind == 'clock'))
    model = TinyDenoiser()
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    model.train()
    for _ in range(250):
        x, y = make_batch(kind, rng)
        loss = ((model(x)-y)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    model.eval(); vals=[]
    with torch.no_grad():
        for _ in range(12):
            x,y=make_batch(kind,rng)
            vals.append(float(((model(x)-y)**2).mean()))
    return float(np.mean(vals)), float(np.std(vals)/np.sqrt(len(vals)))

if __name__ == '__main__':
    out={'iid_gaussian':run_one('iid'), 'tempered_stable_clock_blind':run_one('clock')}
    print(json.dumps(out, indent=2))
    with open('denoiser_results.json','w') as f: json.dump(out,f,indent=2)