from pathlib import Path
import sys

import numpy as np
import torch
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.ace2 import build_model, hard_correct, init_distributed, load_config, seed_all


class Windows(Dataset):
    def __init__(self, path):
        data = np.load(path)
        self.state = data["state"]
        self.forcing = data["forcing"]
        self.windows = [(n, t) for n in range(len(self.state)) for t in range(self.state.shape[1] - 2)]

    def __len__(self):
        return len(self.windows)

    def __getitem__(self, index):
        n, t = self.windows[index]
        return tuple(torch.from_numpy(x.astype(np.float32)) for x in (
            self.state[n, t], self.state[n, t + 1], self.state[n, t + 2],
            self.forcing[n, t + 1], self.forcing[n, t + 2]))


def main():
    cfg = load_config(ROOT)
    seed_all(cfg["seed"])
    distributed, rank, device = init_distributed()
    dataset = Windows(ROOT / cfg["data"]["path"])
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler,
                        shuffle=sampler is None, num_workers=0)
    model = build_model(cfg).to(device)
    if distributed:
        model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None)
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["train"]["learning_rate"])
    for epoch in range(cfg["train"]["epochs"]):
        if sampler:
            sampler.set_epoch(epoch)
        total = 0.0
        for x0, y1, y2, f1, f2 in loader:
            x0, y1, y2, f1, f2 = (x.to(device) for x in (x0, y1, y2, f1, f2))
            p1 = hard_correct(x0, model(x0, f1))
            p2 = hard_correct(p1, model(p1, f2))
            loss = torch.mean((p1 - y1) ** 2) + cfg["train"]["two_step_weight"] * torch.mean((p2 - y2) ** 2)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            total += loss.item()
        if rank == 0:
            print(f"epoch={epoch + 1} two_step_loss={total / len(loader):.7f}")
    if rank == 0:
        path = ROOT / cfg["train"]["checkpoint"]
        path.parent.mkdir(parents=True, exist_ok=True)
        raw_model = model.module if distributed else model
        torch.save({"model": raw_model.state_dict(), "model_config": cfg["model"], "format_version": cfg["data"]["format_version"]}, path)
        metrics = ROOT / "result/training/metrics.json"
        metrics.parent.mkdir(parents=True, exist_ok=True)
        metrics.write_text(__import__("json").dumps({"history": [{"epoch": cfg["train"]["epochs"], "loss": total / len(loader)}], "world_size": int(__import__("os").environ.get("WORLD_SIZE", "1"))}, indent=2))
        print(f"saved {path}")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()