"""Model training and management module."""

import numpy as np
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from typing import Dict, Any, Tuple
from config.settings import MODEL_CONFIGS, RANDOM_SEED


class ModelManager:
    """Manages model training, evaluation, and prediction."""

    def __init__(self, task_type: str = "classification"):
        """Initialize model manager.
        
        Args:
            task_type: "classification" or "regression"
        """
        self.task_type = task_type
        self.models = {}
        self.model_history = {}

    def _get_model_instance(self, model_name: str):
        """Get model instance based on name and task type."""
        if model_name == "logistic_regression":
            if self.task_type == "classification":
                return LogisticRegression(**MODEL_CONFIGS["logistic_regression"]["params"])
            else:
                return LinearRegression()
        
        elif model_name == "random_forest":
            if self.task_type == "classification":
                return RandomForestClassifier(**MODEL_CONFIGS["random_forest"]["params"])
            else:
                return RandomForestRegressor(**MODEL_CONFIGS["random_forest"]["params"])
        
        else:
            raise ValueError(f"Unknown model: {model_name}")

    def train(self, model_name: str, X_train: np.ndarray, y_train: np.ndarray) -> Dict[str, Any]:
        """Train a model.
        
        Args:
            model_name: Name of the model to train
            X_train: Training features
            y_train: Training target
        
        Returns:
            Training info dictionary
        """
        model = self._get_model_instance(model_name)
        model.fit(X_train, y_train)
        
        self.models[model_name] = model
        train_score = model.score(X_train, y_train)
        
        self.model_history[model_name] = {
            "model": model,
            "train_score": train_score,
        }
        
        return {
            "model_name": model_name,
            "train_score": train_score,
            "status": "success",
        }

    def predict(self, model_name: str, X: np.ndarray) -> np.ndarray:
        """Make predictions using a trained model.
        
        Args:
            model_name: Name of the trained model
            X: Features to predict on
        
        Returns:
            Predictions
        """
        if model_name not in self.models:
            raise ValueError(f"Model '{model_name}' not trained yet.")
        
        return self.models[model_name].predict(X)

    def predict_proba(self, model_name: str, X: np.ndarray) -> np.ndarray:
        """Get prediction probabilities (classification only).
        
        Args:
            model_name: Name of the trained model
            X: Features to predict on
        
        Returns:
            Prediction probabilities
        """
        if model_name not in self.models:
            raise ValueError(f"Model '{model_name}' not trained yet.")
        
        if not hasattr(self.models[model_name], "predict_proba"):
            raise ValueError(f"Model '{model_name}' does not support predict_proba.")
        
        return self.models[model_name].predict_proba(X)

    def get_model(self, model_name: str):
        """Get a trained model instance."""
        if model_name not in self.models:
            raise ValueError(f"Model '{model_name}' not trained yet.")
        return self.models[model_name]

    def get_models(self) -> Dict[str, Any]:
        """Get all trained models."""
        return self.models

    def list_available_models(self) -> list:
        """List available model types for this task."""
        if self.task_type == "classification":
            return list(MODEL_CONFIGS.keys())
        else:
            return list(MODEL_CONFIGS.keys())