""" Data processing classes compatible with scikit-learn pipelines """

from typing import Any, Dict, List, Optional, Union

import numpy as np 
import pandas as pd
from collections import Counter
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import LabelEncoder



class ColumnCopyImputer(BaseEstimator, TransformerMixin):
    """
    Imputes values by copying from other columns.
    
    Parameters
    ----------
    copy_mapping : dict
        Mapping: {'target_column': 'source_column'}
    """
    
    def __init__(self, copy_mapping: Dict[str, str]):
        self.copy_mapping = copy_mapping
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "ColumnCopyImputer":
        self.is_fitted_ = True
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        X_transformed = X.copy()
        
        for target_col, source_col in self.copy_mapping.items():
            if target_col in X_transformed.columns and source_col in X_transformed.columns:
                null_mask = X_transformed[target_col].isnull()
                values_to_copy = X_transformed.loc[null_mask, source_col]
                
                if pd.api.types.is_integer_dtype(values_to_copy.dtype):
                    # Jeśli źródło to Int64, konwertuj na float
                    values_to_copy = values_to_copy.astype('float32')
                
                X_transformed.loc[null_mask, target_col] = values_to_copy
        
        return X_transformed
    

class NullImputerWithFlags(BaseEstimator, TransformerMixin):
    """
    Imputes null values and creates flag columns for selected columns.
    
    Parameters
    ----------
    columns : list
        List of columns to impute
    values : dict or any
        Values for imputation. If dict, maps columns to values.
        If single value, uses it for all columns.
    flag_columns : list, optional
        List of columns for which to create flags (whether value existed before imputation)
    flag_suffix : str, default='_isnull_flag'
        Suffix for flag column names
    strategy : str, optional
        Imputation strategy: 'mean', 'median', 'mode' (overrides values)
    """
    
    def __init__(self, 
                 columns: List[str],
                 values: Union[Dict[str, Any], Any] = None,
                 flag_columns: Optional[List[str]] = None,
                 flag_suffix: str = '_isnull_flag',
                 strategy: Optional[str] = None):
        
        self.columns = columns
        self.values = values
        self.flag_columns = flag_columns or []
        self.flag_suffix = flag_suffix
        self.strategy = strategy
        self._fitted_values: Dict[str, Any] = {}
        
        # Walidacja
        if self.values is None and self.strategy is None:
            raise ValueError("Either 'values' or 'strategy' must be provided")
        
        if self.strategy and self.strategy not in ['mean', 'median', 'mode']:
            raise ValueError("strategy must be one of: 'mean', 'median', 'mode'")
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "NullImputerWithFlags":
        """
        Learn values for imputation (if using strategy).
        """
        if self.strategy is not None:
            for column in self.columns:
                if column in X.columns:
                    if self.strategy == 'mean':
                        self._fitted_values[column] = X[column].mean()
                    elif self.strategy == 'median':
                        self._fitted_values[column] = X[column].median()
                    elif self.strategy == 'mode':
                        mode_result = X[column].mode()
                        self._fitted_values[column] = mode_result[0] if len(mode_result) > 0 else 0
        
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        """
        Impute values and create flags.
        """
        X_transformed = X.copy()
        
        # 1. NAJPIERW STWÓRZ FLAGI (przed imputacją!)
        for column in self.flag_columns:
            if column in X_transformed.columns:
                flag_name = f"{column}{self.flag_suffix}"
                X_transformed[flag_name] = X_transformed[column].isnull().astype(int)
        
        # 2. POTEM IMPUTUJ WARTOŚCI
        for column in self.columns:
            if column in X_transformed.columns:
                
                # Określ wartość do imputacji
                if self.strategy is not None and column in self._fitted_values:
                    fill_value = self._fitted_values[column]
                elif isinstance(self.values, dict) and column in self.values:
                    fill_value = self.values[column]
                elif not isinstance(self.values, dict):
                    fill_value = self.values
                else:
                    continue  # Pomiń jeśli nie ma wartości dla tej kolumny
                
                # Imputuj
                X_transformed[column] = X_transformed[column].fillna(fill_value)
        
        return X_transformed
    
class MaxMultiplierImputer(BaseEstimator, TransformerMixin):
    """
    Imputes values based on maximum * multiplier from training data.
    Optionally creates flags for columns that had values before imputation.
    
    Parameters
    ----------
    columns : list
        List of columns to impute
    multiplier : float, default=1.5
        Multiplier for maximum value (max * multiplier)
    flag_columns : list, optional
        List of columns for which to create flags (whether value existed before imputation)
        If None, no flags are created
    flag_suffix : str, default='_flag'
        Suffix for flag column names
    """
    
    def __init__(self,
                 columns: List[str],
                 multiplier: float = 1.5,
                 flag_columns: Optional[List[str]] = None,
                 flag_suffix: str = '_flag'):
        
        self.columns = columns
        self.multiplier = multiplier
        self.flag_columns = flag_columns or []
        self.flag_suffix = flag_suffix
        self._max_values: Dict[str, float] = {}
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "MaxMultiplierImputer":
        """
        Learn maximum values from training data.
        """
        for column in self.columns:
            if column in X.columns:
                existing_values = X[column].dropna()
                if len(existing_values) > 0:
                    self._max_values[column] = existing_values.max() * self.multiplier
                else:
                    self._max_values[column] = 365  # fallback
        
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        """
        Create flags and impute values as max * multiplier.
        """
        X_transformed = X.copy()
        
        # 1. NAJPIERW STWÓRZ FLAGI (przed imputacją!)
        for column in self.flag_columns:
            if column in X_transformed.columns:
                flag_name = f"{column}{self.flag_suffix}"
                X_transformed[flag_name] = (X_transformed[column].isnull()).astype(int)
        
        # 2. POTEM IMPUTUJ WARTOŚCI
        for column in self.columns:
            if column in X_transformed.columns and column in self._max_values:
                X_transformed[column] = X_transformed[column].fillna(self._max_values[column])
        
        return X_transformed

class SportsHybridEncoder(BaseEstimator, TransformerMixin):
    """
    Encodes top N sports as binary features + aggregations for the rest.
    """
    def __init__(self, sport_columns: Optional[List[str]] = None, top_n: int = 5):
        self.sport_columns = sport_columns or ['dominant_sport', 'first_ticket_sport']
        self.top_n = top_n
        self.sport_mappings_: Dict[str, List[str]] = {}
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "SportsHybridEncoder":
        
        for col in self.sport_columns:
            if col not in X.columns:
                continue
            
            all_sports = []
            for sports_array in X[col]:
                if isinstance(sports_array, np.ndarray):
                    all_sports.extend(sports_array.tolist())
            
            sport_counts = Counter(all_sports)
            top_sports = [sport for sport, _ in sport_counts.most_common(self.top_n)]
            self.sport_mappings_[col] = top_sports
            
            print(f"{col}: top {len(top_sports)} sports selected")
        
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        X_result = X.copy()
        
        for col in self.sport_columns:
            if col not in X.columns:
                continue
            
            prefix = col.replace('_sport', '')
            top_sports = self.sport_mappings_[col]
            
            # Binary tylko dla top N sportów
            for sport in top_sports:
                safe_name = (sport.lower().replace(' ', '_')
                            .replace('ł', 'l').replace('ą', 'a'))
                col_name = f'has_{prefix}_{safe_name}'
                
                X_result[col_name] = X[col].apply(
                    lambda x: 1 if isinstance(x, np.ndarray) and sport in x else 0
                )
            
            # Agregacje
            X_result[f'num_{prefix}_sports'] = X[col].apply(
                lambda x: len(x) if isinstance(x, np.ndarray) else 0
            )
            
            X_result[f'num_{prefix}_niche'] = X[col].apply(
                lambda x: sum(1 for s in x if s not in top_sports)
                          if isinstance(x, np.ndarray) else 0
            )
            
            X_result[f'pct_{prefix}_popular'] = X_result.apply(
                lambda row: (row[f'num_{prefix}_sports'] - row[f'num_{prefix}_niche']) 
                           / row[f'num_{prefix}_sports'] 
                           if row[f'num_{prefix}_sports'] > 0 else 0,
                axis=1
            )
            
            # Drop oryginał
            X_result = X_result.drop(col, axis=1)
        
        return X_result

class LabelEncoderTransformer(BaseEstimator, TransformerMixin):
    """
    Label encoding for categorical variables (for LightGBM).
    """
    def __init__(self, columns: List[str]):
        self.columns = columns
        self.label_encoders_: Dict[str, LabelEncoder] = {}
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "LabelEncoderTransformer":
        
        for col in self.columns:
            if col not in X.columns:
                continue
            
            le = LabelEncoder()
            le.fit(X[col].astype(str))
            self.label_encoders_[col] = le
            
            print(f"{col}: {len(le.classes_)} categories")
        
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        X_result = X.copy()
        
        for col in self.columns:
            if col not in X.columns or col not in self.label_encoders_:
                continue
            
            le = self.label_encoders_[col]
            col_str = X[col].astype(str)
        
            # Zamień nieznane kategorie na wartość która była w treningu
            # lub na specjalną wartość -1
            def safe_transform(val):
                if val in le.classes_:
                    return le.transform([val])[0]
                else:
                    # Nieznana wartość -> zwróć -1 lub pierwszą znaną klasę
                    return -1  # LightGBM poradzi sobie z -1 jako "unknown"
            
            X_result[col] = col_str.apply(safe_transform)

            # Oznacz jako category dla LightGBM
            X_result[col] = X_result[col].astype('category')
        
        return X_result


class DerivedFeatureCreator(BaseEstimator, TransformerMixin):
    """
    Creates new columns that are derived from existing columns.
    
    Parameters
    ----------
    derived_features : dict
        Mapping: {'new_column': ('source_column', operation)}
        Operation can be:
        - 'divide_7' - divide by 7
        - 'divide_30' - divide by 30
        - float - divide by given number
    """
    
    def __init__(self, derived_features: Dict[str, tuple]):
        self.derived_features = derived_features
    
    def fit(self, X: pd.DataFrame, y: Optional[Any] = None) -> "DerivedFeatureCreator":
        self.is_fitted_ = True
        return self
    
    def transform(self, X: pd.DataFrame, y: Optional[Any] = None) -> pd.DataFrame:
        X_transformed = X.copy()
        
        for new_col, (source_col, operation) in self.derived_features.items():
            if source_col not in X_transformed.columns:
                print(f"WARNING: Source column '{source_col}' not found, skipping '{new_col}'")
                continue
            
            if operation == 'divide_7':
                X_transformed[new_col] = X_transformed[source_col] / 7
            elif operation == 'divide_30':
                X_transformed[new_col] = X_transformed[source_col] / 30
            elif isinstance(operation, (int, float)):
                X_transformed[new_col] = X_transformed[source_col] / operation
            else:
                print(f"WARNING: Unknown operation '{operation}' for column '{new_col}'")
        
        return X_transformed
