"""Exploratory leakage-safe grouped benchmark for the ENGIN diesel dataset.

EXPERIMENT STATUS: SUPERSEDED. Its metrics describe baseline candidates, not
the model shipped by ``final_pipeline.py``. Use ``ml_polish_benchmark.py`` for
the reproducibility numbers of the final architecture.

The benchmark compares three spectral representations:

* raw: the 21 measured amplitudes,
* relative: cylinder spectrum minus the median spectrum of its engine,
* combined: raw and relative features together.

Every validation split contains complete, previously unseen engines.  Relative
features use only cylinders from the engine currently being transformed and do
not use labels, so the same operation is available for an unseen test engine.
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
from sklearn.model_selection import GroupKFold, StratifiedGroupKFold
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from engin.config import FAULT_LABELS, FREQ_COLS, LABELS, NOT_APPLICABLE, SEVERITIES

RANDOM_STATE = 42
FEATURE_SETS = ("raw", "relative", "combined")
MODEL_NAMES = ("logistic", "extra_trees")


class SpectrumFeatures(BaseEstimator, TransformerMixin):
    """Create raw and engine-relative spectral features.

    NaNs are interpolated independently within each row along frequency.  The
    training-set medians stored by ``fit`` are used only as a final fallback for
    a completely empty row, which is not present in the supplied data.
    """

    def __init__(self, feature_set: str = "raw") -> None:
        self.feature_set = feature_set

    def fit(self, X: pd.DataFrame, y: Iterable[str] | None = None):
        self._validate_input(X)
        if self.feature_set not in FEATURE_SETS:
            raise ValueError(
                f"Unknown feature_set={self.feature_set!r}; choose from {FEATURE_SETS}."
            )
        spectra = X[FREQ_COLS].apply(pd.to_numeric, errors="coerce")
        medians = spectra.median(axis=0).to_numpy(dtype=float)
        if np.isnan(medians).any():
            raise ValueError("At least one frequency column is entirely NaN in training.")
        self.fallback_medians_ = medians
        return self

    def transform(self, X: pd.DataFrame) -> np.ndarray:
        self._validate_input(X)
        if not hasattr(self, "fallback_medians_"):
            raise RuntimeError("SpectrumFeatures must be fitted before transform().")

        spectra = X[FREQ_COLS].apply(pd.to_numeric, errors="coerce")
        spectra = spectra.interpolate(axis=1, limit_direction="both")
        spectra = spectra.fillna(
            pd.Series(self.fallback_medians_, index=FREQ_COLS)
        )

        # The median is calculated separately for every engine in X. During
        # validation X contains complete held-out engines, never training rows.
        engine_median = spectra.groupby(X["engine_id"].to_numpy()).transform("median")
        raw = spectra.to_numpy(dtype=float)
        relative = raw - engine_median.to_numpy(dtype=float)

        if self.feature_set == "raw":
            return raw
        if self.feature_set == "relative":
            return relative
        return np.hstack([raw, relative])

    @staticmethod
    def _validate_input(X: pd.DataFrame) -> None:
        required = {"engine_id", *FREQ_COLS}
        missing = sorted(required.difference(X.columns))
        if missing:
            raise ValueError(f"Missing input columns: {missing}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", type=Path, default=Path("val.csv"))
    parser.add_argument(
        "--output-dir", type=Path, default=Path("benchmark_outputs")
    )
    parser.add_argument("--n-splits", type=int, default=5)
    parser.add_argument("--random-state", type=int, default=RANDOM_STATE)
    parser.add_argument(
        "--splitter",
        choices=("stratified-group", "group"),
        default="stratified-group",
        help="Both splitters keep complete engines together.",
    )
    parser.add_argument(
        "--feature-sets",
        nargs="+",
        choices=FEATURE_SETS,
        default=list(FEATURE_SETS),
    )
    parser.add_argument(
        "--models", nargs="+", choices=MODEL_NAMES, default=list(MODEL_NAMES)
    )
    parser.add_argument(
        "--n-jobs",
        type=int,
        default=-1,
        help="Parallel workers for Extra Trees; -1 uses all available CPU cores.",
    )
    return parser.parse_args()


def validate_data(df: pd.DataFrame, n_splits: int) -> None:
    required = {
        "engine_id",
        "cylinder",
        "n_cylinders",
        "label",
        "severity",
        *FREQ_COLS,
    }
    missing = sorted(required.difference(df.columns))
    if missing:
        raise ValueError(f"Input is missing columns: {missing}")
    if df.duplicated(["engine_id", "cylinder"]).any():
        raise ValueError("Duplicate engine_id + cylinder keys found.")
    unknown_labels = sorted(set(df["label"].dropna()).difference(LABELS))
    if unknown_labels:
        raise ValueError(f"Unexpected labels: {unknown_labels}")
    if df["engine_id"].nunique() < n_splits:
        raise ValueError("n_splits cannot exceed the number of engines.")

    engine_sizes = df.groupby("engine_id").agg(
        rows=("cylinder", "size"),
        expected=("n_cylinders", "first"),
        size_values=("n_cylinders", "nunique"),
        unique_cylinders=("cylinder", "nunique"),
    )
    complete = (
        (engine_sizes["rows"] == engine_sizes["expected"])
        & (engine_sizes["unique_cylinders"] == engine_sizes["expected"])
        & (engine_sizes["size_values"] == 1)
    )
    if not complete.all():
        bad = engine_sizes.index[~complete].tolist()
        raise ValueError(f"Incomplete or inconsistent engines: {bad}")

    invalid_severity = df.loc[df["label"].isin(FAULT_LABELS), "severity"].dropna()
    invalid_severity = sorted(set(invalid_severity).difference(SEVERITIES))
    if invalid_severity:
        raise ValueError(f"Unexpected fault severities: {invalid_severity}")


def make_splits(
    df: pd.DataFrame,
    splitter_name: str,
    n_splits: int,
    random_state: int,
) -> list[tuple[np.ndarray, np.ndarray]]:
    groups = df["engine_id"].to_numpy()
    if splitter_name == "stratified-group":
        splitter = StratifiedGroupKFold(
            n_splits=n_splits, shuffle=True, random_state=random_state
        )
    else:
        splitter = GroupKFold(n_splits=n_splits)

    splits = list(splitter.split(df, df["label"], groups))
    seen_validation_rows: list[int] = []
    for fold, (train_idx, valid_idx) in enumerate(splits, start=1):
        train_engines = set(df.iloc[train_idx]["engine_id"])
        valid_engines = set(df.iloc[valid_idx]["engine_id"])
        overlap = train_engines.intersection(valid_engines)
        if overlap:
            raise AssertionError(f"Engine leakage in fold {fold}: {sorted(overlap)}")
        seen_validation_rows.extend(valid_idx.tolist())

    if sorted(seen_validation_rows) != list(range(len(df))):
        raise AssertionError("Every row must appear in validation exactly once.")
    return splits


def make_pipeline(
    model_name: str,
    feature_set: str,
    random_state: int,
    n_jobs: int,
) -> Pipeline:
    steps: list[tuple[str, object]] = [
        ("features", SpectrumFeatures(feature_set=feature_set))
    ]
    if model_name == "logistic":
        steps.extend(
            [
                ("scale", StandardScaler()),
                (
                    "model",
                    LogisticRegression(
                        C=1.0,
                        class_weight="balanced",
                        max_iter=5_000,
                        random_state=random_state,
                    ),
                ),
            ]
        )
    elif model_name == "extra_trees":
        steps.append(
            (
                "model",
                ExtraTreesClassifier(
                    n_estimators=500,
                    class_weight="balanced",
                    max_features="sqrt",
                    min_samples_leaf=1,
                    n_jobs=n_jobs,
                    random_state=random_state,
                ),
            )
        )
    else:
        raise ValueError(f"Unknown model: {model_name}")
    return Pipeline(steps)


def macro_f1(y_true: pd.Series | np.ndarray, y_pred: np.ndarray) -> float:
    return float(
        f1_score(y_true, y_pred, labels=LABELS, average="macro", zero_division=0)
    )


def raw_score(label_f1: float, severity_accuracy: float) -> float:
    return 0.75 * label_f1 + 0.25 * severity_accuracy


def ml_points(score: float) -> float:
    if score < 0.80:
        return 0.0
    return min(40.0, 40.0 * (score - 0.80) / 0.20)


def safe_name(value: str) -> str:
    return re.sub(r"[^a-zA-Z0-9_-]+", "_", value)


def evaluate_experiment(
    df: pd.DataFrame,
    splits: list[tuple[np.ndarray, np.ndarray]],
    model_name: str,
    feature_set: str,
    random_state: int,
    n_jobs: int,
    output_dir: Path,
) -> tuple[dict[str, object], list[dict[str, object]], list[dict[str, object]]]:
    experiment_id = f"{feature_set}__{model_name}"
    y_label = df["label"].reset_index(drop=True)
    y_severity = df["severity"].reset_index(drop=True)
    true_fault = y_label.isin(FAULT_LABELS).to_numpy()

    oof_label = np.full(len(df), "", dtype=object)
    oof_severity_oracle = np.full(len(df), "", dtype=object)
    oof_severity_submission = np.full(len(df), NOT_APPLICABLE, dtype=object)
    fold_rows: list[dict[str, object]] = []

    for fold, (train_idx, valid_idx) in enumerate(splits, start=1):
        X_train = df.iloc[train_idx]
        X_valid = df.iloc[valid_idx]
        y_train_label = y_label.iloc[train_idx]

        label_pipeline = make_pipeline(
            model_name, feature_set, random_state + fold, n_jobs
        )
        label_pipeline.fit(X_train, y_train_label)
        label_pred = label_pipeline.predict(X_valid)
        oof_label[valid_idx] = label_pred

        severity_train_mask = y_train_label.isin(FAULT_LABELS).to_numpy()
        severity_pipeline = make_pipeline(
            model_name, feature_set, random_state + 100 + fold, n_jobs
        )
        severity_pipeline.fit(
            X_train.iloc[np.flatnonzero(severity_train_mask)],
            y_severity.iloc[train_idx].iloc[np.flatnonzero(severity_train_mask)],
        )
        severity_pred = severity_pipeline.predict(X_valid)
        oof_severity_oracle[valid_idx] = severity_pred

        predicted_fault = np.isin(label_pred, FAULT_LABELS)
        emitted = np.full(len(valid_idx), NOT_APPLICABLE, dtype=object)
        emitted[predicted_fault] = severity_pred[predicted_fault]
        oof_severity_submission[valid_idx] = emitted

        fold_true_label = y_label.iloc[valid_idx]
        fold_true_fault = fold_true_label.isin(FAULT_LABELS).to_numpy()
        fold_f1 = macro_f1(fold_true_label, label_pred)
        fold_severity_oracle = accuracy_score(
            y_severity.iloc[valid_idx].to_numpy()[fold_true_fault],
            severity_pred[fold_true_fault],
        )
        fold_severity_submission = accuracy_score(
            y_severity.iloc[valid_idx].to_numpy()[fold_true_fault],
            emitted[fold_true_fault],
        )
        fold_raw = raw_score(fold_f1, float(fold_severity_submission))

        fold_rows.append(
            {
                "experiment_id": experiment_id,
                "fold": fold,
                "train_engines": X_train["engine_id"].nunique(),
                "valid_engines": X_valid["engine_id"].nunique(),
                "valid_rows": len(valid_idx),
                "valid_fault_rows": int(fold_true_fault.sum()),
                "macro_f1": fold_f1,
                "severity_accuracy_oracle_label": float(fold_severity_oracle),
                "severity_accuracy_submission": float(fold_severity_submission),
                "raw_score": fold_raw,
                "ml_points": ml_points(fold_raw),
            }
        )

    if (oof_label == "").any() or (oof_severity_oracle == "").any():
        raise AssertionError("Some rows did not receive out-of-fold predictions.")

    overall_f1 = macro_f1(y_label, oof_label)
    severity_oracle_accuracy = float(
        accuracy_score(
            y_severity.to_numpy()[true_fault], oof_severity_oracle[true_fault]
        )
    )
    severity_submission_accuracy = float(
        accuracy_score(
            y_severity.to_numpy()[true_fault], oof_severity_submission[true_fault]
        )
    )
    overall_raw = raw_score(overall_f1, severity_submission_accuracy)

    fold_frame = pd.DataFrame(fold_rows)
    summary: dict[str, object] = {
        "experiment_id": experiment_id,
        "feature_set": feature_set,
        "model": model_name,
        "macro_f1_oof": overall_f1,
        "severity_accuracy_oracle_label_oof": severity_oracle_accuracy,
        "severity_accuracy_submission_oof": severity_submission_accuracy,
        "raw_score_oof": overall_raw,
        "ml_points_oof": ml_points(overall_raw),
        "macro_f1_fold_mean": float(fold_frame["macro_f1"].mean()),
        "macro_f1_fold_std": float(fold_frame["macro_f1"].std(ddof=1)),
        "raw_score_fold_mean": float(fold_frame["raw_score"].mean()),
        "raw_score_fold_std": float(fold_frame["raw_score"].std(ddof=1)),
    }

    class_f1_rows: list[dict[str, object]] = []
    per_class = f1_score(
        y_label, oof_label, labels=LABELS, average=None, zero_division=0
    )
    for label, score in zip(LABELS, per_class):
        class_f1_rows.append(
            {
                "experiment_id": experiment_id,
                "label": label,
                "support": int((y_label == label).sum()),
                "f1": float(score),
            }
        )

    matrix = confusion_matrix(y_label, oof_label, labels=LABELS)
    matrix_frame = pd.DataFrame(matrix, index=LABELS, columns=LABELS)
    matrix_frame.index.name = "true_label"
    matrix_frame.to_csv(output_dir / f"confusion__{safe_name(experiment_id)}.csv")

    oof_frame = df[["engine_id", "cylinder"]].copy()
    oof_frame["fold"] = -1
    for fold, (_, valid_idx) in enumerate(splits, start=1):
        oof_frame.loc[valid_idx, "fold"] = fold
    oof_frame["true_label"] = y_label
    oof_frame["predicted_label"] = oof_label
    oof_frame["true_severity"] = y_severity
    oof_frame["predicted_severity_oracle_label"] = oof_severity_oracle
    oof_frame["predicted_severity_submission"] = oof_severity_submission
    oof_frame.to_csv(
        output_dir / f"oof__{safe_name(experiment_id)}.csv", index=False
    )

    return summary, fold_rows, class_f1_rows


def evaluate_hybrid_combinations(
    df: pd.DataFrame,
    experiment_ids: list[str],
    output_dir: Path,
) -> pd.DataFrame:
    """Combine label and severity OOF predictions without retraining.

    All base experiments use identical folds. Therefore selecting one base
    pipeline for labels and another for severity remains fully out-of-fold and
    gives an honest estimate for a two-model submission architecture.
    """

    y_label = df["label"].reset_index(drop=True)
    y_severity = df["severity"].reset_index(drop=True)
    true_fault = y_label.isin(FAULT_LABELS).to_numpy()
    oof_frames = {
        experiment_id: pd.read_csv(
            output_dir / f"oof__{safe_name(experiment_id)}.csv"
        )
        for experiment_id in experiment_ids
    }

    rows: list[dict[str, object]] = []
    for label_experiment in experiment_ids:
        label_frame = oof_frames[label_experiment]
        label_pred = label_frame["predicted_label"].to_numpy(dtype=object)
        label_f1 = macro_f1(y_label, label_pred)

        for severity_experiment in experiment_ids:
            severity_frame = oof_frames[severity_experiment]
            severity_pred = severity_frame[
                "predicted_severity_oracle_label"
            ].to_numpy(dtype=object)
            emitted = np.full(len(df), NOT_APPLICABLE, dtype=object)
            predicted_fault = np.isin(label_pred, FAULT_LABELS)
            emitted[predicted_fault] = severity_pred[predicted_fault]
            severity_accuracy = float(
                accuracy_score(y_severity.to_numpy()[true_fault], emitted[true_fault])
            )

            fold_raw_scores: list[float] = []
            fold_label_scores: list[float] = []
            for fold in sorted(label_frame["fold"].unique()):
                fold_mask = label_frame["fold"].eq(fold).to_numpy()
                fold_true_fault = true_fault & fold_mask
                fold_label_f1 = macro_f1(y_label[fold_mask], label_pred[fold_mask])
                fold_severity_accuracy = float(
                    accuracy_score(
                        y_severity.to_numpy()[fold_true_fault],
                        emitted[fold_true_fault],
                    )
                )
                fold_label_scores.append(fold_label_f1)
                fold_raw_scores.append(
                    raw_score(fold_label_f1, fold_severity_accuracy)
                )

            score = raw_score(label_f1, severity_accuracy)
            rows.append(
                {
                    "label_experiment_id": label_experiment,
                    "severity_experiment_id": severity_experiment,
                    "macro_f1_oof": label_f1,
                    "severity_accuracy_submission_oof": severity_accuracy,
                    "raw_score_oof": score,
                    "ml_points_oof": ml_points(score),
                    "macro_f1_fold_std": float(np.std(fold_label_scores, ddof=1)),
                    "raw_score_fold_std": float(np.std(fold_raw_scores, ddof=1)),
                }
            )

    return pd.DataFrame(rows).sort_values(
        ["raw_score_oof", "macro_f1_oof"], ascending=False
    )


def main() -> None:
    args = parse_args()
    df = pd.read_csv(args.data).reset_index(drop=True)
    validate_data(df, args.n_splits)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    splits = make_splits(
        df,
        splitter_name=args.splitter,
        n_splits=args.n_splits,
        random_state=args.random_state,
    )

    split_rows: list[dict[str, object]] = []
    for fold, (train_idx, valid_idx) in enumerate(splits, start=1):
        split_rows.append(
            {
                "fold": fold,
                "train_engines": df.iloc[train_idx]["engine_id"].nunique(),
                "valid_engines": df.iloc[valid_idx]["engine_id"].nunique(),
                "train_rows": len(train_idx),
                "valid_rows": len(valid_idx),
                "valid_engine_ids": "|".join(
                    sorted(df.iloc[valid_idx]["engine_id"].unique())
                ),
            }
        )
    pd.DataFrame(split_rows).to_csv(
        args.output_dir / "fold_assignments.csv", index=False
    )

    summaries: list[dict[str, object]] = []
    all_fold_rows: list[dict[str, object]] = []
    all_class_f1_rows: list[dict[str, object]] = []
    for feature_set in args.feature_sets:
        for model_name in args.models:
            print(f"Running {feature_set} + {model_name} ...", flush=True)
            summary, fold_rows, class_f1_rows = evaluate_experiment(
                df=df,
                splits=splits,
                model_name=model_name,
                feature_set=feature_set,
                random_state=args.random_state,
                n_jobs=args.n_jobs,
                output_dir=args.output_dir,
            )
            summaries.append(summary)
            all_fold_rows.extend(fold_rows)
            all_class_f1_rows.extend(class_f1_rows)

    summary_frame = pd.DataFrame(summaries).sort_values(
        ["raw_score_oof", "macro_f1_oof"], ascending=False
    )
    summary_frame.to_csv(args.output_dir / "benchmark_summary.csv", index=False)
    pd.DataFrame(all_fold_rows).to_csv(
        args.output_dir / "benchmark_folds.csv", index=False
    )
    pd.DataFrame(all_class_f1_rows).to_csv(
        args.output_dir / "benchmark_class_f1.csv", index=False
    )
    hybrid_frame = evaluate_hybrid_combinations(
        df=df,
        experiment_ids=summary_frame["experiment_id"].tolist(),
        output_dir=args.output_dir,
    )
    hybrid_frame.to_csv(args.output_dir / "benchmark_hybrids.csv", index=False)

    metadata = {
        "data": str(args.data),
        "rows": len(df),
        "engines": int(df["engine_id"].nunique()),
        "splitter": args.splitter,
        "n_splits": args.n_splits,
        "random_state": args.random_state,
        "feature_sets": args.feature_sets,
        "models": args.models,
        "fault_labels_used_for_severity": FAULT_LABELS,
        "note": (
            "severity_accuracy_submission treats nie_dotyczy emitted after an "
            "incorrect ok/unknown label as a severity error on a true fault"
        ),
    }
    (args.output_dir / "metadata.json").write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    columns = [
        "experiment_id",
        "macro_f1_oof",
        "severity_accuracy_submission_oof",
        "raw_score_oof",
        "macro_f1_fold_std",
        "raw_score_fold_std",
        "ml_points_oof",
    ]
    print("\nOOF ranking (complete held-out engines):")
    print(summary_frame[columns].to_string(index=False, float_format="{:.4f}".format))
    print("\nBest two-model combinations (label pipeline + severity pipeline):")
    hybrid_columns = [
        "label_experiment_id",
        "severity_experiment_id",
        "macro_f1_oof",
        "severity_accuracy_submission_oof",
        "raw_score_oof",
        "raw_score_fold_std",
        "ml_points_oof",
    ]
    print(
        hybrid_frame[hybrid_columns]
        .head(5)
        .to_string(index=False, float_format="{:.4f}".format)
    )
    print(f"\nDetailed outputs saved to: {args.output_dir.resolve()}")


if __name__ == "__main__":
    main()