"""Permutation negative control for the best ENGIN severity candidate.

If the grouped validation or feature construction leaked severity labels, the
model would remain accurate after training severities were randomly permuted.
A result near the majority-class baseline, far below the real model, supports
that the measured signal is genuine rather than label leakage.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score

from benchmark_grouped import FAULT_LABELS, make_splits, validate_data
from severity_benchmark import SeverityFeatures, make_estimator


DEFAULT_PERMUTATIONS = [101, 202, 303, 404, 505]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", type=Path, default=Path("val.csv"))
    parser.add_argument("--output", type=Path, default=Path("severity_outputs/negative_control.csv"))
    parser.add_argument("--cv-seed", type=int, default=42)
    parser.add_argument("--n-splits", type=int, default=5)
    parser.add_argument(
        "--permutations", nargs="+", type=int, default=DEFAULT_PERMUTATIONS
    )
    parser.add_argument("--n-jobs", type=int, default=-1)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    df = pd.read_csv(args.data).reset_index(drop=True)
    validate_data(df, args.n_splits)
    y_label = df["label"].reset_index(drop=True)
    y_severity = df["severity"].reset_index(drop=True)
    splits = make_splits(
        df,
        splitter_name="stratified-group",
        n_splits=args.n_splits,
        random_state=args.cv_seed,
    )

    rows = []
    for permutation_seed in args.permutations:
        rng = np.random.default_rng(permutation_seed)
        all_true: list[str] = []
        all_predicted: list[str] = []
        for fold, (train_idx, valid_idx) in enumerate(splits, start=1):
            transformer = SeverityFeatures("deviation").fit(df.iloc[train_idx])
            train_features = transformer.transform(df.iloc[train_idx])
            valid_features = transformer.transform(df.iloc[valid_idx])
            train_fault = y_label.iloc[train_idx].isin(FAULT_LABELS).to_numpy()
            valid_fault = y_label.iloc[valid_idx].isin(FAULT_LABELS).to_numpy()

            shuffled = y_severity.iloc[train_idx].to_numpy(dtype=object)[
                train_fault
            ].copy()
            rng.shuffle(shuffled)
            estimator = make_estimator(
                "extra_trees",
                random_state=9_000 + permutation_seed + fold,
                n_jobs=args.n_jobs,
            ).fit(train_features[train_fault], shuffled)
            predicted = estimator.predict(valid_features[valid_fault])
            all_true.extend(
                y_severity.iloc[valid_idx].to_numpy(dtype=object)[valid_fault]
            )
            all_predicted.extend(predicted)

        accuracy = float(accuracy_score(all_true, all_predicted))
        rows.append(
            {
                "cv_seed": args.cv_seed,
                "permutation_seed": permutation_seed,
                "severity_accuracy": accuracy,
            }
        )
        print(
            f"Permutation {permutation_seed}: severity accuracy = {accuracy:.4f}",
            flush=True,
        )

    result = pd.DataFrame(rows)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    result.to_csv(args.output, index=False)
    majority_baseline = float(
        y_severity[y_label.isin(FAULT_LABELS)].value_counts(normalize=True).max()
    )
    print(f"\nPermutation mean: {result['severity_accuracy'].mean():.4f}")
    print(f"Permutation max:  {result['severity_accuracy'].max():.4f}")
    print(f"Majority baseline: {majority_baseline:.4f}")
    print("Real grouped-CV model reference: approximately 0.9263")
    print(f"Saved to: {args.output.resolve()}")


if __name__ == "__main__":
    main()