"""ML model training utilities."""

from __future__ import annotations

from typing import Any

import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler


def scale_features(
    train: np.ndarray,
    test: np.ndarray | None = None,
    method: str = "standard",
) -> tuple[Any, np.ndarray, np.ndarray | None]:
    """Scale numeric features. Returns (scaler, X_train_scaled, X_test_scaled or None)."""
    if method == "standard":
        scaler = StandardScaler()
    else:
        scaler = StandardScaler()
    X_train = scaler.fit_transform(np.asarray(train))
    X_test = scaler.transform(test) if test is not None else None
    return scaler, X_train, X_test


def encode_labels(
    series: pd.Series | np.ndarray | list,
    method: str = "label",
) -> tuple[Any, np.ndarray]:
    """Encode labels to integers. Returns (encoder, encoded_array)."""
    arr = np.asarray(series).ravel()
    if method == "label":
        enc = LabelEncoder()
        encoded = enc.fit_transform(arr)
    else:
        enc = LabelEncoder()
        encoded = enc.fit_transform(arr)
    return enc, encoded


def build_classification_pipeline(
    vectorizer: Any | None = None,
    classifier: Any | None = None,
    scaler: Any | None = None,
) -> Any:
    """Build a sklearn Pipeline for text classification (optional vectorizer + optional scaler + classifier)."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.pipeline import Pipeline

    steps = []
    if vectorizer is not None:
        steps.append(("vec", vectorizer))
    if scaler is not None:
        steps.append(("scaler", scaler))
    steps.append(("clf", classifier if classifier is not None else LogisticRegression(max_iter=1000)))
    return Pipeline(steps)


def train_and_evaluate(
    pipeline: Any,
    X_train: np.ndarray,
    X_test: np.ndarray,
    y_train: np.ndarray,
    y_test: np.ndarray,
    task: str = "classification",
) -> dict[str, float]:
    """Fit pipeline and return metrics dict (accuracy, f1_weighted for classification; mse, rmse for regression)."""
    pipeline.fit(X_train, y_train)
    y_pred = pipeline.predict(X_test)
    from sklearn.metrics import accuracy_score, f1_score, mean_squared_error

    if task == "classification":
        return {
            "accuracy": float(accuracy_score(y_test, y_pred)),
            "f1_weighted": float(f1_score(y_test, y_pred, average="weighted", zero_division=0)),
        }
    else:
        mse = float(mean_squared_error(y_test, y_pred))
        return {"mse": mse, "rmse": mse ** 0.5}


def build_results_table(all_results: list[dict[str, Any]]) -> pd.DataFrame:
    """Build a DataFrame from list of result dicts (e.g. from train_and_evaluate). Sortable by metric."""
    return pd.DataFrame(all_results)
