"""Model training and prediction helpers."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

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

from meowcontext_lab.data import DEMO_MODEL_PATH, EXPECTED_LABELS, FEATURE_COLUMNS, feature_frame


@dataclass(frozen=True)
class DemoPrediction:
    """One demo prediction result."""

    label: str
    probabilities: dict[str, float]


def train_acoustic5_logistic(df: pd.DataFrame) -> Pipeline:
    """Train the identity-blind acoustic-5 logistic regression demo model."""

    pipeline = Pipeline(
        steps=[
            ("scaler", StandardScaler()),
            (
                "classifier",
                LogisticRegression(
                    max_iter=1000,
                    class_weight="balanced",
                    random_state=7,
                ),
            ),
        ]
    )
    pipeline.fit(feature_frame(df), df["context"])
    return pipeline


def save_demo_model(df: pd.DataFrame, path: Path = DEMO_MODEL_PATH) -> Path:
    """Train and save the demo model bundle."""

    path.parent.mkdir(parents=True, exist_ok=True)
    pipeline = train_acoustic5_logistic(df)
    bundle = {
        "model_name": "Logistic regression, acoustic-5",
        "pipeline": pipeline,
        "feature_columns": list(FEATURE_COLUMNS),
        "labels": list(EXPECTED_LABELS),
        "intended_use": "Predict eliciting recording context from acoustic-5 summaries.",
    }
    joblib.dump(bundle, path)
    return path


def load_demo_model(path: Path = DEMO_MODEL_PATH) -> dict[str, Any]:
    """Load the demo model bundle."""

    return joblib.load(path)


def predict_from_features(bundle: dict[str, Any], features: dict[str, float]) -> DemoPrediction:
    """Predict one label from acoustic-5 feature values."""

    columns = bundle["feature_columns"]
    row = pd.DataFrame([{column: float(features[column]) for column in columns}])
    pipeline = bundle["pipeline"]
    probabilities = pipeline.predict_proba(row)[0]
    classes = list(pipeline.classes_)
    best_idx = int(np.argmax(probabilities))
    return DemoPrediction(
        label=str(classes[best_idx]),
        probabilities={str(label): float(probabilities[idx]) for idx, label in enumerate(classes)},
    )