"""Train the DnCNN denoiser on synthetic clean/noisy patches (self-contained, CPU).

Generates synthetic clean microscopy images, adds Poisson+Gaussian noise, extracts
patches, and trains the net to predict the noise (residual learning). Saves weights.

Usage:
  python -m scripts.train_denoiser <out_path> <n_images> <epochs>
"""
from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

from core import model
from core.synth import add_noise, clean_image


def _patches(n_images: int, patch: int, per_img: int, rng):
    X, Y = [], []
    for i in range(n_images):
        clean = clean_image(seed=i)
        noisy = add_noise(clean, sigma=rng.uniform(0.08, 0.16), seed=1000 + i)
        H, W = clean.shape
        for _ in range(per_img):
            y0, x0 = rng.integers(0, H - patch), rng.integers(0, W - patch)
            c = clean[y0:y0 + patch, x0:x0 + patch]
            n = noisy[y0:y0 + patch, x0:x0 + patch]
            X.append(n)
            Y.append(n - c)                              # target = noise residual
    return np.stack(X)[:, None], np.stack(Y)[:, None]


def main() -> int:
    out_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("dncnn.pt")
    n_images = int(sys.argv[2]) if len(sys.argv) > 2 else 24
    epochs = int(sys.argv[3]) if len(sys.argv) > 3 else 15

    import torch
    torch.manual_seed(0)
    rng = np.random.default_rng(0)

    X, Y = _patches(n_images, patch=64, per_img=40, rng=rng)
    print(f"patches: {len(X)}")
    Xt = torch.from_numpy(X.astype(np.float32))
    Yt = torch.from_numpy(Y.astype(np.float32))

    net = model.build_net()
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    lossf = torch.nn.MSELoss()
    bs = 32
    for ep in range(epochs):
        net.train()
        perm = torch.randperm(len(Xt))
        tot = 0.0
        for j in range(0, len(perm), bs):
            b = perm[j:j + bs]
            opt.zero_grad()
            loss = lossf(net(Xt[b]), Yt[b])
            loss.backward(); opt.step()
            tot += float(loss) * len(b)
        print(f"epoch {ep+1:2d}  mse={tot/len(Xt):.5f}")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(net.state_dict(), out_path)
    print(f"saved {out_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())