from typing import List, TypedDict, Literal, Final, cast
from abc import ABC, abstractmethod
from datetime import date
from random import seed
from math import log

import pandas as pd
from pandas import DataFrame, Series
from sklearn.metrics import roc_curve, roc_auc_score
from matplotlib import pyplot as plt
from toad.plot import bin_plot, badrate_plot


class AutomatedModeling(ABC):
    seed(a=42) # 固定随机种子
    Evaluation = TypedDict(
        "Evaluation",
        {
            "train_ks": float,
            "validation_ks": float,
            "oot_ks": float,
            "train_auc": float,
            "validation_auc": float,
            "oot_auc": float,
            "train_divergence": float,
            "validation_divergence": float,
            "oot_divergence": float,
            "model_psi": float
        }
    )

    def __init__(self, data: DataFrame, target: str, time: str, not_features: List[str]) -> None:
        """传入规定格式的数据以实例化 AutomatedModeling 对象

        Args:
            data (DataFrame): 样本数据
            target (str): 字段值为 0-1 变量
            time (str): 字段值为 date 对象
            not_features (List[str]): 除相应变量外的其他非特征字段（如时间、订单号、包名等）
        """
        self.data: DataFrame = data.reset_index(drop=True, inplace=False) # 样本
        self.target: str = target # 响应变量
        self.time: str = time # 事件日期
        self.not_features: List[str] = not_features # 除响应变量以外的其他非特征字段
        self.features: List[str] = [col for col in self.data.columns if col not in self.not_features + [self.target]] # 除 target 和 not_features 以外的所有特征字段
        self.is_train: str = "is_train" # 新增用以区分训练集、验证集和测试集的字段
        self.train: DataFrame
        self.validation: DataFrame
        self.oot: DataFrame
        self.used_features: List[str] # 最终入模特征
        self.evaluation: AutomatedModeling.Evaluation
        self.score: str # 模型分字段
        self.split_validation_set: bool # 是否划分验证集

    def split_train_oot(self, split_time: date, split_validation_set: bool = False, validation_pct: float = 0.2) -> None:
        """划分训练集和测试集（如不划分验证集则默认将测试集复制一份作为验证集）

        Args:
            split_time (date): 用于划分训练集和测试集的临界日期
            split_validation_set (bool, optional): 是否需要划分验证集. Defaults to False.
            validation_pct (float, optional): 验证集占训练集的样本. Defaults to 0.2.
        """
        self.split_validation_set = split_validation_set
        self.data[self.is_train] = -1 # 在原始数据集中新增一个用于区分训练集、验证集和测试集的字段`is_train`: {1: "训练集", 0: "测试集", 2: "验证集"}
        
        self.train = self.data[self.data[self.time] <= split_time]
        self.oot = self.data[self.data[self.time] > split_time]
        
        if split_validation_set:
            self.validation = self.train.sample(frac=validation_pct, random_state=42)
            self.train.drop(index=self.validation.index, axis=0, inplace=True) # 从训练集中剔除掉验证集样本
        else:
            self.validation = self.oot.copy()
        
        self.train[self.is_train] = 1
        self.validation[self.is_train] = 2
        self.oot[self.is_train] = 0
        self.data = pd.concat(objs=[self.train, self.validation, self.oot], axis=0)

        self.train.reset_index(drop=True, inplace=True)
        self.validation.reset_index(drop=True, inplace=True)
        self.oot.reset_index(drop=True, inplace=True)
        self.data.reset_index(drop=True, inplace=True)

    @abstractmethod
    def eliminate_low_iv_features(self, selected_features: List[str], iv: float = 0.02, train_validation_oot: Literal[0, 1, 2] = 1) -> List[str]:
        """剔除 iv 值较低的变量

        Args:
            selected_features (List[str]): 潜在的入模变量
            iv (float, optional): iv 阈值. Defaults to 0.02.
            train_validation_oot (Literal[0, 1, 2], optional): {1: "训练集", 0: "测试集", 2: "验证集"}. Defaults to 1.

        Returns:
            List[str]: iv 值超过指定阈值的变量
        """
        pass

    @abstractmethod
    def eliminate_unstable_features(self, selected_features: List[str], psi_threshold: float = 0.1, psi_original: bool = True) -> List[str]:
        """剔除稳定性较差的变量

        Args:
            selected_features (List[str]): 潜在的入模变量
            psi_threshold (float, optional): psi筛选的阈值. Defaults to 0.1.
            psi_original (bool, optional): 使用原始数据还是特征工程后的数据（True: 原始数据, False: 特征工程后的数据）. Defaults to True.

        Returns:
            List[str]: 稳定性较好的变量
        """
        pass

    @abstractmethod
    def calculate_model_score_psi(self, n_bins: int) -> float:
        """计算模型分离散化后的 PSI

        Args:
            n_bins (int, optional): 计算模型分 PSI 时的分箱个数

        Returns:
            float: 模型分离散化后在训练集和测试集上的 PSI
        """
        pass

    @abstractmethod
    def evaluate(self, n_bins: int) -> Evaluation:
        """模型评估

        Args:
            n_bins (int, optional): 计算模型分 PSI 时的分箱个数

        Returns:
            Evaluation:  模型的量化指标（KS, AUC, model_score_psi）
        """
        pass

    @abstractmethod
    def ks_bucket(self, train_validation_oot: Literal[0, 1, 2], n_bins: int) -> DataFrame:
        """模型分分箱

        Args:
            train_validation_oot (Literal[0, 1, 2], optional): {1: "训练集", 0: "测试集", 2: "验证集"}
            n_bins (int, optional): 分箱数量

        Raises:
            ValueError: 参数传值错误

        Returns:
            DataFrame: 模型分箱后在各个箱中的效果
        """
        pass
    
    @abstractmethod
    def predict(self, X: DataFrame) -> Series:
        """对新数据进行预测

        Args:
            X (DataFrame): 原始数据

        Returns:
            Series: 模型分
        """
        pass
    
    @staticmethod
    def proba2score(prob: float, pdo: int = 20, rate: float = 2, base_odds: float = 1.22, base_score: int = 600) -> float:
        """将概率转化为分数

        Args:
            prob (float): 响应概率
            pdo (int, optional): pdo. Defaults to 20.
            rate (float, optional): rate. Defaults to 2.
            base_odds (float, optional): base_odds. Defaults to 1.22.
            base_score (int, optional): base_score. Defaults to 600.

        Returns:
            float: 响应分数
        """
        factor: float = pdo / log(rate)
        offset: float = base_score + factor * log(base_odds)
        return offset - factor * log(prob / (1-prob))
    
    @abstractmethod
    def get_features_index_report(self, select_features: List[str]) -> DataFrame:
        """输出变量的 iv、psi 等评价指标

        Args:
            select_features (List[str]): 特征列表

        Returns:
            DataFrame: 变量的 iv、psi 等评价指标
        """
        pass

    @abstractmethod
    def export_model(self, path: str) -> None:
        """将模型以二进制导出为 pkl 格式

        Args:
            path (str): 导出路径
        """
        pass

    @abstractmethod
    def generate_yaml(self, model_file: str, model_name: str) -> None:
        """生成模型配置的标准 yaml 文件

        Args:
            model_file (str): 序列化后的模型文件路径
            model_name (str): 模型名称
        """
        pass
    
    @staticmethod
    def plot_ks_curve(target: Series, proba: Series, title: str = "KS Curve") -> None:
        """绘制 KS 曲线

        Args:
            target (Series): 实际标签（0-1变量）
            proba (Series): 响应概率
        """
        target_name: Final[str] = "target"
        proba_name: Final[str] = "proba"
        
        data: DataFrame = DataFrame(data={target_name: target, proba_name: proba})
        data.sort_values(by=proba_name, inplace=True, ignore_index=True)

        count: int = target.count()
        bads_count: int = target.sum()
        goods_count: int = count - bads_count

        data["cum_bads_prop"] = data[target_name].cumsum() / bads_count
        data["cum_goods_prop"] = (1 - data[target_name]).cumsum() / goods_count
        data["ks"] = abs(data["cum_bads_prop"] - data["cum_goods_prop"])
        data["count"] = 1
        data["cum_total_prop"] = data["count"].cumsum() / count

        max_ks_idx: int = cast(
            typ=int, 
            val=data["ks"].idxmax()
        )
        max_ks: float = data["ks"].max()
        ks_max_population: float = cast(
            typ=float, 
            val=data.loc[max_ks_idx, "cum_total_prop"]
        )

        fig, ax = plt.subplots()
        ax.plot(data["cum_total_prop"], data["cum_bads_prop"], label="cum_bads_prop", color="darkorange") # 累计坏样本分布
        ax.plot(data["cum_total_prop"], data["cum_goods_prop"], label="cum_goods_prop", color="firebrick") # 累计好样本分布
        ax.plot(data["cum_total_prop"], data["ks"], label="ks", color="darkblue") # KS 曲线
        ax.plot([ks_max_population, ks_max_population], [max_ks, 0], linestyle="--", color="red")
        ax.scatter(x=[ks_max_population], y=[max_ks], color="red")
        ax.fill_between( # 填充两条累积曲线之间的区域
            x=data["cum_total_prop"], 
            y1=data["cum_bads_prop"], 
            y2=data["cum_goods_prop"], 
            alpha=0.2, # 覆盖区域的透明度 [0, 1],其值越大，表示越不透明
            color="blue"
        )
        ax.legend()
        
        ax.annotate( # 添加注释文字
            text=f"ks = {max_ks:.4f} at {ks_max_population:.2%}", 
            xy=(ks_max_population, max_ks),
            xytext=(ks_max_population+0.05, max_ks+0.03), # 注释位置
            fontsize=14
        )

        # 设置坐标轴
        ax.set_xlim(left=0, right=1)
        ax.set_ylim(bottom=0, top=1)
        ax.set_xticks(ticks=[i / 10 for i in range(0, 11, 2)])
        ax.set_yticks(ticks=[i / 10 for i in range(0, 11, 2)])

        # 设置标签
        ax.set_xlabel(xlabel="Cum Total Prop", fontsize=14)
        ax.set_ylabel(ylabel="Cum Bads / Goods Prop", fontsize=14)
        ax.set_title(label=title, fontsize=15)
        
        ax.grid(visible=True, alpha=0.3, linestyle="-", linewidth=0.8) # 网格线

        # 设置边框样式
        for spine in ax.spines.values():
            spine.set_color(c="blue")
            spine.set_linewidth(w=2)
        
        fig.patch.set_facecolor(color="white") # 设置图形背景颜色
        plt.show()

        print(f"KS: {max_ks}")

    @staticmethod
    def plot_auc_curve(target: Series, proba: Series, title: str = "ROC Curve") -> None:
        """绘制 ROC 曲线

        Args:
            target (Series): 实际标签（0-1变量）
            proba (Series): 响应概率
        """
        fpr, tpr, _=roc_curve(target, proba, pos_label=None, sample_weight=None, drop_intermediate=True)
        auc: float = float(roc_auc_score(target,proba))

        fig, ax = plt.subplots()
        ax.plot([0,1], [0,1], linestyle="--", linewidth=1.2, label="Base Line")
        ax.plot(fpr, tpr, label="ROC Curve")
        ax.fill_between( # ROC 曲线下的区域
            x=fpr, 
            y1=0, 
            y2=tpr, 
            alpha=0.2, # 覆盖区域的透明度 [0, 1],其值越大，表示越不透明
            color="blue"
        )
        ax.legend()

        ax.annotate( # 添加注释文字
            text=f"auc = {auc:.4f}", 
            xy=(0.6, 0.2), # 注释位置
            fontsize=14
        )

        # 设置坐标轴
        ax.set_xlim(left=0, right=1)
        ax.set_ylim(bottom=0, top=1)
        ax.set_xticks(ticks=[i / 10 for i in range(0, 11, 2)])
        ax.set_yticks(ticks=[i / 10 for i in range(0, 11, 2)])

        # 设置标签
        ax.set_xlabel(xlabel="False Positive Rate", fontsize=14)
        ax.set_ylabel(ylabel="True Positive Rate", fontsize=14)
        ax.set_title(label=title, fontsize=15)
        
        ax.grid(visible=True, alpha=0.3, linestyle="-", linewidth=0.8) # 网格线

        # 设置边框样式
        for spine in ax.spines.values():
            spine.set_color(c="blue")
            spine.set_linewidth(w=2)
        
        fig.patch.set_facecolor(color="white") # 设置图形背景颜色
        plt.show()

        print(f"AUC: {auc}")

    @staticmethod
    def model_effect(data: DataFrame, model_score: str, target: str, time: str, is_train: str, n_bins: int = 10) -> None:
        """模型效果可视化

        Args:
            data (DataFrame): 模型评估数据集（含其他参数字段）
            model_score (str): 模型分字段
            target (str): 标签字段
            time (str): 时间字段
            is_train (str): 数据集切分字段，{1: "训练集", 0: "时间外样本"}
            n_bins (int, optional): 等频分箱箱数. Defaults to 10.
        """
        data[f"{model_score}_bin"] = pd.qcut(x=data[model_score], q=n_bins, duplicates="drop")
        data[f"{model_score}_bin"] = data[f"{model_score}_bin"].astype(dtype=str)
        
        for i in (1, 0):
            label: str = "train" if i == 1 else "oot"
            bin_plot(frame=data[data[is_train]==i], x=f"{model_score}_bin", target=target)
            plt.title(label=label)
        
        badrate_plot(frame=data, x=time, target=target, by=f"{model_score}_bin")
        plt.title(label=model_score)
        plt.show()

    @staticmethod
    def divergence(data: DataFrame, score: str, target: str) -> float:
        """计算 divergence 值

        Args:
            data (DataFrame): 数据集（含其他参数字段）
            score (str): 模型分字段
            target (str): 标签字段

        Returns:
            float: divergence 值
        """
        good_data_score: Series = data[data[target]==0][score]
        bad_data_score: Series = data[data[target]==1][score]
        
        good_score_mean: float = good_data_score.mean()
        bad_score_mean: float = bad_data_score.mean()
        good_score_std: float = good_data_score.std()
        bad_score_std: float = bad_data_score.std()

        result: float = (good_score_mean - bad_score_mean)**2 / ((good_score_std**2 + bad_score_std**2) / 2)
        
        return result
