from __future__ import annotations

import argparse
import csv
import time
from pathlib import Path

import torch

from common import (
    ROOT,
    build_model,
    count_parameters,
    load_config,
    make_normalized_loaders,
    relative_l2_per_sample,
    relative_l2_sum,
    resolve,
    save_json,
    select_device,
    set_seed,
)


def evaluate(model, loader, y_normalizer, device):
    model.eval()
    total = 0.0
    count = 0
    with torch.no_grad():
        for x, y_norm in loader:
            x = x.to(device, non_blocking=True)
            y_norm = y_norm.to(device, non_blocking=True)
            pred = y_normalizer.decode(model(x))
            target = y_normalizer.decode(y_norm)
            rel = relative_l2_per_sample(pred, target)
            total += float(rel.sum().item())
            count += int(rel.numel())
    return total / max(1, count)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train HO-FNO on Pipe horizontal velocity.")
    parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml"))
    args = parser.parse_args()

    cfg = load_config(args.config)
    set_seed(int(cfg["experiment"].get("seed", 2026)))
    output_dir = resolve(cfg["paths"]["output_dir"])
    weight_path = resolve(cfg["paths"]["checkpoint"], output_dir)
    results_dir = resolve(cfg["paths"]["results_dir"], output_dir)
    weight_path.parent.mkdir(parents=True, exist_ok=True)
    results_dir.mkdir(parents=True, exist_ok=True)

    require_accelerator = bool(cfg["runtime"].get("require_accelerator", False))
    device = select_device(require_accelerator=require_accelerator)
    print(f"device={device}", flush=True)

    train_loader, test_loader, x_normalizer, y_normalizer = make_normalized_loaders(cfg)
    x_normalizer.to(device)
    y_normalizer.to(device)

    model = build_model(cfg).to(device)
    print(f"parameters={count_parameters(model)}", flush=True)

    train_cfg = cfg["training"]
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(train_cfg["learning_rate"]),
        weight_decay=float(train_cfg.get("weight_decay", 0.0)),
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=int(train_cfg["epochs"]),
        eta_min=float(train_cfg["final_learning_rate"]),
    )

    history_path = results_dir / "train_history.csv"
    best_metric = float("inf")
    best_epoch = -1
    train_size = int(cfg["data"]["train_size"])
    epochs = int(train_cfg["epochs"])
    start_time = time.time()

    with history_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "test_mean_relative_l2", "lr", "seconds"])
        writer.writeheader()
        for epoch in range(1, epochs + 1):
            model.train()
            epoch_start = time.time()
            train_loss_sum = 0.0
            for x, y_norm in train_loader:
                x = x.to(device, non_blocking=True)
                y_norm = y_norm.to(device, non_blocking=True)
                optimizer.zero_grad(set_to_none=True)
                pred = y_normalizer.decode(model(x))
                target = y_normalizer.decode(y_norm)
                loss = relative_l2_sum(pred, target)
                loss.backward()
                optimizer.step()
                train_loss_sum += float(loss.item())

            scheduler.step()
            train_loss = train_loss_sum / train_size
            test_l2 = evaluate(model, test_loader, y_normalizer, device)
            lr = float(optimizer.param_groups[0]["lr"])
            seconds = time.time() - epoch_start
            writer.writerow(
                {
                    "epoch": epoch,
                    "train_loss": train_loss,
                    "test_mean_relative_l2": test_l2,
                    "lr": lr,
                    "seconds": seconds,
                }
            )
            handle.flush()
            print(
                f"epoch={epoch:04d} train_loss={train_loss:.8f} "
                f"test_mean_relative_l2={test_l2:.8f} lr={lr:.8e}",
                flush=True,
            )

            if test_l2 < best_metric:
                best_metric = test_l2
                best_epoch = epoch
                torch.save(
                    {
                        "model_state": model.state_dict(),
                        "config": cfg,
                        "epoch": epoch,
                        "best_metric": best_metric,
                        "x_normalizer": x_normalizer.state_dict(),
                        "y_normalizer": y_normalizer.state_dict(),
                    },
                    weight_path,
                )
                save_json(
                    results_dir / "best_metrics.json",
                    {"best_epoch": best_epoch, "test_mean_relative_l2": best_metric, "checkpoint": str(weight_path)},
                )

    torch.save(
        {
            "model_state": model.state_dict(),
            "config": cfg,
            "epoch": epochs,
            "best_metric": best_metric,
            "x_normalizer": x_normalizer.state_dict(),
            "y_normalizer": y_normalizer.state_dict(),
        },
        weight_path.parent / "last_model.pth",
    )
    save_json(
        results_dir / "training_summary.json",
        {
            "epochs": epochs,
            "best_epoch": best_epoch,
            "best_test_mean_relative_l2": best_metric,
            "total_seconds": time.time() - start_time,
            "checkpoint": str(weight_path),
        },
    )
    print(f"best_checkpoint={weight_path}", flush=True)
    print(f"best_epoch={best_epoch} best_test_mean_relative_l2={best_metric:.8f}", flush=True)


if __name__ == "__main__":
    main()