from typing import Dict
import numpy as np

from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.metrics import (
    mean_squared_error,
    r2_score,
    accuracy_score,
    f1_score,
)

from quickml.decisions import Decision, DecisionLog


class ModelTrainer:
    def __init__(
        self,
        model_name: str,
        task: str,
        log: DecisionLog,
        random_state: int = 42,
    ):
        self.model_name = model_name
        self.task = task
        self.log = log
        self.random_state = random_state
        self.model = self._init_model()

    # -----------------------------------------------------
    # MODEL INITIALIZATION
    # -----------------------------------------------------
    def _init_model(self):
        if self.task == "regression":

            if self.model_name == "linear":
                model = LinearRegression()

            elif self.model_name == "random_forest":
                model = RandomForestRegressor(
                    random_state=self.random_state
                )

            else:
                raise ValueError(
                    f"Unsupported regression model: {self.model_name}"
                )

        elif self.task == "classification":

            if self.model_name == "logistic":
                model = LogisticRegression(max_iter=1000)

            elif self.model_name == "random_forest":
                model = RandomForestClassifier(
                    random_state=self.random_state
                )

            else:
                raise ValueError(
                    f"Unsupported classification model: {self.model_name}"
                )

        else:
            raise ValueError(f"Unknown task type: {self.task}")

        # Log model initialization
        self.log.add(
            Decision.create(
                category="modeling",
                component="ModelTrainer",
                action="initialize_model",
                reason="Model selected by user",
                parameters={
                    "model": self.model_name,
                    "task": self.task,
                },
                impact="Defines learning algorithm",
                severity="info",
            )
        )

        return model

    # -----------------------------------------------------
    # TRAINING
    # -----------------------------------------------------
    def fit(self, X_train: np.ndarray, y_train: np.ndarray) -> None:
        self.model.fit(X_train, y_train)

        self.log.add(
            Decision.create(
                category="modeling",
                component="ModelTrainer",
                action="fit_model",
                reason="Model trained on preprocessed data",
                parameters={
                    "samples": len(y_train),
                },
                impact="Model parameters learned",
                severity="info",
            )
        )

    # -----------------------------------------------------
    # EVALUATION
    # -----------------------------------------------------
    def evaluate(
        self,
        X_test: np.ndarray,
        y_test: np.ndarray,
    ) -> Dict[str, float]:

        y_pred = self.model.predict(X_test)
        metrics: Dict[str, float] = {}

        # ------------------------
        # REGRESSION METRICS
        # ------------------------
        if self.task == "regression":

            rmse = mean_squared_error(y_test, y_pred) ** 0.5
            r2 = r2_score(y_test, y_pred)

            metrics["rmse"] = rmse
            metrics["r2"] = r2

            # Log invalid R² if undefined
            if np.isnan(r2):
                self.log.add(
                    Decision.create(
                        category="evaluation",
                        component="ModelTrainer",
                        action="invalid_metric",
                        reason="R2 undefined due to insufficient test samples",
                        parameters={"metric": "r2"},
                        impact="Metric not statistically reliable",
                        severity="warning",
                    )
                )

        # ------------------------
        # CLASSIFICATION METRICS
        # ------------------------
        else:

            acc = accuracy_score(y_test, y_pred)
            f1 = f1_score(y_test, y_pred, average="weighted")

            metrics["accuracy"] = acc
            metrics["f1"] = f1

        # Log evaluation summary
        self.log.add(
            Decision.create(
                category="evaluation",
                component="ModelTrainer",
                action="evaluate_model",
                reason="Model evaluated on test data",
                parameters=metrics,
                impact="Performance metrics recorded",
                severity="info",
            )
        )

        return metrics
