import logging
from typing import Any, List, Dict, Union, Tuple, Optional
import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


class OutputHandler:
    """
    Universal output handler that can process predictions from any ML framework
    and convert them to a standardized format for metrics calculation.
    """
    
    def __init__(self):
        self.output_type = None
        self.output_shape = None
        self.is_classification = False
        self.is_regression = False
        self.is_multioutput = False
        self.num_classes = None
        
    def infer_output_type(self, predictions: Any) -> Dict[str, Any]:
        """
        Automatically infer the type and structure of model predictions.
        
        Args:
            predictions: Raw model output in any format
            
        Returns:
            Dictionary with output metadata
        """
        metadata = {
            "type": "unknown",
            "shape": None,
            "is_classification": False,
            "is_regression": False,
            "is_multioutput": False,
            "num_classes": None,
            "framework": None,
        }
        
        # Check for None or empty
        if predictions is None:
            metadata["type"] = "none"
            return metadata
            
        # Detect framework-specific types
        predictions_normalized = self._normalize_framework_output(predictions)
        
        # Convert to numpy for analysis
        if isinstance(predictions_normalized, pd.Series):
            pred_array = predictions_normalized.values
        elif isinstance(predictions_normalized, pd.DataFrame):
            pred_array = predictions_normalized.values
        elif isinstance(predictions_normalized, list):
            pred_array = np.array(predictions_normalized)
        elif isinstance(predictions_normalized, (int, float)):
            pred_array = np.array([predictions_normalized])
        else:
            pred_array = np.array(predictions_normalized)
        
        metadata["shape"] = pred_array.shape
        
        # Determine prediction type
        if len(pred_array.shape) == 1:
            # 1D output
            metadata["is_multioutput"] = False
            if self._is_binary_or_categorical(pred_array):
                metadata["is_classification"] = True
                metadata["type"] = "classification"
                metadata["num_classes"] = len(np.unique(pred_array))
            else:
                metadata["is_regression"] = True
                metadata["type"] = "regression"
                
        elif len(pred_array.shape) == 2:
            # 2D output - could be multi-class probabilities or multi-output
            if pred_array.shape[1] > 1:
                # Check if rows sum to ~1 (probabilities)
                row_sums = np.sum(pred_array, axis=1)
                if np.allclose(row_sums, 1.0, atol=0.01):
                    metadata["is_classification"] = True
                    metadata["type"] = "classification_proba"
                    metadata["num_classes"] = pred_array.shape[1]
                else:
                    metadata["is_multioutput"] = True
                    metadata["type"] = "multioutput"
            else:
                # Single column - treat as regression
                metadata["is_regression"] = True
                metadata["type"] = "regression"
                
        elif len(pred_array.shape) > 2:
            # High-dimensional output (images, sequences, etc.)
            metadata["type"] = "high_dimensional"
            metadata["is_multioutput"] = True
            
        return metadata
    
    def _normalize_framework_output(self, predictions: Any) -> Any:
        """
        Normalize output from different ML frameworks to a common format.
        
        Supports:
        - TensorFlow/Keras tensors
        - PyTorch tensors
        - scikit-learn outputs
        - NumPy arrays
        - Python lists
        - Pandas Series/DataFrames
        """
        # Check for TensorFlow/Keras
        if hasattr(predictions, 'numpy'):
            try:
                # TensorFlow tensor or Keras output
                logger.debug("Detected TensorFlow/Keras output")
                return predictions.numpy()
            except Exception as e:
                logger.warning(f"Failed to convert TensorFlow tensor: {e}")
        
        # Check for PyTorch
        if hasattr(predictions, 'cpu') and hasattr(predictions, 'detach'):
            try:
                logger.debug("Detected PyTorch output")
                return predictions.cpu().detach().numpy()
            except Exception as e:
                logger.warning(f"Failed to convert PyTorch tensor: {e}")
        
        # Check for sparse matrices (scikit-learn)
        if hasattr(predictions, 'toarray'):
            try:
                logger.debug("Detected sparse matrix output")
                return predictions.toarray()
            except Exception as e:
                logger.warning(f"Failed to convert sparse matrix: {e}")
        
        # Already numpy or compatible
        return predictions
    
    def _is_binary_or_categorical(self, array: np.ndarray) -> bool:
        """
        Check if array contains categorical/binary values.
        """
        unique_vals = np.unique(array)
        
        # If only a few unique values and they're integers or binary
        if len(unique_vals) <= 20:
            # Check if values are integers or binary (0/1)
            if np.all(array == array.astype(int)):
                return True
            if np.all((array == 0) | (array == 1)):
                return True
        
        return False
    
    def standardize_predictions(
        self, 
        predictions: Any, 
        task_type: Optional[str] = None
    ) -> List[float]:
        """
        Convert any prediction format to a standardized list of floats.
        
        Args:
            predictions: Raw model output
            task_type: Optional hint about task type ('classification', 'regression', etc.)
            
        Returns:
            List of float values suitable for metrics calculation
        """
        if predictions is None:
            return []
        
        # Normalize framework-specific outputs
        predictions_normalized = self._normalize_framework_output(predictions)
        
        # Infer output type if not provided
        if task_type is None:
            metadata = self.infer_output_type(predictions_normalized)
            task_type = metadata["type"]
        
        # Convert to numpy array
        if isinstance(predictions_normalized, pd.Series):
            pred_array = predictions_normalized.values
        elif isinstance(predictions_normalized, pd.DataFrame):
            pred_array = predictions_normalized.values
        elif isinstance(predictions_normalized, list):
            pred_array = np.array(predictions_normalized)
        elif isinstance(predictions_normalized, (int, float)):
            pred_array = np.array([predictions_normalized])
        else:
            pred_array = np.array(predictions_normalized)
        
        # Handle different output types
        if task_type == "classification_proba":
            # Multi-class probabilities - convert to class predictions or use max probability
            if len(pred_array.shape) == 2:
                # Use the maximum probability as the score
                standardized = np.max(pred_array, axis=1).tolist()
            else:
                standardized = pred_array.flatten().tolist()
                
        elif task_type == "classification":
            # Class labels - already numeric
            standardized = pred_array.flatten().tolist()
            
        elif task_type == "regression":
            # Regression outputs
            standardized = pred_array.flatten().tolist()
            
        elif task_type == "multioutput":
            # Multi-output predictions - use mean across outputs
            if len(pred_array.shape) == 2:
                standardized = np.mean(pred_array, axis=1).tolist()
            else:
                standardized = pred_array.flatten().tolist()
                
        elif task_type == "high_dimensional":
            # High-dimensional outputs (images, etc.) - use mean
            standardized = np.mean(pred_array.reshape(pred_array.shape[0], -1), axis=1).tolist()
            
        else:
            # Unknown type - flatten and hope for the best
            logger.warning(f"Unknown prediction type: {task_type}, flattening output")
            standardized = pred_array.flatten().tolist()
        
        # Convert to float and handle NaN/inf
        result = []
        for val in standardized:
            try:
                float_val = float(val)
                if np.isnan(float_val) or np.isinf(float_val):
                    result.append(None)
                else:
                    result.append(float_val)
            except (ValueError, TypeError):
                result.append(None)
        
        return result
    
    def handle_bounding_boxes(
        self, 
        predictions: Any,
        format_type: str = "xyxy"
    ) -> List[float]:
        """
        Handle object detection outputs with bounding boxes.
        
        Args:
            predictions: Bounding box predictions
            format_type: Format of boxes ('xyxy', 'xywh', 'cxcywh')
            
        Returns:
            List of confidence scores or box areas for drift detection
        """
        predictions_normalized = self._normalize_framework_output(predictions)
        pred_array = np.array(predictions_normalized)
        
        # If predictions include confidence scores, extract them
        if isinstance(predictions_normalized, dict):
            if 'scores' in predictions_normalized:
                return self.standardize_predictions(predictions_normalized['scores'])
            elif 'confidences' in predictions_normalized:
                return self.standardize_predictions(predictions_normalized['confidences'])
        
        # If array format, assume last column is confidence or compute box areas
        if len(pred_array.shape) == 2 and pred_array.shape[1] >= 4:
            # Extract bounding box coordinates
            if pred_array.shape[1] > 4:
                # Assume last column is confidence
                return pred_array[:, -1].tolist()
            else:
                # Compute box areas as a metric
                if format_type == "xyxy":
                    areas = (pred_array[:, 2] - pred_array[:, 0]) * (pred_array[:, 3] - pred_array[:, 1])
                elif format_type in ["xywh", "cxcywh"]:
                    areas = pred_array[:, 2] * pred_array[:, 3]
                else:
                    areas = np.ones(len(pred_array))
                return areas.tolist()
        
        # Fallback to standard handling
        return self.standardize_predictions(predictions_normalized)
    
    def handle_segmentation(self, predictions: Any) -> List[float]:
        """
        Handle semantic/instance segmentation outputs.
        
        Args:
            predictions: Segmentation masks
            
        Returns:
            List of mean prediction values per sample
        """
        predictions_normalized = self._normalize_framework_output(predictions)
        pred_array = np.array(predictions_normalized)
        
        # Compute mean confidence per mask
        if len(pred_array.shape) > 2:
            # Shape: (batch, height, width, [channels])
            batch_size = pred_array.shape[0]
            means = []
            for i in range(batch_size):
                means.append(float(np.mean(pred_array[i])))
            return means
        
        return self.standardize_predictions(predictions_normalized)


def get_robust_predictions(
    model_func: Any,
    input_data: Any,
    output_handler: Optional[OutputHandler] = None
) -> List[float]:
    """
    Robustly get predictions from any model type with comprehensive error handling.
    
    Args:
        model_func: Model or prediction function
        input_data: Input data (DataFrame, array, tensor, etc.)
        output_handler: Optional OutputHandler instance
        
    Returns:
        List of standardized float predictions
    """
    if output_handler is None:
        output_handler = OutputHandler()
    
    try:
        # Handle different model types
        raw_predictions = None
        
        # Check if it's a callable
        if callable(model_func):
            raw_predictions = model_func(input_data)
        # Check for scikit-learn style models
        elif hasattr(model_func, 'predict'):
            raw_predictions = model_func.predict(input_data)
        # Check for PyTorch models
        elif hasattr(model_func, 'forward'):
            import torch
            if not isinstance(input_data, torch.Tensor):
                # Convert to tensor if needed
                if isinstance(input_data, pd.DataFrame):
                    input_tensor = torch.FloatTensor(input_data.values)
                else:
                    input_tensor = torch.FloatTensor(np.array(input_data))
            else:
                input_tensor = input_data
            
            model_func.eval()
            with torch.no_grad():
                raw_predictions = model_func.forward(input_tensor)
        # Check for TensorFlow/Keras models
        elif hasattr(model_func, '__call__'):
            raw_predictions = model_func(input_data)
        else:
            raise ValueError(f"Unsupported model type: {type(model_func)}")
        
        # Standardize the output
        standardized = output_handler.standardize_predictions(raw_predictions)
        
        # Ensure we have the right number of predictions
        if isinstance(input_data, pd.DataFrame):
            expected_len = len(input_data)
        elif hasattr(input_data, '__len__'):
            expected_len = len(input_data)
        else:
            expected_len = 1
        
        if len(standardized) != expected_len:
            logger.warning(
                f"Prediction count mismatch: got {len(standardized)}, expected {expected_len}"
            )
            # Pad or truncate as needed
            if len(standardized) < expected_len:
                standardized.extend([None] * (expected_len - len(standardized)))
            else:
                standardized = standardized[:expected_len]
        
        return standardized
        
    except Exception as e:
        logger.error(f"Error getting predictions: {e}")
        # Return None values for failed predictions
        if isinstance(input_data, pd.DataFrame):
            return [None] * len(input_data)
        elif hasattr(input_data, '__len__'):
            return [None] * len(input_data)
        else:
            return [None]
