# -*-coding:utf-8-*-
import pickle
import warnings
import joblib
from lightgbm import LGBMClassifier
import logging
import numpy as np
import sys
import gc
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score
from sklearn.ensemble import RandomForestClassifier
from sklearn.cluster import KMeans
import lightgbm as lgb
import xgboost as xgb
from catboost import CatBoostClassifier
from sklearn.linear_model import LogisticRegression
import pandas as pd
import os
from sklearn.preprocessing import LabelEncoder

logging.getLogger("lightgbm").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")

TOP_K = 3
N_SPLITS = 5
SEED = 42
RARE_THR = 20
TE_SMOOTH = 50
MAX_EARLY_STOP = 200
LGB_ROUNDS = 2000
XGB_ROUNDS = 2000
CTB_ROUNDS = 2000

# 模型开关
HAS_XGB = True
HAS_CTB = True
HAS_RF = True
HAS_AUTO_ENGINEERING = True
verbose = False
verbose_int = 1 if verbose else 0


# 重构以后得代码
def train_models_refactoring(X, y, features, cat_cols, model_dir):
    """
    训练多个模型并进行预测
    """

    # 计算正负样本比例，用于模型处理数据不平衡逻辑
    def _compute_spw(y_data):
        return y_data.sum() / (len(y) - y_data.sum())

    spw_global = _compute_spw(y)
    print(f"[Info] Pos={int(y.sum())}, Neg={int(len(y) - y.sum())}, Global scale_pos_weight={spw_global:.3f}")
    # 定义启用的模型列表
    active_models = ['lgb', 'xgb', 'ctb', 'rf']
    # oof_preds用于存放所有训练集折外预测结果，models_info用于存放所有模型每折的模型和auc等信息
    oof_preds = {k: np.zeros(len(X)) for k in active_models}
    models_info = {k: [] for k in active_models}
    models_info.update({'features': features, 'cat_cols': cat_cols})
    # 仅用于LGB特征重要性
    fi_list = []
    # 定义交叉验证数据分折
    skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=SEED)

    # 所有模型的参数信息
    def get_params(model_name, fold, y_tr_spw):
        """获取特定模型的参数"""
        seed_current = SEED + fold

        if model_name == 'lgb':
            return {
                'objective': 'binary', 'metric': 'auc', 'boosting_type': 'gbdt',
                'learning_rate': 0.009, 'num_leaves': 165, 'max_depth': -1,
                'min_data_in_leaf': 45, 'feature_fraction': 0.62, 'bagging_fraction': 0.9,
                'bagging_freq': 3, 'lambda_l1': 0.001, 'lambda_l2': 2.41,
                'min_gain_to_split': 0.84, 'verbose': -1, 'seed': seed_current, 'n_jobs': -1,
                'scale_pos_weight': y_tr_spw * 0.5
            }

        elif model_name == 'xgb':
            return {
                'objective': 'binary:logistic', 'eval_metric': 'auc', 'tree_method': 'hist',
                'learning_rate': 0.007, 'max_depth': 7, 'subsample': 0.76,
                'colsample_bytree': 0.83, 'reg_alpha': 0.027, 'reg_lambda': 0.028,
                'gamma': 1.17, 'min_child_weight': 13.11,
                'random_state': seed_current, 'nthread': -1,
                'scale_pos_weight': y_tr_spw * 0.5
            }

        elif model_name == 'ctb':
            return {
                'loss_function': "Logloss", 'eval_metric': "AUC", 'learning_rate': 0.03,
                'depth': 7, 'l2_leaf_reg': 6.0, 'random_seed': seed_current,
                'iterations': CTB_ROUNDS, 'od_type': "Iter", 'od_wait': MAX_EARLY_STOP,
                'verbose': verbose
            }

        elif model_name == 'rf':
            return {
                'n_estimators': 900, 'max_depth': None, 'min_samples_split': 2,
                'min_samples_leaf': 1, 'max_features': "sqrt", 'class_weight': "balanced",
                'oob_score': True, 'n_jobs': -1, 'random_state': seed_current,
                'verbose': verbose_int
            }

    # 单独训练一个模型的代码
    def train_single_model(model_name, X_tr, y_tr, X_va, y_va, params, fold):
        """统一的训练、预测和保存接口"""
        model = None
        pred_val = None

        if model_name == 'lgb':
            d_tr = lgb.Dataset(X_tr, label=y_tr, categorical_feature=cat_cols or None, free_raw_data=False)
            d_va = lgb.Dataset(X_va, label=y_va, categorical_feature=cat_cols or None, reference=d_tr,
                               free_raw_data=False)
            model = lgb.train(params, d_tr, num_boost_round=LGB_ROUNDS, valid_sets=[d_tr, d_va],
                              callbacks=[lgb.early_stopping(stopping_rounds=MAX_EARLY_STOP, verbose=verbose)])
            pred_val = model.predict(X_va, num_iteration=model.best_iteration)

            # 只记录lgb的特征重要性
            fi = pd.DataFrame(
                {"feature": features, "importance": model.feature_importance(importance_type="gain"), "fold": fold})
            fi_list.append(fi)

        elif model_name == 'xgb':
            d_tr = xgb.DMatrix(X_tr, label=y_tr, feature_names=features)
            d_va = xgb.DMatrix(X_va, label=y_va, feature_names=features)
            model = xgb.train(params, d_tr, num_boost_round=XGB_ROUNDS, evals=[(d_tr, "train"), (d_va, "valid")],
                              early_stopping_rounds=MAX_EARLY_STOP, verbose_eval=verbose)
            pred_val = model.predict(d_va)

        elif model_name == 'ctb':
            cat_idx = [features.index(c) for c in cat_cols if c in features]
            model = CatBoostClassifier(**params)
            model.fit(X_tr, y_tr, eval_set=(X_va, y_va), cat_features=cat_idx, use_best_model=True)
            pred_val = model.predict_proba(X_va)[:, 1]

        elif model_name == 'rf':
            model = RandomForestClassifier(**params)
            model.fit(X_tr, y_tr)
            pred_val = model.predict_proba(X_va)[:, 1]

        # 计算当前使用模型的AUC
        auc_score = roc_auc_score(y_va, pred_val)
        # 保存模型
        filename = f'{model_name}_fold{fold}.model'
        save_path = os.path.join(model_dir, filename)
        if model_name == 'rf' or model_name == 'lr':
            joblib.dump(model, save_path, compress=3)
        else:
            model.save_model(save_path)
        # 打印当前模型、折次以及保存地址
        print(f'{model_name}_model_path: {save_path}')
        return model, pred_val, auc_score, save_path

    # N折交叉验证代码循环
    for fold, (tr_idx, val_idx) in enumerate(skf.split(X, y), 1):
        # 数据准备 (每折只做一次)
        X_tr, X_va = X.iloc[tr_idx][features], X.iloc[val_idx][features]
        y_tr, y_va = y.iloc[tr_idx], y.iloc[val_idx]
        spw_fold = _compute_spw(y_tr)

        print(f'fold x {fold}', len(X_tr), len(X_va))

        fold_scores_msg = []

        # 遍历所有启用的模型
        for m_name in active_models:
            try:
                # 获取参数
                params = get_params(m_name, fold, spw_fold)

                # 执行核心逻辑
                model, preds, auc, path = train_single_model(m_name, X_tr, y_tr, X_va, y_va, params, fold)

                # 记录结果
                oof_preds[m_name][val_idx] = preds
                models_info[m_name].append({'model_path': path, 'auc': auc, 'fold': fold})
                fold_scores_msg.append(f"{m_name.upper()} AUC={auc:.5f}")

            except Exception as e:
                print(f"[Error] Fold {fold} Model {m_name} failed: {e}")

        print(f"[Fold {fold}] " + ", ".join(fold_scores_msg))

    # 所有数据的OOF 评估
    for m_name in active_models:
        total_auc = roc_auc_score(y, oof_preds[m_name])
        print(f"[OOF] {m_name.upper()} AUC = {total_auc:.6f}")

    # 特征重要性 (LGB Only)
    if fi_list:
        fi_all = pd.concat(fi_list, ignore_index=True)
        fi_all = fi_all.groupby("feature", as_index=False)["importance"].mean().sort_values("importance",
                                                                                            ascending=False)
        print("\n[Feature Importance - Top 20]")
        print(fi_all.head(20))

    # 保存元数据
    models_info_path = os.path.join(model_dir, "models_info.pkl")
    with open(models_info_path, 'wb') as f:
        pickle.dump(models_info, f)

    print(f"所有模型已保存到目录: {model_dir}")

    # 返回结果 (根据是否有模型存在，返回全零或预测值)
    return (oof_preds.get('lgb'),
            oof_preds.get('xgb') if 'xgb' in active_models else None,
            oof_preds.get('ctb') if 'ctb' in active_models else None,
            oof_preds.get('rf') if 'rf' in active_models else None,
            models_info)


if __name__ == '__main__':
    # 设置随机种子，保证可复现
    np.random.seed(42)

    # 生成连续特征
    cont1 = np.random.normal(loc=0, scale=1, size=100)
    cont2 = np.random.uniform(low=0, high=10, size=100)

    # 生成离散特征
    cat1 = np.random.choice(['A', 'B', 'C'], size=100)
    cat2 = np.random.choice(['X', 'Y'], size=100)

    # 构造 label（随机二分类）
    label = np.random.choice([0, 1], size=100)

    # 整合成 DataFrame
    df = pd.DataFrame({
        'cont1': cont1,
        'cont2': cont2,
        'cat1': cat1,
        'cat2': cat2,
        'label': label
    })

    le_cat1 = LabelEncoder()
    le_cat2 = LabelEncoder()

    df['cat1'] = le_cat1.fit_transform(df['cat1']).astype(int)
    df['cat2'] = le_cat2.fit_transform(df['cat2']).astype(int)

    print(df.head())
    print(df.shape)
    X, y = df[['cont1', 'cont2', 'cat1', 'cat2']], df['label']
    a, b, c, d, e = train_models_refactoring(X, y, ['cont1', 'cont2', 'cat1', 'cat2'], ['cat1', 'cat2'], 'output')
    print(a, b, c, d, e)
