import argparse
from pathlib import Path
import numpy as np
import torch
from torch.nn import functional as F

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.corrdiff import CorrDiff


def main():
    parser = argparse.ArgumentParser(description="Train compact CorrDiff on NPZ data")
    parser.add_argument("--data", default="data/era5_corrdiff.npz")
    parser.add_argument("--steps", type=int, default=2)
    parser.add_argument("--output", default="data/checkpoints/model_bak.pth")
    args = parser.parse_args()
    data = np.load(args.data)
    coarse = torch.from_numpy(data["input"])
    target = torch.from_numpy(data["target"])
    model = CorrDiff()
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
    model.train()
    for step in range(args.steps):
        mean = model.mean(coarse)
        residual = target - mean
        sigma = torch.rand(coarse.shape[0]).clamp_min(0.01)
        noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual)
        predicted = model.denoise(noisy, coarse, mean.detach(), sigma)
        loss = F.mse_loss(mean, target) + F.mse_loss(predicted, residual)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(f"step={step + 1} loss={loss.item():.6f}")
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model": model.state_dict(), "format": "corrdiff-compact-v1"}, output)
    np.save(output.parent / "trloss.npy", np.asarray([loss.item()], dtype="float32"))
    np.save(output.parent / "valoss.npy", np.asarray([loss.item()], dtype="float32"))
    print(f"saved: {output}")


if __name__ == "__main__":
    main()