import pandas as pd
import numpy as np
from sklearn.preprocessing import KBinsDiscretizer
from typing import Dict, List, Tuple
import matplotlib.pyplot as plt
import seaborn as sns

class MultiFeatureRuleMiner:
    """
    多特征交叉规则生成类，用于生成双特征交叉分析结果
    """
    
    def __init__(self, df: pd.DataFrame, target_col: str = 'ISBAD'):
        """
        初始化多特征规则挖掘器
        
        参数:
            df: 输入的数据集
            target_col: 目标字段名，默认为'ISBAD'
        """
        self.df = df.copy()
        self.target_col = target_col
        
        # 筛选出数值型特征
        self.numeric_features = [col for col in self.df.columns 
                                if col != target_col and pd.api.types.is_numeric_dtype(self.df[col])]
        
        # 筛选出类别型特征
        self.categorical_features = [col for col in self.df.columns 
                                    if col != target_col and not pd.api.types.is_numeric_dtype(self.df[col])]
        
    def _bin_numeric_feature(self, feature: str, max_bins: int = 20) -> pd.Series:
        """
        对数值型特征进行分箱处理
        
        参数:
            feature: 特征名
            max_bins: 最大分箱数量，默认为20
            
        返回:
            分箱后的特征值
        """
        data = self.df[feature].dropna().values.reshape(-1, 1)
        discretizer = KBinsDiscretizer(n_bins=max_bins, encode='ordinal', strategy='quantile', quantile_method='averaged_inverted_cdf')
        binned = discretizer.fit_transform(data)
        
        # 转换为Series，保留原始索引
        binned_series = pd.Series(binned.flatten(), index=self.df[feature].dropna().index)
        
        # 填充缺失值
        result = pd.Series(index=self.df.index, dtype=float)
        result.loc[binned_series.index] = binned_series
        
        return result
    
    def _prepare_feature(self, feature: str, max_unique_values: int = 20) -> pd.Series:
        """
        准备特征，对取值较多的特征进行分箱处理
        
        参数:
            feature: 特征名
            max_unique_values: 最大允许的唯一值数量，超过则进行分箱
            
        返回:
            处理后的特征值
        """
        unique_count = self.df[feature].nunique()
        
        if unique_count > max_unique_values:
            # 对数值型特征进行分箱
            if feature in self.numeric_features:
                return self._bin_numeric_feature(feature, max_unique_values)
            # 对类别型特征，保留前max_unique_values个最常见的值，其余归为"其他"
            else:
                top_values = self.df[feature].value_counts().head(max_unique_values).index
                return self.df[feature].apply(lambda x: x if x in top_values else '其他')
        else:
            return self.df[feature]
    
    def generate_cross_matrix(self, feature1: str, feature2: str, 
                             max_unique_values: int = 20) -> pd.DataFrame:
        """
        生成双特征交叉矩阵，包含badrate、样本占比等关键指标
        
        参数:
            feature1: 第一个特征名
            feature2: 第二个特征名
            max_unique_values: 最大允许的唯一值数量，超过则进行分箱
            
        返回:
            交叉矩阵，包含各种统计指标
        """
        # 准备两个特征
        prepared_feature1 = self._prepare_feature(feature1, max_unique_values)
        prepared_feature2 = self._prepare_feature(feature2, max_unique_values)
        
        # 创建交叉表数据
        cross_data = pd.DataFrame({
            'feature1': prepared_feature1,
            'feature2': prepared_feature2,
            'target': self.df[self.target_col]
        })
        
        # 计算基本统计量
        # 计算每个交叉组合的样本数
        count_matrix = pd.crosstab(cross_data['feature1'], cross_data['feature2'], 
                                  rownames=[feature1], colnames=[feature2])
        
        # 计算每个交叉组合的坏样本数
        bad_count_matrix = pd.crosstab(cross_data['feature1'], cross_data['feature2'], 
                                      values=cross_data['target'], aggfunc='sum',
                                      rownames=[feature1], colnames=[feature2])
        
        # 计算每个交叉组合的坏样本率
        badrate_matrix = bad_count_matrix / count_matrix
        
        # 计算每个交叉组合的样本占比
        total_samples = count_matrix.sum().sum()
        sample_ratio_matrix = count_matrix / total_samples
        
        # 计算lift值
        # 总样本坏样本率
        total_badrate = self.df[self.target_col].mean()
        lift_matrix = badrate_matrix / total_badrate
        
        # 整合所有矩阵到一个MultiIndex DataFrame中
        # 创建一个空的MultiIndex DataFrame
        features = [feature1, feature2]
        metrics = ['count', 'bad_count', 'badrate', 'sample_ratio', 'lift']
        
        # 生成行索引（feature1的唯一值）
        rows = count_matrix.index
        # 生成列索引（feature2的唯一值 + 指标）
        cols = pd.MultiIndex.from_product([count_matrix.columns, metrics], 
                                         names=[feature2, 'metric'])
        
        # 填充数据
        cross_matrix = pd.DataFrame(index=rows, columns=cols)
        
        for f2_val in count_matrix.columns:
            cross_matrix[(f2_val, 'count')] = count_matrix[f2_val]
            cross_matrix[(f2_val, 'bad_count')] = bad_count_matrix[f2_val]
            cross_matrix[(f2_val, 'badrate')] = badrate_matrix[f2_val]
            cross_matrix[(f2_val, 'sample_ratio')] = sample_ratio_matrix[f2_val]
            cross_matrix[(f2_val, 'lift')] = lift_matrix[f2_val]
        
        return cross_matrix
    
    def get_cross_rules(self, feature1: str, feature2: str, 
                       top_n: int = 10, metric: str = 'lift',
                       max_unique_values: int = 20) -> pd.DataFrame:
        """
        从交叉矩阵中提取top规则
        
        参数:
            feature1: 第一个特征名
            feature2: 第二个特征名
            top_n: 返回的规则数量，默认为10
            metric: 排序指标，默认为'lift'
            max_unique_values: 最大允许的唯一值数量，超过则进行分箱
            
        返回:
            包含top规则的DataFrame
        """
        # 生成交叉矩阵
        cross_matrix = self.generate_cross_matrix(feature1, feature2, max_unique_values)
        
        # 转换为长格式，方便排序
        long_df = cross_matrix.stack(level=0, future_stack=True).reset_index()
        long_df.rename(columns={feature1: 'feature1_value', feature2: 'feature2_value'}, inplace=True)
        
        # 按指定指标排序，返回top_n
        top_rules = long_df.sort_values(by=metric, ascending=False).head(top_n)
        
        # 添加规则描述
        top_rules['rule_description'] = top_rules.apply(
            lambda x: f"{feature1} = {x['feature1_value']} AND {feature2} = {x['feature2_value']}",
            axis=1
        )
        
        return top_rules
    
    def plot_cross_heatmap(self, feature1: str, feature2: str, 
                          metric: str = 'lift', max_unique_values: int = 20, 
                          figsize: Tuple[int, int] = (12, 10)):
        """
        绘制双特征交叉热力图
        
        参数:
            feature1: 第一个特征名
            feature2: 第二个特征名
            metric: 要可视化的指标，默认为'lift'
            max_unique_values: 最大允许的唯一值数量，超过则进行分箱
            figsize: 图表大小
            
        返回:
            matplotlib.pyplot对象
        """
        # 生成交叉矩阵
        cross_matrix = self.generate_cross_matrix(feature1, feature2, max_unique_values)
        
        # 提取指定指标的矩阵
        metric_matrix = cross_matrix.xs(metric, level='metric', axis=1)
        
        # 创建热力图
        plt.figure(figsize=figsize)
        sns.heatmap(metric_matrix, annot=True, fmt='.3f', cmap='coolwarm', 
                   cbar_kws={'label': metric})
        
        plt.title(f'{feature1} vs {feature2} - {metric} Heatmap')
        plt.tight_layout()
        
        return plt
    
    def generate_all_cross_rules(self, top_n: int = 10, metric: str = 'lift',
                                max_unique_values: int = 20, max_feature_pairs: int = 10) -> pd.DataFrame:
        """
        生成所有特征对的交叉规则，并返回top规则
        
        参数:
            top_n: 返回的规则数量，默认为10
            metric: 排序指标，默认为'lift'
            max_unique_values: 最大允许的唯一值数量，超过则进行分箱
            max_feature_pairs: 最大生成的特征对数量，默认为10
            
        返回:
            包含top规则的DataFrame
        """
        # 生成特征对组合
        from itertools import combinations
        features = self.numeric_features + self.categorical_features
        feature_pairs = list(combinations(features, 2))[:max_feature_pairs]  # 限制特征对数量
        
        all_rules = []
        
        for feature1, feature2 in feature_pairs:
            try:
                rules = self.get_cross_rules(feature1, feature2, top_n=top_n, 
                                            metric=metric, max_unique_values=max_unique_values)
                if not rules.empty:
                    all_rules.append(rules)
            except Exception as e:
                print(f"Error generating rules for {feature1} and {feature2}: {e}")
                continue
        
        if not all_rules:
            return pd.DataFrame()
        
        # 合并所有规则，取top_n
        combined_rules = pd.concat(all_rules, ignore_index=True)
        return combined_rules.sort_values(by=metric, ascending=False).head(top_n)
