
import pandas as pd
import numpy as np
from sklearn.tree import DecisionTreeClassifier
from .config import Config

class BinningProcess:
    def __init__(self, method='chi', max_bins=10, min_bin_pct=0.05, special_values=None):
        self.method = method
        self.max_bins = max_bins
        self.min_bin_pct = min_bin_pct
        self.special_values = special_values if special_values else Config.SPECIAL_VALUES
        self.rules = {}  # Stores binning rules: {col: {'cutoffs': [], 'woe': {}}}
        self.iv_df = None
        

    def set_rules(self, rules):
        """
        Manually set binning rules.
        """
        self.rules = rules

    def recalculate_stats(self, X, y):
        """
        Recalculate WOE and IV for existing rules on new data.
        """
        iv_stats = []
        for col, rule in self.rules.items():
            if col not in X.columns:
                continue
                
            if rule['type'] == 'numeric':
                cutoffs = rule['cutoffs']
                woe_map, iv, bin_stats = self._calculate_woe(X[col], y, cutoffs)
            elif rule['type'] == 'categorical':
                categories = rule['categories']
                woe_map, iv, bin_stats = self._calculate_woe_categorical(X[col], y, categories)
                
            # Update rule with new stats (keep cutoffs/categories)
            self.rules[col]['woe_map'] = woe_map
            self.rules[col]['iv'] = iv
            self.rules[col]['stats'] = bin_stats
            
            iv_stats.append({'variable': col, 'iv': iv})
            
        self.iv_df = pd.DataFrame(iv_stats).sort_values('iv', ascending=False)

    def fit(self, X, y, monotonicity_constraints=None):
        """
        Fit binning rules for each column in X using y.
        monotonicity_constraints: dict {col: 'increasing'|'decreasing'|'u-shape'}
        """
        self.rules = {}
        iv_stats = []
        
        for col in X.columns:
            if X[col].dtype == 'object' or X[col].dtype.name == 'category':
                # Categorical binning
                categories = self._categorical_binning(X[col], y, max_bins=self.max_bins)
                woe_map, iv, bin_stats = self._calculate_woe_categorical(X[col], y, categories)
                
                self.rules[col] = {
                    'categories': categories, # List of lists (bins)
                    'woe_map': woe_map,
                    'iv': iv,
                    'stats': bin_stats,
                    'type': 'categorical'
                }
            else:
                # Numerical binning
                constraint = monotonicity_constraints.get(col) if monotonicity_constraints else None
                
                if self.method == 'chi':
                    cutoffs = self._chimerge(X[col], y, max_bins=self.max_bins, constraint=constraint)
                elif self.method == 'dt':
                    cutoffs = self._dt_binning(X[col], y, max_bins=self.max_bins)
                else:
                    cutoffs = self._quantile_binning(X[col], max_bins=self.max_bins)
                    
                # Calculate WOE and IV
                woe_map, iv, bin_stats = self._calculate_woe(X[col], y, cutoffs)
                
                self.rules[col] = {
                    'cutoffs': cutoffs,
                    'woe_map': woe_map,
                    'iv': iv,
                    'stats': bin_stats,
                    'type': 'numeric'
                }
            
            iv_stats.append({'variable': col, 'iv': iv})
            
        self.iv_df = pd.DataFrame(iv_stats).sort_values('iv', ascending=False)
        return self

    def transform(self, X):
        """
        Apply WOE transformation.
        """
        X_woe = X.copy()
        for col, rule in self.rules.items():
            if col not in X.columns:
                continue
            
            if rule['type'] == 'numeric':
                cutoffs = rule['cutoffs']
                woe_map = rule['woe_map']
                
                # Binning
                bins = [-np.inf] + cutoffs + [np.inf]
                
                X_woe[col] = pd.cut(X[col], bins=bins, labels=False, include_lowest=True)
                X_woe[col] = X_woe[col].map(woe_map).astype(float)
                
                if -1 in woe_map: # Convention: -1 for NaN
                    X_woe[col] = X_woe[col].fillna(woe_map[-1])
                    
            elif rule['type'] == 'categorical':
                woe_map = rule['woe_map']
                # woe_map for categorical is {value: woe}
                # But we stored it by bin index in _calculate_woe_categorical
                # So we need to map value -> bin_index -> woe
                
                # Create a direct map: value -> woe
                value_woe_map = {}
                categories = rule['categories'] # List of lists
                
                for bin_idx, bin_values in enumerate(categories):
                    woe = woe_map[bin_idx]
                    for val in bin_values:
                        value_woe_map[val] = woe
                        
                # Map values
                X_woe[col] = X[col].map(value_woe_map)
                
                # Handle unknown values (treat as missing or new bin?)
                # For now, map to NaN/missing bin if exists
                if -1 in woe_map:
                    X_woe[col] = X_woe[col].fillna(woe_map[-1])
                else:
                    # If no missing bin, fill with 0 or mean woe? 
                    # Let's fill with 0 (neutral) for safety or error?
                    X_woe[col] = X_woe[col].fillna(0) 
                    
                X_woe[col] = X_woe[col].astype(float)

        return X_woe

    def _categorical_binning(self, x, y, max_bins=10):
        """
        Bin categorical features by grouping similar bad rates.
        """
        df = pd.DataFrame({'x': x, 'y': y})
        # Fill NA with 'Missing' for analysis
        df['x'] = df['x'].fillna('Missing')
        
        # Calculate Bad Rate per category
        grouped = df.groupby('x', observed=False)['y'].agg(['count', 'sum'])
        grouped['bad_rate'] = grouped['sum'] / grouped['count']
        grouped = grouped.sort_values('bad_rate')
        
        # We can treat bad_rate as a continuous variable and use ChiMerge?
        # Or just use the ordered categories as "x" for ChiMerge.
        
        # Strategy: 
        # 1. Sort categories by Bad Rate.
        # 2. Treat them as ordinal (0, 1, 2...).
        # 3. Apply ChiMerge on these ordinal values.
        
        categories_ordered = grouped.index.tolist()
        cat_map = {cat: i for i, cat in enumerate(categories_ordered)}
        
        df['x_ordinal'] = df['x'].map(cat_map)
        
        # Use ChiMerge on x_ordinal
        cutoffs = self._chimerge(df['x_ordinal'], df['y'], max_bins=max_bins)
        
        # Convert cutoffs back to groups of categories
        # cutoffs are indices in categories_ordered
        # e.g., cutoffs = [2, 5] means bins: [0,1,2], [3,4,5], [6...]
        
        bins = []
        start_idx = 0
        
        # Ensure cutoffs are integers
        cutoffs = [int(c) for c in cutoffs]
        
        for end_idx in cutoffs:
            # bin includes end_idx (because chimerge cutoffs are inclusive on left in my logic? 
            # Wait, _chimerge returns right edges of intervals from pd.cut(include_lowest=True).
            # pd.cut: (a, b]. 
            # My _chimerge uses pd.cut then aggregates. 
            # Let's assume cutoffs are the split points.
            # If cutoffs=[2], splits are <=2 and >2.
            
            # The indices are 0 to N-1.
            # If cutoff is 2.5, it splits 0,1,2 | 3,4...
            
            # _chimerge returns specific values from x as cutoffs.
            # Since x are integers 0..N, cutoffs will be integers (or floats if mean).
            # In my _chimerge, I take right edge of interval.
            # If unique values, cutoffs are values.
            
            # Let's look at _chimerge again.
            # "cutoffs = grouped.index.tolist()[:-1]" (for small unique)
            # So cutoffs are the values.
            # If values are 0,1,2,3. Cutoffs could be 0, 1, 2.
            # If cutoff=0, split is <=0.
            
            end_idx = int(end_idx)
            current_bin = categories_ordered[start_idx : end_idx + 1]
            bins.append(current_bin)
            start_idx = end_idx + 1
            
        # Add last bin
        if start_idx < len(categories_ordered):
            bins.append(categories_ordered[start_idx:])
            
        return bins

    def _calculate_woe_categorical(self, x, y, categories):
        """
        Calculate WOE for categorical bins.
        """
        # categories is list of lists
        # Map x to bin index
        val_to_bin = {}
        for idx, cat_list in enumerate(categories):
            for val in cat_list:
                val_to_bin[val] = idx
        
        # Handle Missing if not in categories (though we handled it in binning)
        # If x has real NaNs and we mapped 'Missing' in binning, we need to map NaN here too?
        # x is original series.
        
        df = pd.DataFrame({'x': x, 'y': y})
        # We need to temporarily fillna to map
        df['x_temp'] = df['x'].fillna('Missing')
        
        df['bin'] = df['x_temp'].map(val_to_bin)
        
        # Handle values not in map (e.g. new values not seen in fit) -> NaN bin or -1
        df['bin'] = df['bin'].fillna(-1)
        
        total_good = (df['y'] == 0).sum()
        total_bad = (df['y'] == 1).sum()
        
        grouped = df.groupby('bin', observed=False)['y'].agg(['count', 'sum'])
        grouped['bad'] = grouped['sum']
        grouped['good'] = grouped['count'] - grouped['sum']
        
        woe_map = {}
        iv_total = 0
        bin_stats = []
        
        for bin_idx, row in grouped.iterrows():
            good = row['good']
            bad = row['bad']
            
            if good == 0: good = 0.5
            if bad == 0: bad = 0.5
            
            dist_good = good / total_good
            dist_bad = bad / total_bad
            
            woe = np.log(dist_bad / dist_good)
            iv = (dist_bad - dist_good) * woe
            
            woe_map[bin_idx] = woe
            iv_total += iv
            
            # For categorical, min/max are not relevant, but list of values
            if bin_idx != -1 and bin_idx < len(categories):
                bin_label = str(categories[int(bin_idx)])[:50] + "..." # Truncate for display
            else:
                bin_label = "Unmapped"
                
            bin_stats.append({
                'bin_index': bin_idx,
                'min': bin_label, 
                'max': bin_label,
                'good': row['good'],
                'bad': row['bad'],
                'total': row['count'],
                'bad_rate': row['bad'] / row['count'],
                'woe': woe,
                'iv': iv
            })
            
        return woe_map, iv_total, bin_stats



    def _detect_trend(self, x, y):
        """
        Detect optimal trend (1: increasing, -1: decreasing) using Spearman correlation.
        """
        correlation = pd.concat([x, y], axis=1).corr(method='spearman').iloc[0, 1]
        if abs(correlation) < 0.05: # Weak correlation, treat as no constraint? Or weak constraint.
            # Usually we prefer a direction even if weak.
            return 1 if correlation >= 0 else -1
        return 1 if correlation > 0 else -1

    def _enforce_monotonicity(self, stats, cutoffs, constraint):
        """
        Post-processing to enforce monotonicity.
        stats: [[good, bad], ...]
        cutoffs: list of cutoffs
        constraint: 1 (increasing), -1 (decreasing)
        """
        while len(stats) > 2: # Don't merge if only 2 bins
            bad_rates = stats[:, 1] / stats.sum(axis=1)
            violations = []
            
            for i in range(len(bad_rates) - 1):
                if constraint == 1: # Increasing
                    if bad_rates[i] > bad_rates[i+1]:
                        violations.append(i)
                elif constraint == -1: # Decreasing
                    if bad_rates[i] < bad_rates[i+1]:
                        violations.append(i)
            
            if not violations:
                break
                
            # Merge logic:
            # 1. Prefer merging small bins involved in violation
            # 2. Else merge pair with min Chi2
            
            merge_idx = -1
            min_chi2 = np.inf
            
            for idx in violations:
                # Calculate Chi2 for the violating pair
                chi2 = self._calculate_chi2(stats[idx], stats[idx+1])
                if chi2 < min_chi2:
                    min_chi2 = chi2
                    merge_idx = idx
            
            # Execute Merge
            stats[merge_idx] += stats[merge_idx+1]
            stats = np.delete(stats, merge_idx+1, axis=0)
            del cutoffs[merge_idx]
            
        return stats, cutoffs

    def _chimerge(self, x, y, max_bins=10, constraint=None):
        """
        ChiMerge algorithm implementation.
        """
        df = pd.DataFrame({'x': x, 'y': y})
        df = df.dropna()
        
        # Initial fine binning (e.g. 50 bins or unique values if small)
        if df['x'].nunique() > 50:
            _, bins = pd.qcut(df['x'], 50, retbins=True, duplicates='drop')
            df['bin'] = pd.cut(df['x'], bins=bins, include_lowest=True)
            grouped = df.groupby('bin', observed=False)['y'].agg(['count', 'sum'])
            grouped['bad'] = grouped['sum']
            grouped['good'] = grouped['count'] - grouped['sum']
            grouped = grouped.reset_index()
            intervals = grouped['bin'].tolist()
            cutoffs = [i.right for i in intervals]
            cutoffs = cutoffs[:-1]
            stats = grouped[['good', 'bad']].values
        else:
            grouped = df.groupby('x')['y'].agg(['count', 'sum'])
            grouped['bad'] = grouped['sum']
            grouped['good'] = grouped['count'] - grouped['sum']
            stats = grouped[['good', 'bad']].values
            cutoffs = grouped.index.tolist()[:-1]

        total_samples = len(df)
        min_samples = total_samples * self.min_bin_pct
        
        # Merge Logic
        while True:
            # 1. Check for small bins and merge them first
            # Calculate bin sizes
            # stats is [[good, bad], ...]
            # Bin i corresponds to interval (cutoffs[i-1], cutoffs[i]] (logic is roughly this)
            # Actually, len(stats) = len(cutoffs) + 1
            
            bin_sizes = stats.sum(axis=1)
            small_bin_indices = np.where(bin_sizes < min_samples)[0]
            
            if len(stats) <= 2: # Don't merge if only 2 bins left
                break
                
            merge_idx = -1
            
            # If there are small bins, prioritize merging the smallest one
            if len(small_bin_indices) > 0:
                # Find smallest bin
                min_size_idx = np.argmin(bin_sizes)
                # If smallest bin is >= min_samples, then no small bins (logic check)
                if bin_sizes[min_size_idx] >= min_samples:
                    # Proceed to Chi2 check
                    pass
                else:
                    # Merge smallest bin with neighbor having lower Chi2
                    idx = min_size_idx
                    if idx == 0:
                        merge_idx = 0 # Merge 0 and 1
                    elif idx == len(stats) - 1:
                        merge_idx = idx - 1 # Merge last and second last
                    else:
                        # Calculate Chi2 with left and right
                        chi2_left = self._calculate_chi2(stats[idx-1], stats[idx])
                        chi2_right = self._calculate_chi2(stats[idx], stats[idx+1])
                        if chi2_left < chi2_right:
                            merge_idx = idx - 1
                        else:
                            merge_idx = idx
            
            # If no small bins need merging, check max_bins
            if merge_idx == -1:
                if len(stats) <= max_bins:
                    # Satisfied max_bins, now enforce monotonicity if needed
                    # We do this as a post-check to avoid over-merging early?
                    # Or we do it iteratively.
                    
                    # If we are here, we have <= max_bins and no small bins.
                    # Now check trend.
                    break # We break here and do a final pass for monotonicity
                
                # Merge based on min Chi2
                chi2_list = []
                for i in range(len(stats) - 1):
                    chi2_list.append(self._calculate_chi2(stats[i], stats[i+1]))
                
                merge_idx = np.argmin(chi2_list)
            
            # Execute Merge
            stats[merge_idx] += stats[merge_idx+1]
            stats = np.delete(stats, merge_idx+1, axis=0)
            del cutoffs[merge_idx]
            
        # Post-processing: Enforce Monotonicity
        # If constraint is None, detect it automatically
        if constraint is None:
            constraint = self._detect_trend(x, y)
            
        if constraint in [1, -1]:
            stats, cutoffs = self._enforce_monotonicity(stats, cutoffs, constraint)
            
        return cutoffs

    def _calculate_chi2(self, bin1, bin2):
        obs = np.array([bin1, bin2])
        row_sums = obs.sum(axis=1)
        col_sums = obs.sum(axis=0)
        total = obs.sum()
        
        if 0 in row_sums or 0 in col_sums:
            return 0
            
        expected = np.outer(row_sums, col_sums) / total
        with np.errstate(divide='ignore', invalid='ignore'):
            chi2 = ((obs - expected)**2 / expected).sum()
        return chi2


    def _dt_binning(self, x, y, max_bins=10):
        # ... implementation using DecisionTreeClassifier ...
        df = pd.DataFrame({'x': x, 'y': y}).dropna()
        dt = DecisionTreeClassifier(max_leaf_nodes=max_bins, min_samples_leaf=0.05)
        dt.fit(df[['x']], df['y'])
        thresholds = sorted(dt.tree_.threshold[dt.tree_.threshold != -2])
        return thresholds

    def _quantile_binning(self, x, max_bins=10):
        _, bins = pd.qcut(x, max_bins, retbins=True, duplicates='drop')
        return bins[1:-1].tolist()

    def _calculate_woe(self, x, y, cutoffs):
        """
        Calculate WOE and IV for given cutoffs.
        """
        # bins = [-inf] + cutoffs + [inf]
        bins = [-np.inf] + cutoffs + [np.inf]
        
        df = pd.DataFrame({'x': x, 'y': y})
        df['bin'] = pd.cut(df['x'], bins=bins, include_lowest=True, labels=False)
        
        # Handle Missing
        # Map nan to -1
        df['bin'] = df['bin'].fillna(-1)
        
        total_good = (df['y'] == 0).sum()
        total_bad = (df['y'] == 1).sum()
        
        grouped = df.groupby('bin', observed=False)['y'].agg(['count', 'sum'])
        grouped['bad'] = grouped['sum']
        grouped['good'] = grouped['count'] - grouped['sum']
        
        woe_map = {}
        iv_total = 0
        bin_stats = []
        
        for bin_idx, row in grouped.iterrows():
            good = row['good']
            bad = row['bad']
            
            # Smooth to avoid inf
            if good == 0: good = 0.5
            if bad == 0: bad = 0.5
            
            dist_good = good / total_good
            dist_bad = bad / total_bad
            
            woe = np.log(dist_bad / dist_good)
            iv = (dist_bad - dist_good) * woe
            
            woe_map[bin_idx] = woe
            iv_total += iv
            
            bin_stats.append({
                'bin_index': bin_idx,
                'min': bins[int(bin_idx)] if bin_idx != -1 else np.nan,
                'max': bins[int(bin_idx)+1] if bin_idx != -1 else np.nan,
                'good': row['good'],
                'bad': row['bad'],
                'total': row['count'],
                'bad_rate': row['bad'] / row['count'],
                'woe': woe,
                'iv': iv
            })
            
        return woe_map, iv_total, bin_stats
