"""Runtime-only prediction contract for the frozen ENGIN model."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd

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

OOD_OK_TO_UNKNOWN_THRESHOLD_MV = 7.25
OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD = 2.5
SUBMISSION_COLUMNS = ["engine_id", "cylinder", "label", "severity"]
KEY_COLUMNS = ["engine_id", "cylinder"]
INFERENCE_COLUMNS = ["engine_id", "cylinder", "n_cylinders", *FREQ_COLS]


@dataclass
class DiagnosticModels:
    label_pipeline: object
    severity_transformer: object
    severity_estimator: object


def apply_ood_override(
    predicted_label: np.ndarray,
    deviation_features: np.ndarray,
    engine_ids: np.ndarray | pd.Series,
    threshold_mv: float = OOD_OK_TO_UNKNOWN_THRESHOLD_MV,
    threshold_ratio: float = OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Turn an implausibly anomalous ``ok`` into ``unknown``."""

    result = np.asarray(predicted_label, dtype=object).copy()
    absolute_deviation = deviation_features[
        :, len(FREQ_COLS) : 2 * len(FREQ_COLS)
    ]
    anomaly_score = absolute_deviation.mean(axis=1)
    ids = np.asarray(engine_ids, dtype=object)
    if len(ids) != len(result):
        raise ValueError("engine_ids must align with predicted labels.")
    score_frame = pd.DataFrame({"engine_id": ids, "anomaly_score": anomaly_score})
    engine_median = (
        score_frame.groupby("engine_id", sort=False)["anomaly_score"]
        .transform("median")
        .to_numpy(dtype=float)
    )
    anomaly_ratio = anomaly_score / np.maximum(engine_median, 1e-6)
    override = (
        (result == "ok")
        & (anomaly_score > threshold_mv)
        & (anomaly_ratio > threshold_ratio)
    )
    result[override] = "unknown"
    return result, override, anomaly_score, anomaly_ratio


def validate_inference_data(df: pd.DataFrame) -> None:
    required = set(INFERENCE_COLUMNS)
    missing = sorted(required.difference(df.columns))
    if missing:
        raise ValueError(f"Inference data is missing columns: {missing}")
    if df.empty:
        raise ValueError("Inference data is empty.")
    if df[["engine_id", "cylinder", "n_cylinders"]].isna().any().any():
        raise ValueError("Inference identifiers cannot contain NaN.")
    if df.duplicated(KEY_COLUMNS).any():
        raise ValueError("Duplicate engine_id + cylinder keys found in inference data.")

    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)
        & (engine_sizes["rows"] >= 2)
    )
    if not complete.all():
        bad = engine_sizes.index[~complete].tolist()
        raise ValueError(f"Incomplete or inconsistent inference engines: {bad}")


def _prepare_labeled_frame(df: pd.DataFrame, labels: np.ndarray) -> pd.DataFrame:
    result = df.copy()
    result["diagnostic_label"] = labels
    return result


def predict_test(
    models: DiagnosticModels, test: pd.DataFrame
) -> tuple[pd.DataFrame, pd.DataFrame]:
    validate_inference_data(test)

    raw_model_label = models.label_pipeline.predict(test).astype(object)
    label_features = models.label_pipeline.named_steps["features"].transform(test)
    predicted_label, ood_override, anomaly_score, anomaly_ratio = apply_ood_override(
        raw_model_label,
        label_features,
        test["engine_id"],
    )
    label_probabilities = models.label_pipeline.predict_proba(test)
    raw_model_confidence = label_probabilities.max(axis=1)
    label_confidence = raw_model_confidence.copy()
    label_confidence[ood_override] = np.nan
    ordered_probabilities = np.sort(label_probabilities, axis=1)
    label_margin = ordered_probabilities[:, -1] - ordered_probabilities[:, -2]

    test_with_labels = _prepare_labeled_frame(test, predicted_label)
    severity_features = models.severity_transformer.transform(test_with_labels)
    predicted_severity_all = models.severity_estimator.predict(
        severity_features
    ).astype(object)

    predicted_fault = np.isin(predicted_label, FAULT_LABELS)
    emitted_severity = np.full(len(test), NOT_APPLICABLE, dtype=object)
    emitted_severity[predicted_fault] = predicted_severity_all[predicted_fault]

    severity_probabilities = models.severity_estimator.predict_proba(
        severity_features
    )
    severity_confidence = np.full(len(test), np.nan, dtype=float)
    severity_confidence[predicted_fault] = severity_probabilities.max(axis=1)[
        predicted_fault
    ]

    absolute_deviation = severity_features[:, len(FREQ_COLS) : 2 * len(FREQ_COLS)]
    top_frequency_indices = np.argsort(absolute_deviation, axis=1)[:, -3:][:, ::-1]
    top_frequencies = [
        "|".join(str(int(index)) for index in row) for row in top_frequency_indices
    ]

    submission = test[KEY_COLUMNS].copy()
    submission["label"] = predicted_label
    submission["severity"] = emitted_severity

    diagnostics = submission.copy()
    diagnostics["raw_model_label"] = raw_model_label
    diagnostics["decision_source"] = np.where(
        ood_override, "ood_override", "classifier"
    )
    diagnostics["label_confidence"] = label_confidence
    diagnostics["raw_model_confidence"] = raw_model_confidence
    diagnostics["label_margin"] = label_margin
    diagnostics["severity_confidence"] = severity_confidence
    diagnostics["anomaly_score_mean_abs_mv"] = anomaly_score
    diagnostics["anomaly_ratio_to_engine_median"] = anomaly_ratio
    diagnostics["ood_absolute_threshold_mv"] = OOD_OK_TO_UNKNOWN_THRESHOLD_MV
    diagnostics["ood_ratio_threshold"] = OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD
    diagnostics["top_anomalous_frequencies_khz"] = top_frequencies
    diagnostics["missing_spectral_cells"] = (
        test[FREQ_COLS].isna().sum(axis=1).to_numpy(dtype=int)
    )
    return submission, diagnostics


def validate_submission(submission: pd.DataFrame, test: pd.DataFrame) -> None:
    if submission.columns.tolist() != SUBMISSION_COLUMNS:
        raise ValueError(
            f"Submission columns must be exactly {SUBMISSION_COLUMNS}; "
            f"received={submission.columns.tolist()}"
        )
    if len(submission) != len(test):
        raise ValueError("Submission row count does not match test.csv.")
    if submission.isna().any().any():
        raise ValueError("Submission contains NaN.")
    if submission.duplicated(KEY_COLUMNS).any():
        raise ValueError("Submission contains duplicate engine_id + cylinder keys.")
    if not submission[KEY_COLUMNS].reset_index(drop=True).equals(
        test[KEY_COLUMNS].reset_index(drop=True)
    ):
        raise ValueError("Submission keys or row order do not match test.csv.")

    invalid_labels = sorted(set(submission["label"]).difference(LABELS))
    if invalid_labels:
        raise ValueError(f"Submission contains invalid labels: {invalid_labels}")

    fault = submission["label"].isin(FAULT_LABELS)
    invalid_fault_severity = sorted(
        set(submission.loc[fault, "severity"]).difference(SEVERITIES)
    )
    if invalid_fault_severity:
        raise ValueError(
            f"Fault predictions contain invalid severity: {invalid_fault_severity}"
        )
    if not submission.loc[~fault, "severity"].eq(NOT_APPLICABLE).all():
        raise ValueError("ok/unknown predictions must use severity=nie_dotyczy.")


__all__ = [
    "DiagnosticModels",
    "INFERENCE_COLUMNS",
    "KEY_COLUMNS",
    "OOD_OK_TO_UNKNOWN_RATIO_THRESHOLD",
    "OOD_OK_TO_UNKNOWN_THRESHOLD_MV",
    "SUBMISSION_COLUMNS",
    "apply_ood_override",
    "predict_test",
    "validate_inference_data",
    "validate_submission",
]