"""Feature preparation shared by every training model."""

from dataclasses import dataclass, field

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

from lmpm.training.dataset import (
    CATEGORICAL_FEATURE_FIELDS,
    NUMERIC_FEATURE_FIELDS,
    TrainingDataset,
)


@dataclass(frozen=True)
class MissingValueStrategy:
    """Documented, model-independent treatment for absent feature values."""

    numeric: str = "median"
    categorical: str = "most_frequent"


DEFAULT_MISSING_VALUE_STRATEGY = MissingValueStrategy()


@dataclass(frozen=True)
class PreprocessedFeatures:
    """Dense model-ready feature matrix and the columns that produced it."""

    values: np.ndarray
    feature_names: tuple[str, ...]
    experiment_ids: tuple[str, ...]


def _feature_frame(dataset: TrainingDataset) -> pd.DataFrame:
    required_fields = NUMERIC_FEATURE_FIELDS + CATEGORICAL_FEATURE_FIELDS
    rows: list[dict[str, float | str | None]] = []
    for record in dataset.records:
        missing_fields = [
            field for field in required_fields if field not in record.features
        ]
        if missing_fields:
            experiment_id = record.metadata.get("experiment_id", "<unknown>")
            fields = ", ".join(missing_fields)
            raise ValueError(f"{experiment_id}: missing feature fields: {fields}")
        rows.append({field: record.features[field] for field in required_fields})
    return pd.DataFrame(rows, columns=required_fields)


@dataclass
class FeaturePreprocessor:
    """Impute, scale, and encode the canonical training feature columns.

    Numeric values use the median learned from the fitting dataset, then receive
    standard scaling. Categorical values use the most frequent fitting value,
    followed by one-hot encoding. Unknown categories at inference are encoded
    as all-zero columns instead of failing a recommendation request.
    """

    missing_value_strategy: MissingValueStrategy = field(
        default_factory=lambda: DEFAULT_MISSING_VALUE_STRATEGY
    )
    _transformer: ColumnTransformer | None = field(default=None, init=False)

    def _build_transformer(self) -> ColumnTransformer:
        numeric_pipeline = Pipeline(
            steps=[
                (
                    "impute",
                    SimpleImputer(strategy=self.missing_value_strategy.numeric),
                ),
                ("scale", StandardScaler()),
            ]
        )
        categorical_pipeline = Pipeline(
            steps=[
                (
                    "impute",
                    SimpleImputer(
                        strategy=self.missing_value_strategy.categorical,
                        missing_values=None,
                    ),
                ),
                (
                    "encode",
                    OneHotEncoder(handle_unknown="ignore", sparse_output=False),
                ),
            ]
        )
        return ColumnTransformer(
            transformers=[
                ("numeric", numeric_pipeline, list(NUMERIC_FEATURE_FIELDS)),
                (
                    "categorical",
                    categorical_pipeline,
                    list(CATEGORICAL_FEATURE_FIELDS),
                ),
            ],
            sparse_threshold=0,
        )

    def fit(self, dataset: TrainingDataset) -> "FeaturePreprocessor":
        feature_frame = _feature_frame(dataset)
        if feature_frame.empty:
            raise ValueError("cannot fit a feature preprocessor with an empty dataset")
        self._transformer = self._build_transformer()
        self._transformer.fit(feature_frame)
        return self

    def transform(self, dataset: TrainingDataset) -> PreprocessedFeatures:
        if self._transformer is None:
            raise RuntimeError("feature preprocessor must be fitted before transform")
        matrix = self._transformer.transform(_feature_frame(dataset))
        return PreprocessedFeatures(
            values=np.asarray(matrix, dtype=float),
            feature_names=tuple(self._transformer.get_feature_names_out()),
            experiment_ids=tuple(dataset.experiment_ids),
        )

    def fit_transform(self, dataset: TrainingDataset) -> PreprocessedFeatures:
        return self.fit(dataset).transform(dataset)