#!/usr/bin/env python3
"""Run the preregistered grouped model comparison for the UNNS Gravity Bridge."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import GroupKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

REQUIRED_PROFILE_COLUMNS = {"galaxy_id", "radius_kpc", "g_bar", "g_obs"}


def load_profiles(path: Path, sample_flag: str = "is_primary") -> pd.DataFrame:
    if path.is_dir():
        files = sorted(path.glob("*.csv"))
        if not files:
            raise FileNotFoundError(f"No CSV files in {path}")
        frame = pd.concat([pd.read_csv(p) for p in files], ignore_index=True)
    else:
        frame = pd.read_csv(path)
    missing = REQUIRED_PROFILE_COLUMNS - set(frame.columns)
    if missing:
        raise ValueError(f"Profile data missing columns: {sorted(missing)}")
    frame = frame.copy()
    frame["galaxy_id"] = frame["galaxy_id"].astype(str)
    if sample_flag and sample_flag.upper() != "ALL":
        if sample_flag not in frame.columns:
            raise ValueError(f"Profile data missing registered sample flag: {sample_flag}")
        flag = frame[sample_flag]
        if flag.dtype != bool:
            flag = flag.astype(str).str.lower().map({"true": True, "false": False})
        if flag.isna().any():
            raise ValueError(f"Sample flag {sample_flag} contains non-boolean values")
        frame = frame.loc[flag.to_numpy(dtype=bool)].copy()
        if frame.empty:
            raise ValueError(f"No profile rows remain for sample flag {sample_flag}")
    for col in ("radius_kpc", "g_bar", "g_obs"):
        frame[col] = pd.to_numeric(frame[col], errors="coerce")
    mask = np.isfinite(frame["radius_kpc"]) & np.isfinite(frame["g_bar"]) & np.isfinite(frame["g_obs"])
    mask &= (frame["radius_kpc"] > 0) & (frame["g_bar"] > 0) & (frame["g_obs"] > 0)
    frame = frame.loc[mask].copy()
    if frame.empty:
        raise ValueError("No finite-positive profile rows remain")
    frame["r_norm"] = frame["radius_kpc"] / frame.groupby("galaxy_id")["radius_kpc"].transform("max")
    frame["log_g_bar"] = np.log10(frame["g_bar"])
    frame["target_y"] = np.log10(frame["g_obs"] / frame["g_bar"])
    return frame


def load_inputs(profiles_path: Path, signatures_path: Path, sample_flag: str = "is_primary") -> pd.DataFrame:
    profiles = load_profiles(profiles_path, sample_flag=sample_flag)
    signatures = pd.read_csv(signatures_path)
    if signatures.empty:
        raise ValueError("Signature table is empty")
    signatures["galaxy_id"] = signatures["galaxy_id"].astype(str)
    merged = profiles.merge(signatures, on="galaxy_id", how="inner", validate="many_to_one")
    missing_galaxies = sorted(set(profiles["galaxy_id"]) - set(merged["galaxy_id"]))
    if missing_galaxies:
        raise ValueError(f"Profiles without signature rows: {missing_galaxies[:20]}")
    return merged


def feature_sets(spec: dict[str, Any]) -> dict[str, tuple[list[str], list[str]]]:
    ordinary = spec["ordinary_features"]
    i_num = spec["struc_i_numeric_features"]
    i_cat = spec["struc_i_categorical_features"]
    p_num = spec["struc_perc_numeric_features"]
    p_cat = spec["struc_perc_categorical_features"]
    return {
        "M0": (["log_g_bar"], []),
        "M1": (["log_g_bar", "r_norm"], []),
        "M2": (["log_g_bar", "r_norm", *ordinary], []),
        "M_I": (["log_g_bar", "r_norm", *i_num], i_cat),
        "M_P": (["log_g_bar", "r_norm", *p_num], p_cat),
        "M_IP": (["log_g_bar", "r_norm", *i_num, *p_num], [*i_cat, *p_cat]),
    }


def check_features(frame: pd.DataFrame, numeric: list[str], categorical: list[str], model_name: str) -> None:
    missing = [c for c in numeric + categorical if c not in frame.columns]
    if missing:
        raise ValueError(f"{model_name}: missing columns {missing}")


def make_pipeline(numeric: list[str], categorical: list[str], alpha: float) -> Pipeline:
    transformers = []
    if numeric:
        transformers.append(("numeric", Pipeline([
            ("imputer", SimpleImputer(strategy="median", keep_empty_features=True)),
            ("scale", StandardScaler()),
        ]), numeric))
    if categorical:
        transformers.append(("categorical", Pipeline([
            ("imputer", SimpleImputer(strategy="most_frequent", keep_empty_features=True)),
            ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
        ]), categorical))
    prep = ColumnTransformer(transformers=transformers, remainder="drop")
    return Pipeline([("preprocess", prep), ("model", Ridge(alpha=alpha))])


def grouped_predictions(frame: pd.DataFrame, numeric: list[str], categorical: list[str], alpha: float, requested_splits: int) -> tuple[pd.DataFrame, pd.DataFrame]:
    groups = frame["galaxy_id"].astype(str)
    unique_groups = groups.nunique()
    if unique_groups < 3:
        raise ValueError(f"Need at least three galaxies; found {unique_groups}")
    n_splits = min(requested_splits, unique_groups)
    splitter = GroupKFold(n_splits=n_splits)
    x = frame[numeric + categorical]
    y = frame["target_y"].to_numpy(dtype=float)

    pred_rows = []
    fold_rows = []
    for fold, (train_idx, test_idx) in enumerate(splitter.split(x, y, groups), start=1):
        train_groups = set(groups.iloc[train_idx])
        test_groups = set(groups.iloc[test_idx])
        overlap = train_groups & test_groups
        if overlap:
            raise RuntimeError(f"Group leakage in fold {fold}: {sorted(overlap)}")
        pipe = make_pipeline(numeric, categorical, alpha)
        pipe.fit(x.iloc[train_idx], y[train_idx])
        pred = pipe.predict(x.iloc[test_idx])
        truth = y[test_idx]
        rmse = math.sqrt(mean_squared_error(truth, pred))
        mae = mean_absolute_error(truth, pred)
        r2 = r2_score(truth, pred) if len(truth) > 1 else float("nan")
        fold_rows.append({
            "fold": fold, "n_train_rows": len(train_idx), "n_test_rows": len(test_idx),
            "n_train_galaxies": len(train_groups), "n_test_galaxies": len(test_groups),
            "rmse": rmse, "mae": mae, "r2": r2,
        })
        test = frame.iloc[test_idx]
        for pos, (_, row) in enumerate(test.iterrows()):
            pred_rows.append({
                "fold": fold,
                "galaxy_id": row["galaxy_id"],
                "radius_kpc": row["radius_kpc"],
                "r_norm": row["r_norm"],
                "log_g_bar": row["log_g_bar"],
                "y_true": truth[pos],
                "y_pred": pred[pos],
                "residual": truth[pos] - pred[pos],
            })
    return pd.DataFrame(fold_rows), pd.DataFrame(pred_rows)


def evaluate_model(frame: pd.DataFrame, model_name: str, numeric: list[str], categorical: list[str], alpha: float, requested_splits: int) -> tuple[pd.DataFrame, pd.DataFrame, dict[str, Any]]:
    check_features(frame, numeric, categorical, model_name)
    folds, predictions = grouped_predictions(frame, numeric, categorical, alpha, requested_splits)
    summary = {
        "model": model_name,
        "status": "COMPLETE",
        "numeric_features": ";".join(numeric),
        "categorical_features": ";".join(categorical),
        "folds": len(folds),
        "rmse_mean": folds["rmse"].mean(),
        "rmse_sd": folds["rmse"].std(ddof=1) if len(folds) > 1 else 0.0,
        "mae_mean": folds["mae"].mean(),
        "mae_sd": folds["mae"].std(ddof=1) if len(folds) > 1 else 0.0,
        "r2_mean": folds["r2"].mean(),
        "r2_sd": folds["r2"].std(ddof=1) if len(folds) > 1 else 0.0,
    }
    folds.insert(0, "model", model_name)
    predictions.insert(0, "model", model_name)
    return folds, predictions, summary


def run_all(frame: pd.DataFrame, spec: dict[str, Any]) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    alpha = float(spec["estimator"]["alpha"])
    requested_splits = int(spec["cross_validation"]["requested_splits"])
    all_folds = []
    all_predictions = []
    summaries = []
    for model_name, (numeric, categorical) in feature_sets(spec).items():
        folds, predictions, summary = evaluate_model(frame, model_name, numeric, categorical, alpha, requested_splits)
        all_folds.append(folds)
        all_predictions.append(predictions)
        summaries.append(summary)
    return pd.DataFrame(summaries), pd.concat(all_folds, ignore_index=True), pd.concat(all_predictions, ignore_index=True)


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[1]
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--profiles", type=Path, default=root / "data/prepared_profiles/galaxy_profiles.csv")
    p.add_argument("--signatures", type=Path, default=root / "signatures/chamber_signatures.csv")
    p.add_argument("--spec", type=Path, default=root / "config/analysis_specification.json")
    p.add_argument("--output-dir", type=Path, default=root / "outputs/tables")
    p.add_argument(
        "--sample-flag",
        default="is_primary",
        help="Boolean profile column selecting the registered sample; use ALL to disable filtering.",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    spec = json.loads(args.spec.read_text(encoding="utf-8"))
    frame = load_inputs(args.profiles, args.signatures, sample_flag=args.sample_flag)
    summaries, folds, predictions = run_all(frame, spec)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    summaries.to_csv(args.output_dir / "model_comparison.csv", index=False)
    folds.to_csv(args.output_dir / "model_comparison_folds.csv", index=False)
    predictions.to_csv(args.output_dir / "heldout_predictions.csv", index=False)
    assignments = predictions[["fold", "galaxy_id"]].drop_duplicates().sort_values(["fold", "galaxy_id"])
    assignments.to_csv(args.output_dir / "fold_assignments.csv", index=False)
    print(summaries[["model", "rmse_mean", "mae_mean", "r2_mean"]].to_string(index=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
