import pandas as pd
import numpy as np
from typing import List, Union, Dict, Optional
from pathlib import Path 
import json
import logging 


from .processors import Processor


logger = logging.getLogger(__name__) 


class NumericalBinningProcessor(Processor):
    """
    A processor that performs numerical binning on a specified column using
    either equal-width or quantile strategies, outputting categorical bin labels.
    """
    
    def __init__(
        self, 
        column_name: str,
        n_bins: int = 5,
        strategy: str = 'quantile', 
        bin_labels: Optional[Union[List[str], bool]] = None, 
        output_column_name: Optional[str] = None, 
        handle_missing_value: Optional[str] = "as_is", 
        handle_out_of_range: Optional[str] = "boundary_bins" 
    ):
        super().__init__()
        self.processor_name = 'numerical_binning_processor'
        self.function_name_list = ['process', 'transform', 'fit']
        
        if not isinstance(column_name, str) or not column_name:
            raise ValueError("column_name must be a non-empty string.")
        self.column_name = column_name

        if not isinstance(n_bins, int) or n_bins <= 0:
            raise ValueError("n_bins must be a positive integer.")
        self.n_bins_requested = n_bins

        if strategy not in ['quantile', 'equal-width']:
            raise ValueError("strategy must be either 'quantile' or 'equal-width'.")
        self.strategy = strategy

        if bin_labels is not None and not isinstance(bin_labels, (list, bool)):
            raise ValueError("bin_labels must be a list of strings, boolean, or None.")
        self.bin_labels_config = bin_labels
        
        self.output_column_name = output_column_name if output_column_name else f"{self.column_name}_binned"

        if not isinstance(handle_missing_value, str):
            raise ValueError("handle_missing_value must be a string (e.g., 'as_is', 'Missing').")
        self.handle_missing_value = handle_missing_value

        if not isinstance(handle_out_of_range, str):
            raise ValueError("handle_out_of_range must be a string (e.g., 'boundary_bins', 'OutOfRange').")
        self.handle_out_of_range = handle_out_of_range

        self.bin_edges_: Optional[np.ndarray] = None
        self.actual_labels_: Optional[Union[List[str], bool]] = None
        self.n_bins_actual_: Optional[int] = None 
        self.min_fitted_value_: Optional[float] = np.nan # Initialize to nan
        self.max_fitted_value_: Optional[float] = np.nan # Initialize to nan
        self.is_fitted = False

    def fit(self, data: pd.DataFrame) -> 'NumericalBinningProcessor':
        if not isinstance(data, pd.DataFrame):
            raise TypeError("fit() requires a pandas DataFrame.")
        if self.column_name not in data.columns:
            raise ValueError(f"Column '{self.column_name}' not found in input data for fitting.")
        
        column_data = data[self.column_name].dropna()
        if column_data.empty:
            logger.warning(f"Column '{self.column_name}' has no valid data after dropping NaNs for fitting. "
                           "Processor will be fitted with a single default bin covering all values.")
            self.min_fitted_value_ = np.nan 
            self.max_fitted_value_ = np.nan
            self.bin_edges_ = np.array([-np.inf, np.inf]) 
            self.n_bins_actual_ = 1
            if isinstance(self.bin_labels_config, list) and len(self.bin_labels_config) == 1:
                self.actual_labels_ = self.bin_labels_config
            elif self.bin_labels_config is True or self.bin_labels_config is None:
                self.actual_labels_ = ["Bin_0"]
            elif self.bin_labels_config is False:
                self.actual_labels_ = False
            else: 
                logger.warning(f"Bin labels config '{self.bin_labels_config}' incompatible with single bin for empty/NaN data. Using default 'Bin_0'.")
                self.actual_labels_ = ["Bin_0"]
            self.is_fitted = True
            return self

        self.min_fitted_value_ = float(column_data.min())
        self.max_fitted_value_ = float(column_data.max())

        current_strategy = self.strategy
        n_bins_to_try = self.n_bins_requested

        if current_strategy == 'quantile':
            try:
                if column_data.nunique() < n_bins_to_try:
                    logger.warning(f"Column '{self.column_name}' has fewer unique values ({column_data.nunique()}) "
                                   f"than requested n_bins ({n_bins_to_try}). Quantile binning might result in fewer bins.")
                _, self.bin_edges_ = pd.qcut(column_data, n_bins_to_try, retbins=True, duplicates='drop')
            except ValueError as e: 
                logger.warning(f"Quantile binning failed for column '{self.column_name}' with {n_bins_to_try} bins (reason: {e}). "
                               f"Falling back to equal-width binning.")
                current_strategy = 'equal-width' 
        
        if current_strategy == 'equal-width': 
            if self.min_fitted_value_ == self.max_fitted_value_:
                logger.warning(f"Column '{self.column_name}' has a single unique value ({self.min_fitted_value_}). Creating one bin encompassing this value.")
                epsilon = max(1e-9, abs(self.min_fitted_value_ * 1e-6))  # Ensure epsilon is small but non-zero
                self.bin_edges_ = np.array([self.min_fitted_value_ - epsilon, self.min_fitted_value_ + epsilon])
            else:
                _, self.bin_edges_ = pd.cut(column_data, bins=n_bins_to_try, retbins=True, include_lowest=True, right=True)
        
        self.bin_edges_ = np.unique(self.bin_edges_)
        
        self.n_bins_actual_ = len(self.bin_edges_) - 1
        if self.n_bins_actual_ <= 0: 
            logger.warning(f"Could not create valid bins for column '{self.column_name}' (actual bins: {self.n_bins_actual_}). Defaulting to a single overarching bin.")
            self.bin_edges_ = np.array([-np.inf, np.inf])
            self.n_bins_actual_ = 1

        if self.n_bins_actual_ != self.n_bins_requested:
            logger.warning(f"Number of bins for column '{self.column_name}' was adjusted from {self.n_bins_requested} "
                           f"to {self.n_bins_actual_} due to data distribution or strategy constraints.")

        if isinstance(self.bin_labels_config, list):
            if len(self.bin_labels_config) == self.n_bins_actual_:
                self.actual_labels_ = self.bin_labels_config
            else:
                logger.warning(f"Provided bin_labels length ({len(self.bin_labels_config)}) "
                               f"does not match the actual number of bins ({self.n_bins_actual_}). Using default labels.")
                self.actual_labels_ = [f"Bin_{i}" for i in range(self.n_bins_actual_)]
        elif self.bin_labels_config is True or self.bin_labels_config is None:
            self.actual_labels_ = [f"Bin_{i}" for i in range(self.n_bins_actual_)]
        elif self.bin_labels_config is False: 
            self.actual_labels_ = False 
            
        self.is_fitted = True
        return self

    def process(self, input_value: Union[int, float, np.number]) -> Optional[str]:
        if not self.is_fitted:
            raise RuntimeError("NumericalBinningProcessor must be fitted before processing.")
        
        if pd.isna(input_value):
            return self.handle_missing_value if self.handle_missing_value != "as_is" else None
        
        val = float(input_value)

        is_out_of_fitted_range = False
        if not pd.isna(self.min_fitted_value_) and not pd.isna(self.max_fitted_value_):
            if val < self.min_fitted_value_ or val > self.max_fitted_value_:
                is_out_of_fitted_range = True
        
        if is_out_of_fitted_range and self.handle_out_of_range != "boundary_bins":
            return self.handle_out_of_range
        
        binned_series = pd.cut(
            pd.Series([val]), 
            bins=self.bin_edges_, 
            labels=self.actual_labels_, 
            include_lowest=True, 
            right=True
        )
        binned_label = binned_series[0]

        if pd.isna(binned_label): # Value didn't fall into any bin from pd.cut
            if self.handle_out_of_range == "boundary_bins" and isinstance(self.actual_labels_, list) and self.n_bins_actual_ > 0:
                if val <= self.bin_edges_[0]: # Catches values below or equal to the first edge
                    return str(self.actual_labels_[0])
                # For values > last edge, pd.cut with include_lowest=True and right=True on a single value
                # should place it in the last bin if labels are provided.
                # If it's still NaN, it implies it was truly outside even the last bin's extended range.
                # Or if only one bin [-inf, inf], it should be caught.
                # This explicit check for > last edge might be redundant if pd.cut handles it.
                # However, if labels=False, pd.cut can create an interval like (edge_n-1, edge_n],
                # and a value exactly on edge_n-1 might be an issue if not include_lowest on that specific interval.
                # For safety with "boundary_bins":
                if val >= self.bin_edges_[-1]: # Check if it's at or beyond the last edge
                     # If it's exactly on the last edge, pd.cut (right=True) includes it in the last bin.
                     # If it's greater, it should also be in the last bin conceptually for "boundary_bins".
                    return str(self.actual_labels_[-1])

            # If still NaN and not handled by boundary_bins logic above, or if not boundary_bins
            if is_out_of_fitted_range and self.handle_out_of_range != "boundary_bins": # Should have been caught earlier
                 return self.handle_out_of_range # Redundant but safe
            return None # Default for unbinnable values
            
        return str(binned_label)


    def transform(self, data: Union[pd.DataFrame, pd.Series]) -> Union[pd.DataFrame, pd.Series]:
        if not self.is_fitted:
            raise RuntimeError("NumericalBinningProcessor must be fitted before transforming.")

        output_data: Union[pd.DataFrame, pd.Series]
        series_to_bin: pd.Series

        if isinstance(data, pd.DataFrame):
            if self.column_name not in data.columns:
                raise ValueError(f"Column '{self.column_name}' not found in input DataFrame for transform.")
            series_to_bin = data[self.column_name].copy()
            output_data = data.copy()
        elif isinstance(data, pd.Series):
            series_to_bin = data.copy()
            output_data = series_to_bin 
        else:
            raise TypeError("Transform input must be a pandas DataFrame or Series.")

        original_nan_mask = series_to_bin.isna()

        binned_series_cat = pd.cut(
            series_to_bin.dropna(), # Apply cut only on non-NaN values first
            bins=self.bin_edges_,
            labels=self.actual_labels_,
            include_lowest=True,
            right=True
        )
        
        # Initialize final binned series with object dtype to allow various assignments
        final_binned_series = pd.Series(index=series_to_bin.index, dtype=object)
        final_binned_series[series_to_bin.notna()] = binned_series_cat

        # 1. Handle original NaNs
        if self.handle_missing_value != "as_is":
            final_binned_series[original_nan_mask] = self.handle_missing_value
        else:
            final_binned_series[original_nan_mask] = np.nan # Explicitly keep as NaN

        # 2. Handle out-of-range values (that were not NaN originally but might be NaN after cut, or outside fitted range)
        if not pd.isna(self.min_fitted_value_) and not pd.isna(self.max_fitted_value_):
            # Identify values that were originally numbers but are outside the fitted range
            true_out_of_range_mask = ~original_nan_mask & \
                                     ((series_to_bin < self.min_fitted_value_) | \
                                      (series_to_bin > self.max_fitted_value_))
        else: # No valid fit range, assume nothing is out of range based on fit
            true_out_of_range_mask = pd.Series([False] * len(series_to_bin), index=series_to_bin.index)

        if self.handle_out_of_range == "boundary_bins":
            if isinstance(self.actual_labels_, list) and self.n_bins_actual_ > 0:
                # Assign values below min_fitted to the first bin's label
                final_binned_series[~original_nan_mask & (series_to_bin < self.min_fitted_value_)] = self.actual_labels_[0]
                # Assign values above max_fitted to the last bin's label
                final_binned_series[~original_nan_mask & (series_to_bin > self.max_fitted_value_)] = self.actual_labels_[-1]
        else: # Custom string label for out-of-range
            final_binned_series[true_out_of_range_mask] = self.handle_out_of_range
        
        # Values that were numeric, within fitted range, but still NaN after cut (edge cases)
        # These are rare if bins cover the fitted range.
        still_nan_after_cut_mask = series_to_bin.notna() & final_binned_series.isna() & ~true_out_of_range_mask
        if self.handle_out_of_range == "boundary_bins" and isinstance(self.actual_labels_, list) and self.n_bins_actual_ > 0:
            # Attempt to place these into boundary bins as a last resort if they are near edges
            final_binned_series.loc[still_nan_after_cut_mask & (series_to_bin <= self.bin_edges_[0])] = self.actual_labels_[0]
            final_binned_series.loc[still_nan_after_cut_mask & (series_to_bin >= self.bin_edges_[-1])] = self.actual_labels_[-1]
        elif self.handle_out_of_range != "boundary_bins":
            final_binned_series.loc[still_nan_after_cut_mask] = self.handle_out_of_range


        # Determine final output type
        if self.actual_labels_ is False:  # Interval notation
            final_binned_series = pd.Series(final_binned_series, dtype=pd.CategoricalDtype())  # Ensure categorical dtype
        else:  # String labels
            final_binned_series = pd.Series(final_binned_series, dtype=pd.CategoricalDtype(categories=self.actual_labels_, ordered=True))

        if isinstance(output_data, pd.DataFrame):
            output_data[self.output_column_name] = final_binned_series
            return output_data
        else: 
            return final_binned_series


    def get_params(self) -> Dict:
        return {
            "column_name": self.column_name,
            "n_bins_requested": self.n_bins_requested,
            "n_bins_actual": self.n_bins_actual_,
            "strategy": self.strategy,
            "bin_labels_config": self.bin_labels_config,
            "output_column_name": self.output_column_name,
            "handle_missing_value": self.handle_missing_value,
            "handle_out_of_range": self.handle_out_of_range,
            "bin_edges": self.bin_edges_.tolist() if self.bin_edges_ is not None else None,
            "actual_labels": self.actual_labels_ if isinstance(self.actual_labels_, list) else str(self.actual_labels_),
            "min_fitted_value": self.min_fitted_value_,
            "max_fitted_value": self.max_fitted_value_
        }

    def save_params(self, output_dir: Union[str, Path]) -> None:
        if not self.is_fitted:
            raise RuntimeError("Processor must be fitted before saving parameters.")
        
        output_dir_path = Path(output_dir)
        output_dir_path.mkdir(parents=True, exist_ok=True)
        params_to_save = self.get_params()
        filepath = output_dir_path / f"{self.processor_name}_{self.column_name}_params.json"
        with open(filepath, 'w') as f:
            json.dump(params_to_save, f, indent=4)
        logger.info(f"Parameters for '{self.column_name}' saved to {filepath}")

    @classmethod
    def load_params(cls, source: Union[str, Path, Dict]) -> 'NumericalBinningProcessor':
        params: Dict
        if isinstance(source, dict):
            params = source
            logger.info(f"Parameters loaded directly from dictionary for column '{params.get('column_name', 'Unknown')}'.")
        elif isinstance(source, (str, Path)):
            filepath_path = Path(source)
            if not filepath_path.exists():
                raise FileNotFoundError(f"Parameter file not found: {filepath_path}")
            with open(filepath_path, 'r') as f:
                params = json.load(f)
            logger.info(f"Parameters loaded from file {filepath_path}")
        else:
            raise TypeError("source must be a filepath (str or Path) or a dictionary.")
        
        required_keys = ["column_name", "n_bins_requested", "strategy", "bin_edges", "actual_labels"]
        if not all(key in params for key in required_keys):
            missing = [key for key in required_keys if key not in params]
            raise ValueError(f"Loaded parameters are missing required keys: {missing}")

        processor = cls(
            column_name=params["column_name"],
            n_bins=params["n_bins_requested"],
            strategy=params["strategy"],
            bin_labels=params.get("bin_labels_config"),
            output_column_name=params.get("output_column_name"),
            handle_missing_value=params.get("handle_missing_value", "as_is"),
            handle_out_of_range=params.get("handle_out_of_range", "boundary_bins")
        )
        
        processor.bin_edges_ = np.array(params["bin_edges"]) if params["bin_edges"] is not None else None
        
        loaded_actual_labels = params["actual_labels"]
        if isinstance(loaded_actual_labels, str) and loaded_actual_labels.lower() == 'false':
            processor.actual_labels_ = False
        elif loaded_actual_labels is None and isinstance(params.get("bin_labels_config"), bool) and not params.get("bin_labels_config"):
            processor.actual_labels_ = False
        else:
            processor.actual_labels_ = loaded_actual_labels

        processor.n_bins_actual_ = params.get("n_bins_actual")
        processor.min_fitted_value_ = params.get("min_fitted_value")
        processor.max_fitted_value_ = params.get("max_fitted_value")
        
        if processor.bin_edges_ is not None and processor.n_bins_actual_ is not None : 
            processor.is_fitted = True
        
        return processor