from __future__ import annotations

import argparse
import copy
from pathlib import Path

import torch

from common import (
    append_history,
    build_model,
    collect_predictions,
    count_parameters,
    get_device,
    load_config,
    make_loaders,
    mean_relative_l2_percent,
    relative_normalized_mse,
    resolve_output_path,
    save_json,
    set_seed,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml"))
    parser.add_argument("--dry-run", action="store_true", help="Run one forward/backward pass for preflight only.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    set_seed(cfg["training"]["seed"], cfg["runtime"]["deterministic"], cfg["runtime"]["benchmark"])
    device = get_device(cfg)

    output_dir = Path(cfg["paths"]["output_dir"])
    weight_path = resolve_output_path(cfg, cfg["paths"]["checkpoint"])
    results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    weight_path.parent.mkdir(parents=True, exist_ok=True)
    results_dir.mkdir(parents=True, exist_ok=True)

    train_loader, val_loader, test_loader = make_loaders(cfg)
    model = build_model(cfg).to(device)
    params = count_parameters(model)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=cfg["training"]["learning_rate"],
        weight_decay=cfg["training"]["weight_decay"],
    )
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer,
        step_size=cfg["training"]["scheduler_step_size"],
        gamma=cfg["training"]["scheduler_gamma"],
    )

    if args.dry_run:
        model.train()
        x, y = next(iter(train_loader))
        x = x.to(device)
        y = y.to(device)
        loss = relative_normalized_mse(model(x), y)
        loss.backward()
        print(f"DRY_RUN forward/backward passed loss={loss.item():.8f} params={params}", flush=True)
        return

    best_val = float("inf")
    best_test = float("inf")
    best_epoch = -1
    best_state = None
    patience = cfg["training"]["early_stopping_patience"]
    stale_epochs = 0
    history_path = results_dir / "history.csv"

    for epoch in range(1, cfg["training"]["epochs"] + 1):
        model.train()
        running_loss = 0.0
        batches = 0
        for x, y in train_loader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            pred = model(x)
            loss = relative_normalized_mse(pred, y)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            batches += 1

        train_loss = running_loss / max(1, batches)
        val_l2 = mean_relative_l2_percent(model, val_loader, device)
        test_l2 = mean_relative_l2_percent(model, test_loader, device)
        lr = optimizer.param_groups[0]["lr"]
        row = {
            "epoch": epoch,
            "train_relative_normalized_mse": train_loss,
            "val_mean_relative_l2_percent": val_l2,
            "test_mean_relative_l2_percent": test_l2,
            "lr": lr,
        }
        append_history(history_path, row)
        print(
            f"epoch={epoch:04d} train_loss={train_loss:.8f} "
            f"val_mean_relative_l2={val_l2:.6f}% test_mean_relative_l2={test_l2:.6f}% lr={lr:.8g}",
            flush=True,
        )

        if val_l2 < best_val:
            best_val = val_l2
            best_test = test_l2
            best_epoch = epoch
            best_state = copy.deepcopy(model.state_dict())
            torch.save(
                {
                    "model_state": best_state,
                    "config": cfg,
                    "epoch": best_epoch,
                    "best_val_mean_relative_l2_percent": best_val,
                    "test_mean_relative_l2_percent_at_best_val": best_test,
                    "parameter_count": params,
                },
                weight_path,
            )
            stale_epochs = 0
        else:
            stale_epochs += 1

        save_json(
            results_dir / "metrics.json",
            {
                "status": "running",
                "epoch": epoch,
                "best_epoch": best_epoch,
                "best_val_mean_relative_l2_percent": best_val,
                "test_mean_relative_l2_percent_at_best_val": best_test,
                "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"],
                "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"],
                "parameter_count": params,
            },
        )
        scheduler.step()

        if stale_epochs >= patience:
            print(
                f"early_stopping epoch={epoch} best_epoch={best_epoch} "
                f"best_val={best_val:.6f}% best_test={best_test:.6f}%",
                flush=True,
            )
            break

    if best_state is not None:
        model.load_state_dict(best_state)
    inputs, preds, targets = collect_predictions(model, test_loader, device)
    pred_path = results_dir / "test_predictions_from_train.npz"
    import numpy as np

    np.savez_compressed(pred_path, inputs=inputs, predictions=preds, targets=targets)
    save_json(
        results_dir / "metrics.json",
        {
            "status": "complete",
            "best_epoch": best_epoch,
            "best_val_mean_relative_l2_percent": best_val,
            "test_mean_relative_l2_percent_at_best_val": best_test,
            "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"],
            "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"],
            "parameter_count": params,
            "checkpoint": str(weight_path),
            "prediction_file": str(pred_path),
        },
    )
    print(
        f"training_complete best_epoch={best_epoch} best_val={best_val:.6f}% "
        f"test_at_best_val={best_test:.6f}% checkpoint={weight_path}",
        flush=True,
    )


if __name__ == "__main__":
    main()