"""Fit the paper's six station-wise model configurations."""

import json
import os
import sys
from pathlib import Path

import numpy as np
import torch
import yaml
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.globalsurgeml import GlobalSurgeML


CONFIGURATIONS = {
    "LR-RS": ("rs_daily", "linear"),
    "LR-RS-lag": ("rs_lagged", "linear"),
    "RF-RS-lag": ("rs_lagged", "random_forest"),
    "LR-AR": ("ar_daily", "linear"),
    "LR-AR-lag": ("ar_lagged", "linear"),
    "RF-AR-lag": ("ar_lagged", "random_forest"),
}


class SurgeDataset(Dataset):
    def __init__(self, path, config):
        self.data = np.load(path)
        expected = config["data"]
        if str(self.data["format_version"]) != expected["format_version"]:
            raise ValueError("incompatible storm-surge data format")
        dimensions = {"rs_daily": expected["rs_daily_features"], "rs_lagged": expected["rs_lagged_features"],
                      "ar_daily": expected["ar_daily_features"], "ar_lagged": expected["ar_lagged_features"]}
        for key, width in dimensions.items():
            if self.data[key].shape[1:] != (int(width),):
                raise ValueError(f"{key} must have shape [N,{width}]")
        if self.data["targets_m"].shape[1:] != (1,):
            raise ValueError("targets_m must have shape [N,1]")

    def __len__(self):
        return len(self.data["targets_m"])

    def __getitem__(self, index):
        return {key: torch.from_numpy(self.data[key][index]).float() for key in
                ("rs_daily", "rs_lagged", "ar_daily", "ar_lagged", "targets_m")}


def device_from_config(config, rank=0):
    requested = config["runtime"]["device"]
    if requested == "auto":
        return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu")
    return torch.device(requested)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    seed = int(config["seed"])
    np.random.seed(seed)
    torch.manual_seed(seed)
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
    rank = torch.distributed.get_rank() if distributed else 0
    device = device_from_config(config, local_rank)
    if device.type == "cuda":
        torch.cuda.set_device(device)
    dataset = SurgeDataset(ROOT / config["data"]["root"] / "train.npz", config)
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler,
                        shuffle=sampler is None, num_workers=int(config["train"]["num_workers"]))
    full = dataset.data
    means, scales, states, history = {}, {}, {}, []
    models = {}
    for model_index, (name, (feature_key, method)) in enumerate(CONFIGURATIONS.items()):
        array = full[feature_key].astype(np.float32)
        means[name] = array.mean(0).astype(np.float32)
        scales[name] = array.std(0).clip(1e-6).astype(np.float32)
        standardized = (array - means[name]) / scales[name]
        model = GlobalSurgeML(array.shape[1], method, config["model"], seed + model_index).to(device)
        if method == "linear":
            model.regressor.select_features(standardized, full["targets_m"][:, 0],
                                             float(config["model"]["p_value_threshold"]))
            wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model
            optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"]),
                                         weight_decay=float(config["train"]["weight_decay"]))
            for epoch in range(int(config["train"]["epochs"])):
                if sampler:
                    sampler.set_epoch(epoch)
                total, steps = 0.0, 0
                for batch in loader:
                    features = (batch[feature_key].to(device) - torch.from_numpy(means[name]).to(device)) / torch.from_numpy(scales[name]).to(device)
                    target = batch["targets_m"].to(device)
                    prediction = wrapped(features)
                    loss = torch.nn.functional.mse_loss(prediction, target)
                    optimizer.zero_grad(set_to_none=True)
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(wrapped.parameters(), float(config["train"]["gradient_clip_norm"]))
                    optimizer.step()
                    total += float(loss.detach())
                    steps += 1
                history.append({"model": name, "epoch": epoch + 1, "mse_m2": total / max(steps, 1)})
            model = wrapped.module if distributed else wrapped
        else:
            model.regressor.fit(standardized, full["targets_m"][:, 0])
            prediction = model(torch.from_numpy(standardized).to(device))
            history.append({"model": name, "epoch": 1,
                            "mse_m2": float(torch.nn.functional.mse_loss(prediction.cpu(), torch.from_numpy(full["targets_m"])).item())})
        models[name] = model
        states[name] = model.state_dict()
    if rank == 0:
        checkpoint = ROOT / config["paths"]["checkpoint"]
        metrics = ROOT / config["paths"]["training_metrics"]
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        metrics.parent.mkdir(parents=True, exist_ok=True)
        torch.save({"model": states, "model_config": config["model"], "configurations": CONFIGURATIONS,
                    "feature_means": means, "feature_scales": scales,
                    "format_version": config["data"]["format_version"], "target_unit": "m"}, checkpoint)
        metrics.write_text(json.dumps({"history": history}, indent=2) + "\n")
        print(f"checkpoint={checkpoint.relative_to(ROOT)} models={len(states)}")
    if distributed:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()