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

logger = logging.getLogger(__name__)

class InputHandler:
    """
    Universal input handler that can prepare data for any ML framework.
    Handles data conversion and validation for different model types.
    """
    
    def __init__(self):
        self.input_type = None
        self.framework = None
        
    def detect_framework(self, model_func: Any) -> str:
        """
        Detect the ML framework of the model.
        
        Returns:
            Framework name: 'pytorch', 'tensorflow', 'sklearn', 'custom'
        """
        # Check for PyTorch
        if hasattr(model_func, 'parameters') or hasattr(model_func, 'forward'):
            return 'pytorch'
        
        # Check for TensorFlow/Keras
        if hasattr(model_func, 'layers') or hasattr(model_func, 'get_weights'):
            return 'tensorflow'
        
        # Check for scikit-learn
        if hasattr(model_func, 'predict') and hasattr(model_func, 'fit'):
            return 'sklearn'
        
        # Check if it's just a callable
        if callable(model_func):
            return 'custom'
        
        return 'unknown'
    
    def prepare_input(
        self, 
        data: Any, 
        model_func: Any,
        framework: Optional[str] = None
    ) -> Any:
        """
        Prepare input data for the specific model framework.
        
        Args:
            data: Input data in any format
            model_func: The model to prepare data for
            framework: Optional framework hint
            
        Returns:
            Data in the format expected by the model
        """
        if framework is None:
            framework = self.detect_framework(model_func)
        
        # Normalize data to DataFrame first if possible
        normalized_data = self._normalize_to_dataframe(data)
        
        # Convert based on framework
        if framework == 'pytorch':
            return self._prepare_pytorch_input(normalized_data)
        elif framework == 'tensorflow':
            return self._prepare_tensorflow_input(normalized_data)
        elif framework == 'sklearn':
            return self._prepare_sklearn_input(normalized_data)
        else:
            # Custom or unknown - try to keep as DataFrame or numpy
            if isinstance(normalized_data, pd.DataFrame):
                return normalized_data
            return np.array(normalized_data)
    
    def _normalize_to_dataframe(self, data: Any) -> pd.DataFrame:
        """
        Normalize various input formats to pandas DataFrame.
        """
        if isinstance(data, pd.DataFrame):
            return data
        
        if isinstance(data, pd.Series):
            return data.to_frame()
        
        if isinstance(data, np.ndarray):
            if len(data.shape) == 1:
                return pd.DataFrame({0: data})
            return pd.DataFrame(data)
        
        if isinstance(data, list):
            return pd.DataFrame(data)
        
        if isinstance(data, dict):
            return pd.DataFrame([data])
        
        # Try to convert anything else
        try:
            return pd.DataFrame(data)
        except Exception as e:
            logger.warning(f"Could not convert input to DataFrame: {e}")
            return pd.DataFrame()
    
    def _prepare_pytorch_input(self, data: pd.DataFrame) -> Any:
        """
        Prepare input for PyTorch models.
        """
        try:
            import torch
            
            # Convert DataFrame to tensor
            if isinstance(data, pd.DataFrame):
                tensor_data = torch.FloatTensor(data.values)
            else:
                tensor_data = torch.FloatTensor(np.array(data))
            
            logger.debug(f"Prepared PyTorch input with shape: {tensor_data.shape}")
            return tensor_data
            
        except ImportError:
            logger.warning("PyTorch not available, returning numpy array")
            return data.values if isinstance(data, pd.DataFrame) else np.array(data)
    
    def _prepare_tensorflow_input(self, data: pd.DataFrame) -> Any:
        """
        Prepare input for TensorFlow/Keras models.
        """
        try:
            import tensorflow as tf
            
            # TensorFlow models typically accept numpy arrays or DataFrames
            if isinstance(data, pd.DataFrame):
                return data.values.astype(np.float32)
            return np.array(data).astype(np.float32)
            
        except ImportError:
            logger.warning("TensorFlow not available, returning numpy array")
            return data.values if isinstance(data, pd.DataFrame) else np.array(data)
    
    def _prepare_sklearn_input(self, data: pd.DataFrame) -> Any:
        """
        Prepare input for scikit-learn models.
        """
        # scikit-learn expects numpy arrays or DataFrames
        if isinstance(data, pd.DataFrame):
            return data
        return np.array(data)
    
    def validate_input(self, data: Any) -> bool:
        """
        Validate that input data is in a usable format.
        
        Returns:
            True if valid, False otherwise
        """
        if data is None:
            logger.error("Input data is None")
            return False
        
        # Check if we can convert to array
        try:
            if isinstance(data, (pd.DataFrame, pd.Series)):
                if len(data) == 0:
                    logger.error("Input data is empty")
                    return False
                return True
            
            arr = np.array(data)
            if arr.size == 0:
                logger.error("Input data is empty")
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"Cannot validate input data: {e}")
            return False
    
    def handle_missing_values(
        self, 
        data: pd.DataFrame, 
        strategy: str = 'drop'
    ) -> pd.DataFrame:
        """
        Handle missing values in input data.
        
        Args:
            data: Input DataFrame
            strategy: 'drop', 'mean', 'median', 'zero', 'forward_fill'
            
        Returns:
            DataFrame with missing values handled
        """
        if not isinstance(data, pd.DataFrame):
            data = self._normalize_to_dataframe(data)
        
        if data.isnull().sum().sum() == 0:
            return data
        
        logger.info(f"Handling missing values with strategy: {strategy}")
        
        if strategy == 'drop':
            return data.dropna()
        elif strategy == 'mean':
            return data.fillna(data.mean())
        elif strategy == 'median':
            return data.fillna(data.median())
        elif strategy == 'zero':
            return data.fillna(0)
        elif strategy == 'forward_fill':
            return data.fillna(method='ffill').fillna(method='bfill')
        else:
            logger.warning(f"Unknown strategy '{strategy}', using 'drop'")
            return data.dropna()
