"""
Trainer Executors

Executors for training ML models.
"""

from __future__ import annotations

from typing import Any, TYPE_CHECKING

import pandas as pd
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import (
    RandomForestClassifier,
    RandomForestRegressor,
    GradientBoostingClassifier,
    GradientBoostingRegressor,
)
from sklearn.svm import SVC, SVR
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

from .base import NodeExecutor, executor_function

if TYPE_CHECKING:
    from ..schemas import Node
    from ..storage import StorageManager


class TrainerExecutor(NodeExecutor):
    """Executor for training ML models."""
    
    MODEL_MAP = {
        # Classification
        "logistic_regression": LogisticRegression,
        "random_forest_classifier": RandomForestClassifier,
        "svm_classifier": SVC,
        "knn_classifier": KNeighborsClassifier,
        "decision_tree_classifier": DecisionTreeClassifier,
        "gradient_boosting_classifier": GradientBoostingClassifier,
        # Regression
        "linear_regression": LinearRegression,
        "random_forest_regressor": RandomForestRegressor,
        "svr": SVR,
        "knn_regressor": KNeighborsRegressor,
        "decision_tree_regressor": DecisionTreeRegressor,
        "gradient_boosting_regressor": GradientBoostingRegressor,
    }
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> dict[str, Any]:
        """
        Train an ML model.
        
        Parameters:
            - model_type: Type of model to train
            - target_column: Column to predict (if not already split)
            - hyperparameters: Model hyperparameters
        
        Input can be:
            - DataFrame with target_column specified
            - Dict with X_train, y_train, X_test, y_test
        """
        params = node.data.params
        
        model_type = params.get("model_type", "random_forest_classifier")
        hyperparameters = params.get("hyperparameters", {})
        target_column = params.get("target_column")
        
        # Get model class
        model_class = self.MODEL_MAP.get(model_type)
        if model_class is None:
            raise ValueError(f"Unknown model type: {model_type}")
        
        # Process input data
        input_value = self.get_single_input(input_data)
        
        if isinstance(input_value, dict) and "X_train" in input_value:
            # Already split
            X_train = input_value["X_train"]
            y_train = input_value["y_train"]
            X_test = input_value.get("X_test")
            y_test = input_value.get("y_test")
        elif isinstance(input_value, pd.DataFrame):
            # Need to split
            if not target_column:
                raise ValueError("target_column required when input is DataFrame")
            
            if target_column not in input_value.columns:
                raise ValueError(f"Column '{target_column}' not found in DataFrame")
            
            X_train = input_value.drop(columns=[target_column])
            y_train = input_value[target_column]
            X_test = None
            y_test = None
        else:
            raise ValueError(f"Unexpected input type: {type(input_value)}")
        
        # Ensure hyperparameters are valid for the model
        valid_params = self._filter_valid_params(model_class, hyperparameters)
        
        # Create and train model
        model = model_class(**valid_params)
        model.fit(X_train, y_train)
        
        # Save model
        model_path = storage.save_model(
            run_id=run_id,
            model_name=f"{node.id}_{model_type}",
            model=model,
        )
        
        # Create result
        result = {
            "model": model,
            "model_type": model_type,
            "model_path": model_path,
            "X_train": X_train,
            "y_train": y_train,
            "feature_names": list(X_train.columns),
        }
        
        if X_test is not None:
            result["X_test"] = X_test
            result["y_test"] = y_test
        
        # Get feature importances if available
        if hasattr(model, "feature_importances_"):
            result["feature_importances"] = dict(
                zip(X_train.columns, model.feature_importances_)
            )
        elif hasattr(model, "coef_"):
            coef = model.coef_
            if len(coef.shape) == 1:
                result["coefficients"] = dict(zip(X_train.columns, coef))
            else:
                result["coefficients"] = dict(zip(X_train.columns, coef[0]))
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="model_trained",
            details={
                "model_type": model_type,
                "hyperparameters": valid_params,
                "training_samples": len(X_train),
                "features": len(X_train.columns),
                "model_path": model_path,
            },
        )
        
        return result
    
    def _filter_valid_params(self, model_class, params: dict) -> dict:
        """Filter parameters to only those valid for the model class."""
        import inspect
        
        sig = inspect.signature(model_class)
        valid_param_names = set(sig.parameters.keys())
        
        return {k: v for k, v in params.items() if k in valid_param_names}


# Functional wrapper
trainer_executor = executor_function(TrainerExecutor)
