"""Train the selected ENGIN models and create a validated submission.

Selected, leakage-safe architecture:

* label: signed/absolute leave-one-cylinder-out deviations, ratio deltas and
  summary statistics with balanced Logistic Regression,
* severity: leave-one-cylinder-out deviation features with Extra Trees.

The script trains only on labeled ``val.csv``. The unlabeled archive is not
used because its benefit has not been established in grouped validation.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

from benchmark_grouped import validate_data
from engin.config import FAULT_LABELS, LABELS, NOT_APPLICABLE, SEVERITIES
from engin.features import SeverityFeatures
from engin.inference import (
    KEY_COLUMNS,
    OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD,
    OOD_OK_TO_UNKNOWN_THRESHOLD_MV,
    DiagnosticModels,
    apply_ood_override,
    predict_test,
    validate_inference_data,
    validate_submission,
)
from severity_benchmark import (
    CANDIDATE_BY_ID,
    fit_candidate,
    prepare_labeled_frame,
)

__all__ = [
    "DiagnosticModels",
    "OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD",
    "OOD_OK_TO_UNKNOWN_THRESHOLD_MV",
    "apply_ood_override",
    "predict_test",
    "run_pipeline",
    "train_models",
    "validate_inference_data",
    "validate_submission",
]

LABEL_MODEL_NAME = "deviation_logistic_c10"
LABEL_FEATURE_SET = "deviation"
SEVERITY_CANDIDATE_ID = "deviation_extra_trees_mf03"


def make_final_label_pipeline(random_state: int = 42) -> Pipeline:
    """Build the multi-seed winner used for final label predictions."""

    return Pipeline(
        [
            ("features", SeverityFeatures(LABEL_FEATURE_SET)),
            ("scale", StandardScaler()),
            (
                "model",
                LogisticRegression(
                    C=10.0,
                    class_weight="balanced",
                    max_iter=5_000,
                    random_state=random_state,
                ),
            ),
        ]
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--val", type=Path, default=Path("val.csv"))
    parser.add_argument("--test", type=Path, default=Path("test.csv"))
    parser.add_argument(
        "--sample-submit", type=Path, default=Path("sample_submit.csv")
    )
    parser.add_argument("--output", type=Path, default=Path("predictions.csv"))
    parser.add_argument(
        "--diagnostics",
        type=Path,
        default=Path("prediction_diagnostics.csv"),
        help="Auxiliary model scores and anomaly data for the application.",
    )
    parser.add_argument("--random-state", type=int, default=42)
    parser.add_argument("--n-jobs", type=int, default=-1)
    return parser.parse_args()


def validate_sample_keys(sample: pd.DataFrame, test: pd.DataFrame) -> None:
    missing = sorted(set(KEY_COLUMNS).difference(sample.columns))
    if missing:
        raise ValueError(f"Sample submission is missing key columns: {missing}")
    if sample.duplicated(KEY_COLUMNS).any():
        raise ValueError("Sample submission contains duplicate keys.")
    sample_keys = set(map(tuple, sample[KEY_COLUMNS].to_numpy()))
    test_keys = set(map(tuple, test[KEY_COLUMNS].to_numpy()))
    if sample_keys != test_keys or len(sample) != len(test):
        raise ValueError("sample_submit.csv keys do not match test.csv.")


def train_models(
    val: pd.DataFrame,
    random_state: int = 42,
    n_jobs: int = -1,
) -> DiagnosticModels:
    validate_data(val, n_splits=2)
    y_label = val["label"].reset_index(drop=True)
    y_severity = val["severity"].reset_index(drop=True)

    label_pipeline = make_final_label_pipeline(random_state=random_state).fit(
        val, y_label
    )

    fault_mask = y_label.isin(FAULT_LABELS).to_numpy()
    observed_severities = set(y_severity.to_numpy()[fault_mask])
    if observed_severities != set(SEVERITIES):
        raise ValueError(
            "Training faults must contain every allowed severity; "
            f"observed={sorted(observed_severities)}"
        )

    candidate = CANDIDATE_BY_ID[SEVERITY_CANDIDATE_ID]
    val_with_labels = prepare_labeled_frame(
        val, y_label.to_numpy(dtype=object)
    )
    transformer, estimator = fit_candidate(
        candidate=candidate,
        X_train_full=val_with_labels,
        train_fault_mask=fault_mask,
        y_severity_fault=y_severity.to_numpy(dtype=object)[fault_mask],
        y_fault_label=y_label.to_numpy(dtype=object)[fault_mask],
        random_state=random_state + 1_002,
        n_jobs=n_jobs,
    )
    return DiagnosticModels(
        label_pipeline=label_pipeline,
        severity_transformer=transformer,
        severity_estimator=estimator,
    )


def run_pipeline(
    val: pd.DataFrame,
    test: pd.DataFrame,
    sample_submit: pd.DataFrame | None = None,
    random_state: int = 42,
    n_jobs: int = -1,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    if sample_submit is not None:
        validate_sample_keys(sample_submit, test)
    models = train_models(val, random_state=random_state, n_jobs=n_jobs)
    submission, diagnostics = predict_test(models, test)
    validate_submission(submission, test)
    return submission, diagnostics


def main() -> None:
    args = parse_args()
    val = pd.read_csv(args.val).reset_index(drop=True)
    test = pd.read_csv(args.test).reset_index(drop=True)
    sample_submit = pd.read_csv(args.sample_submit).reset_index(drop=True)

    submission, diagnostics = run_pipeline(
        val=val,
        test=test,
        sample_submit=sample_submit,
        random_state=args.random_state,
        n_jobs=args.n_jobs,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.diagnostics.parent.mkdir(parents=True, exist_ok=True)
    submission.to_csv(args.output, index=False)
    diagnostics.to_csv(args.diagnostics, index=False)

    print(f"Saved validated submission: {args.output.resolve()}")
    print(f"Saved application diagnostics: {args.diagnostics.resolve()}")
    print(f"Rows: {len(submission)} | engines: {submission['engine_id'].nunique()}")
    print("\nLabel distribution:")
    print(submission["label"].value_counts().reindex(LABELS, fill_value=0).to_string())
    print("\nSeverity distribution:")
    print(
        submission["severity"]
        .value_counts()
        .reindex([NOT_APPLICABLE, *SEVERITIES], fill_value=0)
        .to_string()
    )


if __name__ == "__main__":
    main()