"""
Evaluator Executors

Executors for evaluating trained models.
"""

from __future__ import annotations

from typing import Any, TYPE_CHECKING

import numpy as np
import pandas as pd
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
    confusion_matrix,
    classification_report,
    mean_squared_error,
    mean_absolute_error,
    r2_score,
)
from sklearn.model_selection import cross_val_score

from .base import NodeExecutor, executor_function

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


class EvaluatorExecutor(NodeExecutor):
    """Executor for evaluating ML models."""
    
    CLASSIFICATION_METRICS = {
        "accuracy": accuracy_score,
        "precision": lambda y_true, y_pred: precision_score(y_true, y_pred, average="weighted", zero_division=0),
        "recall": lambda y_true, y_pred: recall_score(y_true, y_pred, average="weighted", zero_division=0),
        "f1": lambda y_true, y_pred: f1_score(y_true, y_pred, average="weighted", zero_division=0),
    }
    
    REGRESSION_METRICS = {
        "mse": mean_squared_error,
        "rmse": lambda y_true, y_pred: np.sqrt(mean_squared_error(y_true, y_pred)),
        "mae": mean_absolute_error,
        "r2": r2_score,
    }
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> dict[str, Any]:
        """
        Evaluate a trained model.
        
        Parameters:
            - metrics: List of metrics to compute
            - cross_validation: Whether to perform CV
            - cv_folds: Number of CV folds
        
        Input should be a dict with model, X_test, y_test
        """
        params = node.data.params
        
        metrics = params.get("metrics", ["accuracy", "f1"])
        cross_validation = params.get("cross_validation", False)
        cv_folds = params.get("cv_folds", 5)
        
        # Get input
        input_value = self.get_single_input(input_data)
        
        if not isinstance(input_value, dict) or "model" not in input_value:
            raise ValueError("Evaluator requires input with 'model' key from Trainer")
        
        model = input_value["model"]
        
        # Determine if classification or regression
        is_classifier = self._is_classifier(model)
        
        results = {
            "model_type": input_value.get("model_type", type(model).__name__),
            "is_classifier": is_classifier,
            "metrics": {},
            "model_path": input_value.get("model_path"),
        }
        
        # Evaluate on test set if available
        if "X_test" in input_value and "y_test" in input_value:
            X_test = input_value["X_test"]
            y_test = input_value["y_test"]
            
            y_pred = model.predict(X_test)
            
            # Store predictions for visualization
            results["y_test"] = y_test.tolist() if hasattr(y_test, "tolist") else list(y_test)
            results["y_pred"] = y_pred.tolist() if hasattr(y_pred, "tolist") else list(y_pred)
            
            # Calculate residuals for regression
            if not is_classifier:
                residuals = np.array(y_test) - np.array(y_pred)
                results["residuals"] = residuals.tolist()
                
                # Calculate quantiles for QQ plot data
                sorted_residuals = np.sort(residuals)
                n = len(sorted_residuals)
                theoretical_quantiles = [float(np.percentile(np.random.randn(1000), (i + 0.5) / n * 100)) for i in range(n)]
                results["qq_data"] = {
                    "theoretical": theoretical_quantiles[:100],  # Limit to 100 points
                    "sample": sorted_residuals[:100].tolist() if n > 100 else sorted_residuals.tolist(),
                }
                
                # Prediction vs actual scatter data (sample if too many points)
                sample_size = min(100, len(y_test))
                indices = np.random.choice(len(y_test), sample_size, replace=False)
                y_test_arr = np.array(y_test)
                y_pred_arr = np.array(y_pred)
                results["scatter_data"] = [
                    {"actual": float(y_test_arr[i]), "predicted": float(y_pred_arr[i])}
                    for i in indices
                ]
            
            # Compute requested metrics
            metric_funcs = self.CLASSIFICATION_METRICS if is_classifier else self.REGRESSION_METRICS
            
            for metric_name in metrics:
                if metric_name in metric_funcs:
                    try:
                        score = metric_funcs[metric_name](y_test, y_pred)
                        results["metrics"][metric_name] = float(score)
                    except Exception as e:
                        results["metrics"][metric_name] = f"Error: {e}"
            
            # Additional classification metrics
            if is_classifier:
                try:
                    results["confusion_matrix"] = confusion_matrix(y_test, y_pred).tolist()
                    results["classification_report"] = classification_report(
                        y_test, y_pred, output_dict=True, zero_division=0
                    )
                except Exception:
                    pass
                
                # Try ROC AUC if binary
                try:
                    if len(np.unique(y_test)) == 2:
                        if hasattr(model, "predict_proba"):
                            y_prob = model.predict_proba(X_test)[:, 1]
                            results["metrics"]["roc_auc"] = float(roc_auc_score(y_test, y_prob))
                except Exception:
                    pass
            
            results["test_samples"] = len(X_test)
            results["predictions"] = y_pred.tolist()[:100]  # First 100 predictions
        
        # Cross-validation if requested
        if cross_validation and "X_train" in input_value and "y_train" in input_value:
            X_train = input_value["X_train"]
            y_train = input_value["y_train"]
            
            scoring = "accuracy" if is_classifier else "r2"
            cv_scores = cross_val_score(model, X_train, y_train, cv=cv_folds, scoring=scoring)
            
            results["cross_validation"] = {
                "scoring": scoring,
                "folds": cv_folds,
                "scores": cv_scores.tolist(),
                "mean": float(cv_scores.mean()),
                "std": float(cv_scores.std()),
            }
        
        # Include feature importances if available
        if "feature_importances" in input_value:
            results["feature_importances"] = input_value["feature_importances"]
        elif "coefficients" in input_value:
            results["coefficients"] = input_value["coefficients"]
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="model_evaluated",
            details={
                "metrics": results["metrics"],
                "is_classifier": is_classifier,
            },
        )
        
        return results
    
    def _is_classifier(self, model) -> bool:
        """Check if a model is a classifier."""
        from sklearn.base import is_classifier
        return is_classifier(model)


class PredictionExecutor(NodeExecutor):
    """Executor for making predictions with a trained model."""
    
    def execute(
        self,
        node: "Node",
        input_data: dict[str, Any],
        storage: "StorageManager",
        run_id: str,
    ) -> pd.DataFrame:
        """
        Make predictions using a trained model.
        
        Input should include model and new data to predict on.
        """
        # Gather inputs - need model from one input, data from another
        model = None
        data = None
        
        for key, value in input_data.items():
            if isinstance(value, dict) and "model" in value:
                model = value["model"]
            elif isinstance(value, pd.DataFrame):
                data = value
        
        if model is None:
            raise ValueError("No model found in inputs")
        if data is None:
            raise ValueError("No data found for prediction")
        
        # Make predictions
        predictions = model.predict(data)
        
        # Create result DataFrame
        result = data.copy()
        result["prediction"] = predictions
        
        # Add probabilities if available
        if hasattr(model, "predict_proba"):
            try:
                probas = model.predict_proba(data)
                for i, col in enumerate(model.classes_):
                    result[f"proba_{col}"] = probas[:, i]
            except Exception:
                pass
        
        storage.log_event(
            run_id=run_id,
            node_id=node.id,
            event="predictions_made",
            details={
                "num_predictions": len(predictions),
            },
        )
        
        return result


# Functional wrappers
evaluator_executor = executor_function(EvaluatorExecutor)
prediction_executor = executor_function(PredictionExecutor)
