"""
Optimus Research Center - Laboratory 1: Predictive Algorithms
Executable implementations of Ridge, Gradient Boosting, Exponential Smoothing, and KNN.
"""

import numpy as np
from typing import Dict, Any, Tuple
from sklearn.linear_model import Ridge
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score


def _safe_r2(y_true: np.ndarray, y_pred: np.ndarray) -> Tuple[Any, str]:
    """
    Computes R^2 only when mathematically well-defined (n_samples >= 2 and non-zero target variance).
    Prevents UndefinedMetricWarning on single-sample test evaluations without fabricating artificial scores.
    """
    if len(y_true) >= 2 and float(np.var(y_true)) > 0.0:
        return round(float(r2_score(y_true, y_pred)), 6), "DEFINED"
    return None, "UNDEFINED_LESS_THAN_2_SAMPLES"


class RidgePredictor:
    def __init__(self, alpha: float = 1.0, random_seed: int = 42):
        self.alpha = alpha
        self.seed = random_seed
        self.model = Ridge(alpha=alpha, random_state=random_seed)

    def train_and_eval(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_test: np.ndarray,
        y_test: np.ndarray
    ) -> Dict[str, Any]:
        self.model.fit(X_train, y_train)
        preds = self.model.predict(X_test)
        mse = float(mean_squared_error(y_test, preds))
        mae = float(mean_absolute_error(y_test, preds))
        r2_val, r2_status = _safe_r2(y_test, preds)
        return {
            "algorithm": "RidgeRegressionPredictor",
            "mse": round(mse, 6),
            "mae": round(mae, 6),
            "r2": r2_val,
            "r2_status": r2_status,
            "coefficients_count": len(self.model.coef_)
        }


class GradientBoostingPredictor:
    def __init__(self, n_estimators: int = 50, learning_rate: float = 0.1, random_seed: int = 42):
        self.model = GradientBoostingRegressor(
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            random_state=random_seed
        )

    def train_and_eval(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_test: np.ndarray,
        y_test: np.ndarray
    ) -> Dict[str, Any]:
        self.model.fit(X_train, y_train)
        preds = self.model.predict(X_test)
        mse = float(mean_squared_error(y_test, preds))
        mae = float(mean_absolute_error(y_test, preds))
        r2_val, r2_status = _safe_r2(y_test, preds)
        return {
            "algorithm": "GradientBoostingPredictor",
            "mse": round(mse, 6),
            "mae": round(mae, 6),
            "r2": r2_val,
            "r2_status": r2_status
        }


class KNNPredictor:
    def __init__(self, n_neighbors: int = 5):
        self.model = KNeighborsRegressor(n_neighbors=n_neighbors)

    def train_and_eval(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_test: np.ndarray,
        y_test: np.ndarray
    ) -> Dict[str, Any]:
        self.model.fit(X_train, y_train)
        preds = self.model.predict(X_test)
        mse = float(mean_squared_error(y_test, preds))
        mae = float(mean_absolute_error(y_test, preds))
        r2_val, r2_status = _safe_r2(y_test, preds)
        return {
            "algorithm": "KNNRegressorPredictor",
            "mse": round(mse, 6),
            "mae": round(mae, 6),
            "r2": r2_val,
            "r2_status": r2_status
        }


class ExponentialSmoothingPredictor:
    """Holt-Winters double exponential smoothing for 1D time series."""

    def __init__(self, alpha: float = 0.3, beta: float = 0.1):
        self.alpha = alpha
        self.beta = beta

    def forecast_and_eval(self, series: np.ndarray, test_horizon: int = 10) -> Dict[str, Any]:
        if len(series) < test_horizon + 5:
            raise ValueError("Time series too short for evaluation")

        train = series[:-test_horizon]
        actual = series[-test_horizon:]

        level = train[0]
        trend = train[1] - train[0]

        for val in train[1:]:
            last_level = level
            level = self.alpha * val + (1 - self.alpha) * (level + trend)
            trend = self.beta * (level - last_level) + (1 - self.beta) * trend

        # Forecast horizon steps
        forecast = [level + (i + 1) * trend for i in range(test_horizon)]
        forecast = np.array(forecast)

        mse = float(np.mean((actual - forecast) ** 2))
        mae = float(np.mean(np.abs(actual - forecast)))

        return {
            "algorithm": "ExponentialSmoothingPredictor",
            "test_horizon": test_horizon,
            "mse": round(mse, 6),
            "mae": round(mae, 6),
            "forecast": forecast.tolist(),
            "actual": actual.tolist()
        }