"""Stable feature transformers shared by model training and runtime inference."""

from __future__ import annotations

from typing import Iterable

import numpy as np
import pandas as pd

from .config import FAULT_LABELS, FREQ_COLS

FEATURE_SETS = ("raw", "deviation", "all")


class SeverityFeatures:
    """Create engine-relative spectral features without using labels.

    Engine-relative values use a leave-one-cylinder-out median reference. Input
    to ``transform`` must therefore contain every cylinder of each engine.
    ``diagnostic_label`` is optional and may contain only a model-predicted
    label at validation or inference time.
    """

    def __init__(self, feature_set: str, include_fault_type: bool = False) -> None:
        if feature_set not in FEATURE_SETS:
            raise ValueError(f"Unknown feature_set={feature_set!r}")
        self.feature_set = feature_set
        self.include_fault_type = include_fault_type

    def fit(
        self, X: pd.DataFrame, y: Iterable[str] | None = None
    ) -> "SeverityFeatures":
        self._validate(X)
        spectra = X[FREQ_COLS].apply(pd.to_numeric, errors="coerce")
        fallback = spectra.median(axis=0).to_numpy(dtype=float)
        if np.isnan(fallback).any():
            raise ValueError("A frequency column is entirely NaN in training.")
        self.fallback_medians_ = fallback
        return self

    def transform(self, X: pd.DataFrame) -> np.ndarray:
        self._validate(X)
        if not hasattr(self, "fallback_medians_"):
            raise RuntimeError("SeverityFeatures must be fitted before transform().")
        raw = self._clean_spectra(X)
        reference = self._leave_one_out_engine_median(
            raw, X["engine_id"].to_numpy()
        )
        relative = raw - reference
        absolute_relative = np.abs(relative)
        ratio_delta = raw / np.maximum(np.abs(reference), 1e-6) - 1.0

        if self.feature_set == "raw":
            features = raw
        elif self.feature_set == "deviation":
            features = np.hstack(
                [
                    relative,
                    absolute_relative,
                    ratio_delta,
                    self._summary_features(raw, relative, ratio_delta),
                ]
            )
        else:
            features = np.hstack(
                [
                    raw,
                    relative,
                    absolute_relative,
                    ratio_delta,
                    self._summary_features(raw, relative, ratio_delta),
                ]
            )

        if self.include_fault_type:
            if "diagnostic_label" not in X.columns:
                raise ValueError(
                    "diagnostic_label is required when include_fault_type=True"
                )
            labels = X["diagnostic_label"].to_numpy(dtype=object)
            one_hot = np.column_stack(
                [labels == label for label in FAULT_LABELS]
            ).astype(float)
            features = np.hstack([features, one_hot])
        if np.isnan(features).any() or np.isinf(features).any():
            raise ValueError("Severity features contain NaN or infinity.")
        return features

    def fit_transform(
        self, X: pd.DataFrame, y: Iterable[str] | None = None
    ) -> np.ndarray:
        return self.fit(X).transform(X)

    def _clean_spectra(self, X: pd.DataFrame) -> np.ndarray:
        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))
        return spectra.to_numpy(dtype=float)

    @staticmethod
    def _leave_one_out_engine_median(
        raw: np.ndarray, groups: np.ndarray
    ) -> np.ndarray:
        reference = np.empty_like(raw)
        for engine_id in pd.unique(groups):
            positions = np.flatnonzero(groups == engine_id)
            if len(positions) < 2:
                raise ValueError(
                    f"Engine {engine_id!r} has fewer than two cylinders in transform()."
                )
            engine_values = raw[positions]
            for local_position, global_position in enumerate(positions):
                keep = np.arange(len(positions)) != local_position
                reference[global_position] = np.median(engine_values[keep], axis=0)
        return reference

    @staticmethod
    def _summary_features(
        raw: np.ndarray, relative: np.ndarray, ratio_delta: np.ndarray
    ) -> np.ndarray:
        absolute_relative = np.abs(relative)
        gradients = np.diff(raw, axis=1)
        frequency = np.arange(raw.shape[1], dtype=float)
        centered_frequency = frequency - frequency.mean()
        slope = raw @ centered_frequency / np.sum(centered_frequency**2)
        if hasattr(np, "trapezoid"):
            spectral_auc = np.trapezoid(raw, axis=1)
        else:  # pragma: no cover - compatibility with NumPy 1.x
            spectral_auc = np.trapz(raw, axis=1)

        columns = [
            raw.mean(axis=1),
            raw.std(axis=1),
            raw.min(axis=1),
            raw.max(axis=1),
            np.ptp(raw, axis=1),
            spectral_auc,
            raw.argmax(axis=1) / 20.0,
            raw.argmin(axis=1) / 20.0,
            slope,
            np.abs(gradients).mean(axis=1),
            np.abs(gradients).max(axis=1),
            relative.mean(axis=1),
            relative.std(axis=1),
            relative.min(axis=1),
            relative.max(axis=1),
            absolute_relative.mean(axis=1),
            absolute_relative.max(axis=1),
            np.sqrt(np.mean(relative**2, axis=1)),
            ratio_delta.mean(axis=1),
            ratio_delta.std(axis=1),
        ]
        for start, stop in ((0, 5), (5, 10), (10, 15), (15, 21)):
            columns.extend(
                [
                    raw[:, start:stop].mean(axis=1),
                    relative[:, start:stop].mean(axis=1),
                    absolute_relative[:, start:stop].mean(axis=1),
                ]
            )
        return np.column_stack(columns)

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


__all__ = ["FEATURE_SETS", "SeverityFeatures"]