"""Run a restartable sequence-model benchmark family from the command line."""

from __future__ import annotations

import argparse
from pathlib import Path
import sys

import numpy as np
import pandas as pd

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.experiments.deep import run_grouped_sequence_benchmark

MODEL_FAMILIES = {
    "recurrent": ("vanilla_lstm", "bidirectional_lstm", "gru", "attention_lstm"),
    "transformer": ("battery_gpt", "temporal_fusion_transformer", "itransformer", "physics_itransformer"),
    "graph_variational": ("dynamic_graph_itransformer", "vae_lstm"),
}


def _read_csv_or_empty(path: Path) -> pd.DataFrame:
    return pd.read_csv(path) if path.exists() else pd.DataFrame()


def _atomic_write_csv(frame: pd.DataFrame, path: Path) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    frame.to_csv(temporary, index=False)
    temporary.replace(path)


def _model_is_complete(
    metrics: pd.DataFrame,
    predictions: pd.DataFrame,
    model_id: str,
    seeds: tuple[int, ...],
    n_splits: int,
) -> bool:
    if metrics.empty or predictions.empty or "model" not in metrics or "model" not in predictions:
        return False
    subset = metrics[metrics["model"] == model_id]
    pred_subset = predictions[predictions["model"] == model_id]
    expected = {(seed, fold) for seed in seeds for fold in range(1, n_splits + 1)}
    observed = set(zip(subset.get("seed", []), subset.get("fold", [])))
    return observed == expected and not pred_subset.empty


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("dataset", choices=("nasa", "calce", "oxford"))
    parser.add_argument("family", choices=tuple(MODEL_FAMILIES))
    parser.add_argument("--project-root", type=Path, default=PROJECT_ROOT)
    parser.add_argument("--max-epochs", type=int, default=200)
    parser.add_argument("--patience", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--seeds", type=int, nargs="+", default=(17, 42, 2026))
    args = parser.parse_args()

    feature_dir = args.project_root / "artifacts" / "v3" / "features" / args.dataset
    result_dir = args.project_root / "artifacts" / "v3" / "results"
    result_dir.mkdir(parents=True, exist_ok=True)
    X = np.load(feature_dir / "sequences.npz")["X"]
    index = pd.read_csv(feature_dir / "sequence_index.csv")
    prefix = result_dir / f"{args.dataset}_{args.family}"
    metrics_path = Path(f"{prefix}_fold_metrics.csv")
    predictions_path = Path(f"{prefix}_predictions.csv")
    metrics = _read_csv_or_empty(metrics_path)
    predictions = _read_csv_or_empty(predictions_path)
    seeds = tuple(args.seeds)
    n_splits = min(5, index["battery_id"].nunique())
    dataset_name = args.dataset.upper() if args.dataset == "nasa" else args.dataset.title()

    for model_id in MODEL_FAMILIES[args.family]:
        if _model_is_complete(metrics, predictions, model_id, seeds, n_splits):
            print(f"[{dataset_name}] model={model_id} checkpoint complete; skipping", flush=True)
            continue
        model_metrics, model_predictions = run_grouped_sequence_benchmark(
            X,
            index,
            dataset_name=dataset_name,
            n_splits=n_splits,
            seeds=seeds,
            max_epochs=args.max_epochs,
            patience=args.patience,
            batch_size=args.batch_size,
            model_ids=(model_id,),
        )
        if not metrics.empty and "model" in metrics:
            metrics = metrics[metrics["model"] != model_id]
        if not predictions.empty and "model" in predictions:
            predictions = predictions[predictions["model"] != model_id]
        metrics = pd.concat([metrics, model_metrics], ignore_index=True)
        predictions = pd.concat([predictions, model_predictions], ignore_index=True)
        _atomic_write_csv(metrics, metrics_path)
        _atomic_write_csv(predictions, predictions_path)
        print(f"[{dataset_name}] model={model_id} checkpoint written", flush=True)

    print(metrics.groupby("model")[["mae", "rmse", "r2", "within_5pp", "epochs"]].mean().sort_values("mae"))


if __name__ == "__main__":
    main()