from abc import ABC, abstractmethod
from typing import List, Any, Optional, Dict
import numpy as np
from ddi_fw.ml.model_wrapper import ModelWrapper, Result
from ddi_fw.ml.evaluation_helper import evaluate, Metrics
from ddi_fw.utils import utils

class EnsembleStrategy(ABC):
    """Base class for ensemble strategies"""
    
    @abstractmethod
    def combine_predictions(self, predictions: List[np.ndarray]) -> np.ndarray:
        """Combine individual predictions into ensemble prediction"""
        pass

    @abstractmethod
    def get_strategy_name(self) -> str:
        """Return strategy name"""
        pass


class VotingStrategy(EnsembleStrategy):
    """Majority voting strategy for classification"""
    
    def combine_predictions(self, predictions: List[np.ndarray]) -> np.ndarray:
        predictions_array = np.array(predictions)
        if predictions_array.ndim == 2:  # Probabilities
            return np.mean(predictions_array, axis=0)
        else:  # Class labels
            return np.apply_along_axis(
                lambda x: np.bincount(x.astype(int)).argmax(),
                axis=0,
                arr=predictions_array
            )
    
    def get_strategy_name(self) -> str:
        return "voting"


class AveragingStrategy(EnsembleStrategy):
    """Averaging strategy for regression or probability outputs"""
    
    def combine_predictions(self, predictions: List[np.ndarray]) -> np.ndarray:
        return np.mean(np.array(predictions), axis=0)
    
    def get_strategy_name(self) -> str:
        return "averaging"


class StackingStrategy(EnsembleStrategy):
    """Stacking strategy using a meta-learner"""
    
    def __init__(self, meta_learner: ModelWrapper):
        self.meta_learner = meta_learner
    
    def combine_predictions(self, predictions: List[np.ndarray]) -> np.ndarray:
        # Stack predictions as features for meta-learner
        stacked_features = np.column_stack(predictions)
        return self.meta_learner.predict(stacked_features)
    
    def get_strategy_name(self) -> str:
        return "stacking"


class GenericEnsembleWrapper(ModelWrapper):
    """Generic ensemble wrapper supporting multiple strategies"""
    
    def __init__(
        self,
        base_wrappers: List[ModelWrapper],
        strategy: EnsembleStrategy,
        name: str = "ensemble",
        tracking_service: Optional[Any] = None,
        date: Optional[str] = None
    ):
        self.base_wrappers = base_wrappers
        self.strategy = strategy
        self.name = name
        self.tracking_service = tracking_service
        self.date = date or utils.utc_time_as_string_simple_format()
        
        self.individual_metrics = {}
        self.ensemble_metrics = None
        self.metrics_log = []
        self.result = Result()

    def fit(self):
        """Train all base wrappers"""
        for idx, wrapper in enumerate(self.base_wrappers):
            descriptor = wrapper.descriptor
            print(f"Training {descriptor}...")
            wrapper.fit()

    def predict(self):
        """Generate ensemble predictions"""
        individual_predictions = []
        
        # Get predictions from each base wrapper
        for idx, wrapper in enumerate(self.base_wrappers):
            descriptor = wrapper.descriptor
            
            y_pred = wrapper.predict()
            individual_predictions.append(y_pred)
            
            # Evaluate individual wrapper if labels provided
            wrapper_metrics = evaluate(self.test_label, y_pred)
            self.individual_metrics[descriptor] = wrapper_metrics
            
            self.metrics_log.append({
                "type": "individual",
                "name": descriptor,
                "metrics": self._metrics_to_dict(wrapper_metrics)
            })
    
        # Combine predictions using strategy
        ensemble_pred = self.strategy.combine_predictions(individual_predictions)
        
        # Evaluate ensemble if labels provided
        self.ensemble_metrics = evaluate(self.test_label, ensemble_pred)
        self.metrics_log.append({
            "type": "ensemble",
            "name": self.name,
            "strategy": self.strategy.get_strategy_name(),
            "metrics": self._metrics_to_dict(self.ensemble_metrics)
        })
        
        return ensemble_pred

    def fit_and_evaluate(self):
        """Fit base wrappers and evaluate ensemble"""
        self.fit()
        ensemble_pred = self.predict()
        return self.metrics_log, self.ensemble_metrics, ensemble_pred

    def _metrics_to_dict(self, metrics: Metrics) -> Dict:
        """Convert Metrics object to dict"""
        if isinstance(metrics, dict):
            return metrics
        result = {}
        for attr in dir(metrics):
            if not attr.startswith('_'):
                value = getattr(metrics, attr)
                if not callable(value):
                    try:
                        result[attr] = float(value)
                    except Exception:
                        result[attr] = value
        return result

    def get_metrics_summary(self) -> Dict:
        """Return complete metrics summary"""
        return {
            "individual": self.individual_metrics,
            "ensemble": self.ensemble_metrics,
            "strategy": self.strategy.get_strategy_name(),
            "log": self.metrics_log
        }