from __future__ import annotations

import csv
import ctypes
import copy
import gc
import json
import math
import os
import random
import time
from contextlib import nullcontext
from datetime import timedelta
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple

import matplotlib
import numpy as np
import optuna
import pandas as pd
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import autocast, GradScaler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader, DistributedSampler

if os.name != "nt" and not os.environ.get("DISPLAY") and not os.environ.get("MPLBACKEND"):
    matplotlib.use("Agg")
import matplotlib.pyplot as plt
# Optional: unify plotting with shared plotting package
try:
    from ..plotting import curves as plot_curves_common
    from ..plotting.diagnostics import plot_loss_curve as plot_loss_curve_common
except Exception:  # pragma: no cover
    try:
        from ins_pricing.plotting import curves as plot_curves_common
        from ins_pricing.plotting.diagnostics import plot_loss_curve as plot_loss_curve_common
    except Exception:  # pragma: no cover
        plot_curves_common = None
        plot_loss_curve_common = None
# 通过限制内存块拆分让 PyTorch 分配器减少碎片化，缓解显存 OOM
# 如需自定义，可通过环境变量 PYTORCH_CUDA_ALLOC_CONF 覆盖默认设置
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:256")


# 常量与工具模块
# =============================================================================
torch.backends.cudnn.benchmark = True
EPS = 1e-8


def set_global_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


class IOUtils:
    # 文件与路径处理的小工具集合。

    @staticmethod
    def csv_to_dict(file_path: str) -> List[Dict[str, Any]]:
        with open(file_path, mode='r', encoding='utf-8') as file:
            reader = csv.DictReader(file)
            return [
                dict(filter(lambda item: item[0] != '', row.items()))
                for row in reader
            ]

    @staticmethod
    def _sanitize_params_dict(params: Dict[str, Any]) -> Dict[str, Any]:
        # 过滤 pandas 读写中可能出现的 "Unnamed: 0" 等索引列
        return {
            k: v
            for k, v in (params or {}).items()
            if k and not str(k).startswith("Unnamed")
        }

    @staticmethod
    def load_params_file(path: str) -> Dict[str, Any]:
        """从 JSON/CSV/TSV 文件加载参数字典。

        - JSON：支持直接 dict，或形如 {"best_params": {...}} 的包装结构
        - CSV/TSV：读取第一行作为参数
        """
        file_path = Path(path).expanduser().resolve()
        if not file_path.exists():
            raise FileNotFoundError(f"params file not found: {file_path}")
        suffix = file_path.suffix.lower()
        if suffix == ".json":
            payload = json.loads(file_path.read_text(
                encoding="utf-8", errors="replace"))
            if isinstance(payload, dict) and "best_params" in payload:
                payload = payload.get("best_params") or {}
            if not isinstance(payload, dict):
                raise ValueError(
                    f"Invalid JSON params file (expect dict): {file_path}")
            return IOUtils._sanitize_params_dict(dict(payload))
        if suffix in (".csv", ".tsv"):
            df = pd.read_csv(file_path, sep="\t" if suffix == ".tsv" else ",")
            if df.empty:
                raise ValueError(f"Empty params file: {file_path}")
            params = df.iloc[0].to_dict()
            return IOUtils._sanitize_params_dict(params)
        raise ValueError(
            f"Unsupported params file type '{suffix}': {file_path}")

    @staticmethod
    def ensure_parent_dir(file_path: str) -> None:
        # 若目标文件所在目录不存在则自动创建
        directory = os.path.dirname(file_path)
        if directory:
            os.makedirs(directory, exist_ok=True)


class TrainingUtils:
    # 训练阶段常用的小型辅助函数集合。

    @staticmethod
    def compute_batch_size(data_size: int, learning_rate: float, batch_num: int, minimum: int) -> int:
        estimated = int((learning_rate / 1e-4) ** 0.5 *
                        (data_size / max(batch_num, 1)))
        return max(1, min(data_size, max(minimum, estimated)))

    @staticmethod
    def tweedie_loss(pred, target, p=1.5, eps=1e-6, max_clip=1e6):
        # 为确保稳定性先将预测值裁剪为正数
        pred_clamped = torch.clamp(pred, min=eps)
        if p == 1:
            term1 = target * torch.log(target / pred_clamped + eps)  # 泊松
            term2 = -target + pred_clamped
            term3 = 0
        elif p == 0:
            term1 = 0.5 * torch.pow(target - pred_clamped, 2)  # 高斯
            term2 = 0
            term3 = 0
        elif p == 2:
            term1 = torch.log(pred_clamped / target + eps)  # 伽马
            term2 = -target / pred_clamped + 1
            term3 = 0
        else:
            term1 = torch.pow(target, 2 - p) / ((1 - p) * (2 - p))
            term2 = target * torch.pow(pred_clamped, 1 - p) / (1 - p)
            term3 = torch.pow(pred_clamped, 2 - p) / (2 - p)
        return torch.nan_to_num(  # Tweedie 负对数似然（忽略常数项）
            2 * (term1 - term2 + term3),
            nan=eps,
            posinf=max_clip,
            neginf=-max_clip
        )

    @staticmethod
    def free_cuda() -> None:
        print(">>> 将所有模型迁移到 CPU...")
        for obj in gc.get_objects():
            try:
                if hasattr(obj, "to") and callable(obj.to):
                    obj.to("cpu")
            except Exception:
                pass

        print(">>> 删除张量/优化器/DataLoader 引用...")
        gc.collect()

        print(">>> 清空 CUDA cache...")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            print(">>> CUDA 显存已释放。")
        else:
            print(">>> CUDA 不可用，已跳过清理。")


class DistributedUtils:
    _cached_state: Optional[tuple] = None

    @staticmethod
    def setup_ddp():
        """初始化用于分布式训练的 DDP 进程组。"""
        if dist.is_initialized():
            if DistributedUtils._cached_state is None:
                rank = dist.get_rank()
                world_size = dist.get_world_size()
                local_rank = int(os.environ.get("LOCAL_RANK", 0))
                DistributedUtils._cached_state = (
                    True,
                    local_rank,
                    rank,
                    world_size,
                )
            return DistributedUtils._cached_state

        if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
            rank = int(os.environ["RANK"])
            world_size = int(os.environ["WORLD_SIZE"])
            local_rank = int(os.environ["LOCAL_RANK"])

            if os.name == "nt" and torch.cuda.is_available() and world_size > 1:
                print(
                    ">>> DDP Setup Disabled: Windows CUDA DDP is not supported. "
                    "Falling back to single process."
                )
                return False, 0, 0, 1

            if torch.cuda.is_available():
                torch.cuda.set_device(local_rank)

            timeout_seconds = int(os.environ.get(
                "BAYESOPT_DDP_TIMEOUT_SECONDS", "1800"))
            timeout = timedelta(seconds=max(1, timeout_seconds))
            backend = "gloo"
            if torch.cuda.is_available() and os.name != "nt":
                try:
                    if getattr(dist, "is_nccl_available", lambda: False)():
                        backend = "nccl"
                except Exception:
                    backend = "gloo"

            dist.init_process_group(
                backend=backend, init_method="env://", timeout=timeout)
            print(
                f">>> DDP Initialized ({backend}, timeout={timeout_seconds}s): "
                f"Rank {rank}/{world_size}, Local Rank {local_rank}"
            )
            DistributedUtils._cached_state = (
                True,
                local_rank,
                rank,
                world_size,
            )
            return DistributedUtils._cached_state
        else:
            print(
                f">>> DDP Setup Failed: RANK or WORLD_SIZE not found in env. Keys found: {list(os.environ.keys())}"
            )
            print(
                ">>> Hint: launch with torchrun --nproc_per_node=<N> <script.py>"
            )
        return False, 0, 0, 1

    @staticmethod
    def cleanup_ddp():
        """销毁已创建的 DDP 进程组并清理本地缓存状态。"""
        if dist.is_initialized():
            dist.destroy_process_group()
        DistributedUtils._cached_state = None

    @staticmethod
    def is_main_process():
        return not dist.is_initialized() or dist.get_rank() == 0

    @staticmethod
    def world_size() -> int:
        return dist.get_world_size() if dist.is_initialized() else 1


class PlotUtils:
    # 多种模型共享的绘图辅助工具。

    @staticmethod
    def split_data(data: pd.DataFrame, col_nme: str, wgt_nme: str, n_bins: int = 10) -> pd.DataFrame:
        data_sorted = data.sort_values(by=col_nme, ascending=True).copy()
        data_sorted['cum_weight'] = data_sorted[wgt_nme].cumsum()
        w_sum = data_sorted[wgt_nme].sum()
        if w_sum <= EPS:
            data_sorted.loc[:, 'bins'] = 0
        else:
            data_sorted.loc[:, 'bins'] = np.floor(
                data_sorted['cum_weight'] * float(n_bins) / w_sum
            )
            data_sorted.loc[(data_sorted['bins'] == n_bins),
                            'bins'] = n_bins - 1
        return data_sorted.groupby(['bins'], observed=True).sum(numeric_only=True)

    @staticmethod
    def plot_lift_ax(ax, plot_data, title, pred_label='Predicted', act_label='Actual', weight_label='Earned Exposure'):
        ax.plot(plot_data.index, plot_data['act_v'],
                label=act_label, color='red')
        ax.plot(plot_data.index, plot_data['exp_v'],
                label=pred_label, color='blue')
        ax.set_title(title, fontsize=8)
        ax.set_xticks(plot_data.index)
        ax.set_xticklabels(plot_data.index, rotation=90, fontsize=6)
        ax.tick_params(axis='y', labelsize=6)
        ax.legend(loc='upper left', fontsize=5, frameon=False)
        ax.margins(0.05)
        ax2 = ax.twinx()
        ax2.bar(plot_data.index, plot_data['weight'],
                alpha=0.5, color='seagreen',
                label=weight_label)
        ax2.tick_params(axis='y', labelsize=6)
        ax2.legend(loc='upper right', fontsize=5, frameon=False)

    @staticmethod
    def plot_dlift_ax(ax, plot_data, title, label1, label2, act_label='Actual', weight_label='Earned Exposure'):
        ax.plot(plot_data.index, plot_data['act_v'],
                label=act_label, color='red')
        ax.plot(plot_data.index, plot_data['exp_v1'],
                label=label1, color='blue')
        ax.plot(plot_data.index, plot_data['exp_v2'],
                label=label2, color='black')
        ax.set_title(title, fontsize=8)
        ax.set_xticks(plot_data.index)
        ax.set_xticklabels(plot_data.index, rotation=90, fontsize=6)
        ax.set_xlabel(f'{label1} / {label2}', fontsize=6)
        ax.tick_params(axis='y', labelsize=6)
        ax.legend(loc='upper left', fontsize=5, frameon=False)
        ax.margins(0.1)
        ax2 = ax.twinx()
        ax2.bar(plot_data.index, plot_data['weight'],
                alpha=0.5, color='seagreen',
                label=weight_label)
        ax2.tick_params(axis='y', labelsize=6)
        ax2.legend(loc='upper right', fontsize=5, frameon=False)

    @staticmethod
    def plot_lift_list(pred_model, w_pred_list, w_act_list,
                       weight_list, tgt_nme, n_bins: int = 10,
                       fig_nme: str = 'Lift Chart'):
        if plot_curves_common is not None:
            save_path = os.path.join(
                os.getcwd(), 'plot', f'05_{tgt_nme}_{fig_nme}.png')
            plot_curves_common.plot_lift_curve(
                pred_model,
                w_act_list,
                weight_list,
                n_bins=n_bins,
                title=f'Lift Chart of {tgt_nme}',
                pred_label='Predicted',
                act_label='Actual',
                weight_label='Earned Exposure',
                pred_weighted=False,
                actual_weighted=True,
                save_path=save_path,
                show=False,
            )
            return
        lift_data = pd.DataFrame()
        lift_data.loc[:, 'pred'] = pred_model
        lift_data.loc[:, 'w_pred'] = w_pred_list
        lift_data.loc[:, 'act'] = w_act_list
        lift_data.loc[:, 'weight'] = weight_list
        plot_data = PlotUtils.split_data(lift_data, 'pred', 'weight', n_bins)
        plot_data['exp_v'] = plot_data['w_pred'] / plot_data['weight']
        plot_data['act_v'] = plot_data['act'] / plot_data['weight']
        plot_data.reset_index(inplace=True)

        fig = plt.figure(figsize=(7, 5))
        ax = fig.add_subplot(111)
        PlotUtils.plot_lift_ax(ax, plot_data, f'Lift Chart of {tgt_nme}')
        plt.subplots_adjust(wspace=0.3)

        save_path = os.path.join(
            os.getcwd(), 'plot', f'05_{tgt_nme}_{fig_nme}.png')
        IOUtils.ensure_parent_dir(save_path)
        plt.savefig(save_path, dpi=300)
        plt.close(fig)

    @staticmethod
    def plot_dlift_list(pred_model_1, pred_model_2,
                        model_nme_1, model_nme_2,
                        tgt_nme,
                        w_list, w_act_list, n_bins: int = 10,
                        fig_nme: str = 'Double Lift Chart'):
        if plot_curves_common is not None:
            save_path = os.path.join(
                os.getcwd(), 'plot', f'06_{tgt_nme}_{fig_nme}.png')
            plot_curves_common.plot_double_lift_curve(
                pred_model_1,
                pred_model_2,
                w_act_list,
                w_list,
                n_bins=n_bins,
                title=f'Double Lift Chart of {tgt_nme}',
                label1=model_nme_1,
                label2=model_nme_2,
                pred1_weighted=False,
                pred2_weighted=False,
                actual_weighted=True,
                save_path=save_path,
                show=False,
            )
            return
        lift_data = pd.DataFrame()
        lift_data.loc[:, 'pred1'] = pred_model_1
        lift_data.loc[:, 'pred2'] = pred_model_2
        lift_data.loc[:, 'diff_ly'] = lift_data['pred1'] / lift_data['pred2']
        lift_data.loc[:, 'act'] = w_act_list
        lift_data.loc[:, 'weight'] = w_list
        lift_data.loc[:, 'w_pred1'] = lift_data['pred1'] * lift_data['weight']
        lift_data.loc[:, 'w_pred2'] = lift_data['pred2'] * lift_data['weight']
        plot_data = PlotUtils.split_data(
            lift_data, 'diff_ly', 'weight', n_bins)
        plot_data['exp_v1'] = plot_data['w_pred1'] / plot_data['act']
        plot_data['exp_v2'] = plot_data['w_pred2'] / plot_data['act']
        plot_data['act_v'] = plot_data['act']/plot_data['act']
        plot_data.reset_index(inplace=True)

        fig = plt.figure(figsize=(7, 5))
        ax = fig.add_subplot(111)
        PlotUtils.plot_dlift_ax(
            ax, plot_data, f'Double Lift Chart of {tgt_nme}', model_nme_1, model_nme_2)
        plt.subplots_adjust(bottom=0.25, top=0.95, right=0.8)

        save_path = os.path.join(
            os.getcwd(), 'plot', f'06_{tgt_nme}_{fig_nme}.png')
        IOUtils.ensure_parent_dir(save_path)
        plt.savefig(save_path, dpi=300)
        plt.close(fig)


def infer_factor_and_cate_list(train_df: pd.DataFrame,
                               test_df: pd.DataFrame,
                               resp_nme: str,
                               weight_nme: str,
                               binary_resp_nme: Optional[str] = None,
                               factor_nmes: Optional[List[str]] = None,
                               cate_list: Optional[List[str]] = None,
                               infer_categorical_max_unique: int = 50,
                               infer_categorical_max_ratio: float = 0.05) -> Tuple[List[str], List[str]]:
    """在特征列名未知时，推断 factor_nmes/cate_list。

    规则：
      - factor_nmes：默认取 train/test 共有列，排除目标列/权重列/(可选二分类目标列)。
      - cate_list：默认取 object/category/bool，外加“低基数”的整数列。
      - 始终与 train/test 的共有列取交集，避免训练/测试列不一致。
    """
    excluded = {resp_nme, weight_nme}
    if binary_resp_nme:
        excluded.add(binary_resp_nme)

    common_cols = [c for c in train_df.columns if c in test_df.columns]
    if factor_nmes is None:
        factors = [c for c in common_cols if c not in excluded]
    else:
        factors = [
            c for c in factor_nmes if c in common_cols and c not in excluded]

    if cate_list is not None:
        cats = [c for c in cate_list if c in factors]
        return factors, cats

    n_rows = max(1, len(train_df))
    cats: List[str] = []
    for col in factors:
        s = train_df[col]
        if pd.api.types.is_bool_dtype(s) or pd.api.types.is_object_dtype(s) or isinstance(s.dtype, pd.CategoricalDtype):
            cats.append(col)
            continue
        if pd.api.types.is_integer_dtype(s):
            nunique = int(s.nunique(dropna=True))
            if nunique <= infer_categorical_max_unique or (nunique / n_rows) <= infer_categorical_max_ratio:
                cats.append(col)
    return factors, cats


# 向后兼容的函数式封装
def csv_to_dict(file_path: str) -> List[Dict[str, Any]]:
    return IOUtils.csv_to_dict(file_path)


def ensure_parent_dir(file_path: str) -> None:
    IOUtils.ensure_parent_dir(file_path)


def compute_batch_size(data_size: int, learning_rate: float, batch_num: int, minimum: int) -> int:
    return TrainingUtils.compute_batch_size(data_size, learning_rate, batch_num, minimum)


# 定义在 PyTorch 环境下的 Tweedie 偏差损失函数
# 参考文档：https://scikit-learn.org/stable/modules/model_evaluation.html#mean-poisson-gamma-and-tweedie-deviances
def tweedie_loss(pred, target, p=1.5, eps=1e-6, max_clip=1e6):
    return TrainingUtils.tweedie_loss(pred, target, p=p, eps=eps, max_clip=max_clip)


# 定义释放CUDA内存函数
def free_cuda():
    TrainingUtils.free_cuda()


class TorchTrainerMixin:
    # 面向 Torch 表格训练器的共享工具方法。

    def _device_type(self) -> str:
        return getattr(self, "device", torch.device("cpu")).type

    def _resolve_resource_profile(self) -> str:
        profile = getattr(self, "resource_profile", None)
        if not profile:
            profile = os.environ.get("BAYESOPT_RESOURCE_PROFILE", "auto")
        profile = str(profile).strip().lower()
        if profile in {"cpu", "mps", "cuda"}:
            profile = "auto"
        if profile not in {"auto", "throughput", "memory_saving"}:
            profile = "auto"
        if profile == "auto" and self._device_type() == "cuda":
            profile = "throughput"
        return profile

    def _log_resource_summary_once(self, profile: str) -> None:
        if getattr(self, "_resource_summary_logged", False):
            return
        if dist.is_initialized() and not DistributedUtils.is_main_process():
            return
        self._resource_summary_logged = True
        device = getattr(self, "device", torch.device("cpu"))
        device_type = self._device_type()
        cpu_count = os.cpu_count() or 1
        cuda_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
        mps_available = bool(getattr(torch.backends, "mps", None) and torch.backends.mps.is_available())
        ddp_enabled = bool(getattr(self, "is_ddp_enabled", False))
        data_parallel = bool(getattr(self, "use_data_parallel", False))
        print(
            f">>> Resource summary: device={device}, device_type={device_type}, "
            f"cpu_count={cpu_count}, cuda_count={cuda_count}, mps={mps_available}, "
            f"ddp={ddp_enabled}, data_parallel={data_parallel}, profile={profile}"
        )

    def _available_system_memory(self) -> Optional[int]:
        if os.name == "nt":
            class _MemStatus(ctypes.Structure):
                _fields_ = [
                    ("dwLength", ctypes.c_ulong),
                    ("dwMemoryLoad", ctypes.c_ulong),
                    ("ullTotalPhys", ctypes.c_ulonglong),
                    ("ullAvailPhys", ctypes.c_ulonglong),
                    ("ullTotalPageFile", ctypes.c_ulonglong),
                    ("ullAvailPageFile", ctypes.c_ulonglong),
                    ("ullTotalVirtual", ctypes.c_ulonglong),
                    ("ullAvailVirtual", ctypes.c_ulonglong),
                    ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
                ]
            status = _MemStatus()
            status.dwLength = ctypes.sizeof(_MemStatus)
            if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
                return int(status.ullAvailPhys)
            return None
        try:
            pages = os.sysconf("SC_AVPHYS_PAGES")
            page_size = os.sysconf("SC_PAGE_SIZE")
            return int(pages * page_size)
        except Exception:
            return None

    def _available_cuda_memory(self) -> Optional[int]:
        if not torch.cuda.is_available():
            return None
        try:
            free_mem, _total_mem = torch.cuda.mem_get_info()
        except Exception:
            return None
        return int(free_mem)

    def _estimate_sample_bytes(self, dataset) -> Optional[int]:
        try:
            if len(dataset) == 0:
                return None
            sample = dataset[0]
        except Exception:
            return None

        def _bytes(obj) -> int:
            if obj is None:
                return 0
            if torch.is_tensor(obj):
                return int(obj.element_size() * obj.nelement())
            if isinstance(obj, np.ndarray):
                return int(obj.nbytes)
            if isinstance(obj, (list, tuple)):
                return int(sum(_bytes(item) for item in obj))
            if isinstance(obj, dict):
                return int(sum(_bytes(item) for item in obj.values()))
            return 0

        sample_bytes = _bytes(sample)
        return int(sample_bytes) if sample_bytes > 0 else None

    def _cap_batch_size_by_memory(self, dataset, batch_size: int, profile: str) -> int:
        if batch_size <= 1:
            return batch_size
        sample_bytes = self._estimate_sample_bytes(dataset)
        if sample_bytes is None:
            return batch_size
        device_type = self._device_type()
        if device_type == "cuda":
            available = self._available_cuda_memory()
            if available is None:
                return batch_size
            if profile == "throughput":
                budget_ratio = 0.8
                overhead = 8.0
            elif profile == "memory_saving":
                budget_ratio = 0.5
                overhead = 14.0
            else:
                budget_ratio = 0.6
                overhead = 12.0
        else:
            available = self._available_system_memory()
            if available is None:
                return batch_size
            if profile == "throughput":
                budget_ratio = 0.4
                overhead = 1.8
            elif profile == "memory_saving":
                budget_ratio = 0.25
                overhead = 3.0
            else:
                budget_ratio = 0.3
                overhead = 2.6
        budget = int(available * budget_ratio)
        per_sample = int(sample_bytes * overhead)
        if per_sample <= 0:
            return batch_size
        max_batch = max(1, int(budget // per_sample))
        if max_batch < batch_size:
            print(
                f">>> Memory cap: batch_size {batch_size} -> {max_batch} "
                f"(per_sample~{sample_bytes}B, budget~{budget // (1024**2)}MB)"
            )
        return min(batch_size, max_batch)

    def _resolve_num_workers(self, max_workers: int, profile: Optional[str] = None) -> int:
        if os.name == 'nt':
            return 0
        if getattr(self, "is_ddp_enabled", False):
            return 0
        profile = profile or self._resolve_resource_profile()
        if profile == "memory_saving":
            return 0
        worker_cap = min(int(max_workers), os.cpu_count() or 1)
        if self._device_type() == "mps":
            worker_cap = min(worker_cap, 2)
        return worker_cap

    def _build_dataloader(self,
                          dataset,
                          N: int,
                          base_bs_gpu: tuple,
                          base_bs_cpu: tuple,
                          min_bs: int = 64,
                          target_effective_cuda: int = 1024,
                          target_effective_cpu: int = 512,
                          large_threshold: int = 200_000,
                          mid_threshold: int = 50_000):
        profile = self._resolve_resource_profile()
        self._log_resource_summary_once(profile)
        batch_size = TrainingUtils.compute_batch_size(
            data_size=len(dataset),
            learning_rate=self.learning_rate,
            batch_num=self.batch_num,
            minimum=min_bs
        )
        gpu_large, gpu_mid, gpu_small = base_bs_gpu
        cpu_mid, cpu_small = base_bs_cpu

        if self._device_type() == 'cuda':
            device_count = torch.cuda.device_count()
            if getattr(self, "is_ddp_enabled", False):
                device_count = 1
            # 多卡环境下，适当增大最小批量，确保每张卡都能分到足够数据
            if device_count > 1:
                min_bs = min_bs * device_count
                print(
                    f">>> Multi-GPU detected: {device_count} devices. Adjusted min_bs to {min_bs}.")

            if N > large_threshold:
                base_bs = gpu_large * device_count
            elif N > mid_threshold:
                base_bs = gpu_mid * device_count
            else:
                base_bs = gpu_small * device_count
        else:
            base_bs = cpu_mid if N > mid_threshold else cpu_small

        # 重新计算 batch_size，确保不小于调整后的 min_bs
        batch_size = TrainingUtils.compute_batch_size(
            data_size=len(dataset),
            learning_rate=self.learning_rate,
            batch_num=self.batch_num,
            minimum=min_bs
        )
        batch_size = min(batch_size, base_bs, N)
        batch_size = self._cap_batch_size_by_memory(
            dataset, batch_size, profile)

        target_effective_bs = target_effective_cuda if self._device_type(
        ) == 'cuda' else target_effective_cpu
        if getattr(self, "is_ddp_enabled", False):
            world_size = max(1, DistributedUtils.world_size())
            target_effective_bs = max(1, target_effective_bs // world_size)

        world_size = getattr(self, "world_size", 1) if getattr(
            self, "is_ddp_enabled", False) else 1
        samples_per_rank = math.ceil(
            N / max(1, world_size)) if world_size > 1 else N
        steps_per_epoch = max(
            1, math.ceil(samples_per_rank / max(1, batch_size)))
        # 限制梯度累积步数：避免用超过实际 batch 数的系数去缩放 loss。
        desired_accum = max(1, target_effective_bs // max(1, batch_size))
        accum_steps = max(1, min(desired_accum, steps_per_epoch))

        # Linux (posix) 采用 fork 更高效；Windows (nt) 使用 spawn，开销更大。
        workers = self._resolve_num_workers(8, profile=profile)
        prefetch_factor = None
        if workers > 0:
            prefetch_factor = 4 if profile == "throughput" else 2
        persistent = workers > 0 and profile != "memory_saving"
        print(
            f">>> DataLoader config: Batch Size={batch_size}, Accum Steps={accum_steps}, "
            f"Workers={workers}, Prefetch={prefetch_factor or 'off'}, Profile={profile}")
        sampler = None
        if dist.is_initialized():
            sampler = DistributedSampler(dataset, shuffle=True)
            shuffle = False  # 交由 DistributedSampler 完成随机打乱
        else:
            shuffle = True

        dataloader = DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=shuffle,
            sampler=sampler,
            num_workers=workers,
            pin_memory=(self._device_type() == 'cuda'),
            persistent_workers=persistent,
            **({"prefetch_factor": prefetch_factor} if prefetch_factor is not None else {}),
        )
        return dataloader, accum_steps

    def _build_val_dataloader(self, dataset, train_dataloader, accum_steps):
        profile = self._resolve_resource_profile()
        val_bs = accum_steps * train_dataloader.batch_size
        val_workers = self._resolve_num_workers(4, profile=profile)
        prefetch_factor = None
        if val_workers > 0:
            prefetch_factor = 2
        return DataLoader(
            dataset,
            batch_size=val_bs,
            shuffle=False,
            num_workers=val_workers,
            pin_memory=(self._device_type() == 'cuda'),
            persistent_workers=(val_workers > 0 and profile != "memory_saving"),
            **({"prefetch_factor": prefetch_factor} if prefetch_factor is not None else {}),
        )

    def _compute_losses(self, y_pred, y_true, apply_softplus: bool = False):
        task = getattr(self, "task_type", "regression")
        if task == 'classification':
            loss_fn = nn.BCEWithLogitsLoss(reduction='none')
            return loss_fn(y_pred, y_true).view(-1)
        if apply_softplus:
            y_pred = F.softplus(y_pred)
        y_pred = torch.clamp(y_pred, min=1e-6)
        power = getattr(self, "tw_power", 1.5)
        return tweedie_loss(y_pred, y_true, p=power).view(-1)

    def _compute_weighted_loss(self, y_pred, y_true, weights, apply_softplus: bool = False):
        losses = self._compute_losses(
            y_pred, y_true, apply_softplus=apply_softplus)
        weighted_loss = (losses * weights.view(-1)).sum() / \
            torch.clamp(weights.sum(), min=EPS)
        return weighted_loss

    def _early_stop_update(self, val_loss, best_loss, best_state, patience_counter, model,
                           ignore_keys: Optional[List[str]] = None):
        if val_loss < best_loss:
            ignore_keys = ignore_keys or []
            state_dict = {
                k: (v.clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v))
                for k, v in model.state_dict().items()
                if not any(k.startswith(ignore_key) for ignore_key in ignore_keys)
            }
            return val_loss, state_dict, 0, False
        patience_counter += 1
        should_stop = best_state is not None and patience_counter >= getattr(
            self, "patience", 0)
        return best_loss, best_state, patience_counter, should_stop

    def _train_model(self,
                     model,
                     dataloader,
                     accum_steps,
                     optimizer,
                     scaler,
                     forward_fn,
                     val_forward_fn=None,
                     apply_softplus: bool = False,
                     clip_fn=None,
                     trial: Optional[optuna.trial.Trial] = None,
                     loss_curve_path: Optional[str] = None):
        device_type = self._device_type()
        best_loss = float('inf')
        best_state = None
        patience_counter = 0
        stop_training = False
        train_history: List[float] = []
        val_history: List[float] = []

        is_ddp_model = isinstance(model, DDP)

        for epoch in range(1, getattr(self, "epochs", 1) + 1):
            epoch_start_ts = time.time()
            val_weighted_loss = None
            if hasattr(self, 'dataloader_sampler') and self.dataloader_sampler is not None:
                self.dataloader_sampler.set_epoch(epoch)

            model.train()
            optimizer.zero_grad()

            epoch_loss_sum = None
            epoch_weight_sum = None
            for step, batch in enumerate(dataloader):
                is_update_step = ((step + 1) % accum_steps == 0) or \
                    ((step + 1) == len(dataloader))
                sync_cm = model.no_sync if (
                    is_ddp_model and not is_update_step) else nullcontext

                with sync_cm():
                    with autocast(enabled=(device_type == 'cuda')):
                        y_pred, y_true, w = forward_fn(batch)
                        weighted_loss = self._compute_weighted_loss(
                            y_pred, y_true, w, apply_softplus=apply_softplus)
                        loss_for_backward = weighted_loss / accum_steps

                    batch_weight = torch.clamp(
                        w.detach().sum(), min=EPS).to(dtype=torch.float32)
                    loss_val = weighted_loss.detach().to(dtype=torch.float32)
                    if epoch_loss_sum is None:
                        epoch_loss_sum = torch.zeros(
                            (), device=batch_weight.device, dtype=torch.float32)
                        epoch_weight_sum = torch.zeros(
                            (), device=batch_weight.device, dtype=torch.float32)
                    epoch_loss_sum = epoch_loss_sum + loss_val * batch_weight
                    epoch_weight_sum = epoch_weight_sum + batch_weight
                    scaler.scale(loss_for_backward).backward()

                if is_update_step:
                    if clip_fn is not None:
                        clip_fn()
                    scaler.step(optimizer)
                    scaler.update()
                    optimizer.zero_grad()

            if epoch_loss_sum is None or epoch_weight_sum is None:
                train_epoch_loss = 0.0
            else:
                train_epoch_loss = (
                    epoch_loss_sum / torch.clamp(epoch_weight_sum, min=EPS)
                ).item()
            train_history.append(float(train_epoch_loss))

            if val_forward_fn is not None:
                should_compute_val = (not dist.is_initialized()
                                      or DistributedUtils.is_main_process())
                val_device = getattr(self, "device", torch.device("cpu"))
                if not isinstance(val_device, torch.device):
                    val_device = torch.device(val_device)
                loss_tensor_device = val_device if device_type == 'cuda' else torch.device(
                    "cpu")
                val_loss_tensor = torch.zeros(1, device=loss_tensor_device)

                if should_compute_val:
                    model.eval()
                    with torch.no_grad(), autocast(enabled=(device_type == 'cuda')):
                        val_result = val_forward_fn()
                        if isinstance(val_result, tuple) and len(val_result) == 3:
                            y_val_pred, y_val_true, w_val = val_result
                            val_weighted_loss = self._compute_weighted_loss(
                                y_val_pred, y_val_true, w_val, apply_softplus=apply_softplus)
                        else:
                            val_weighted_loss = val_result
                    val_loss_tensor[0] = float(val_weighted_loss)

                if dist.is_initialized():
                    dist.broadcast(val_loss_tensor, src=0)
                val_weighted_loss = float(val_loss_tensor.item())

                val_history.append(val_weighted_loss)

                best_loss, best_state, patience_counter, stop_training = self._early_stop_update(
                    val_weighted_loss, best_loss, best_state, patience_counter, model)

                prune_flag = False
                is_main_rank = DistributedUtils.is_main_process()
                if trial is not None and (not dist.is_initialized() or is_main_rank):
                    trial.report(val_weighted_loss, epoch)
                    prune_flag = trial.should_prune()

                if dist.is_initialized():
                    prune_device = getattr(self, "device", torch.device("cpu"))
                    if not isinstance(prune_device, torch.device):
                        prune_device = torch.device(prune_device)
                    prune_tensor = torch.zeros(1, device=prune_device)
                    if is_main_rank:
                        prune_tensor.fill_(1 if prune_flag else 0)
                    dist.broadcast(prune_tensor, src=0)
                    prune_flag = bool(prune_tensor.item())

                if prune_flag:
                    raise optuna.TrialPruned()

                if stop_training:
                    break

            should_log_epoch = (not dist.is_initialized()
                                or DistributedUtils.is_main_process())
            if should_log_epoch:
                elapsed = int(time.time() - epoch_start_ts)
                if val_weighted_loss is None:
                    print(
                        f"[Training] Epoch {epoch}/{getattr(self, 'epochs', 1)} "
                        f"train_loss={float(train_epoch_loss):.6f} elapsed={elapsed}s",
                        flush=True,
                    )
                else:
                    print(
                        f"[Training] Epoch {epoch}/{getattr(self, 'epochs', 1)} "
                        f"train_loss={float(train_epoch_loss):.6f} "
                        f"val_loss={float(val_weighted_loss):.6f} elapsed={elapsed}s",
                        flush=True,
                    )

        history = {"train": train_history, "val": val_history}
        self._plot_loss_curve(history, loss_curve_path)
        return best_state, history

    def _plot_loss_curve(self, history: Dict[str, List[float]], save_path: Optional[str]) -> None:
        if not save_path:
            return
        if dist.is_initialized() and not DistributedUtils.is_main_process():
            return
        train_hist = history.get("train", []) if history else []
        val_hist = history.get("val", []) if history else []
        if not train_hist and not val_hist:
            return
        if plot_loss_curve_common is not None:
            plot_loss_curve_common(
                history=history,
                title="Loss vs. Epoch",
                save_path=save_path,
                show=False,
            )
        else:
            ensure_parent_dir(save_path)
            epochs = range(1, max(len(train_hist), len(val_hist)) + 1)
            fig = plt.figure(figsize=(8, 4))
            ax = fig.add_subplot(111)
            if train_hist:
                ax.plot(range(1, len(train_hist) + 1), train_hist,
                        label='Train Loss', color='tab:blue')
            if val_hist:
                ax.plot(range(1, len(val_hist) + 1), val_hist,
                        label='Validation Loss', color='tab:orange')
            ax.set_xlabel('Epoch')
            ax.set_ylabel('Weighted Loss')
            ax.set_title('Loss vs. Epoch')
            ax.grid(True, linestyle='--', alpha=0.3)
            ax.legend()
            plt.tight_layout()
            plt.savefig(save_path, dpi=300)
            plt.close(fig)
        print(f"[Training] Loss curve saved to {save_path}")


# =============================================================================
# 绘图辅助模块
# =============================================================================

def split_data(data, col_nme, wgt_nme, n_bins=10):
    return PlotUtils.split_data(data, col_nme, wgt_nme, n_bins)

# 定义提纯曲线（Lift）绘制函数


def plot_lift_list(pred_model, w_pred_list, w_act_list,
                   weight_list, tgt_nme, n_bins=10,
                   fig_nme='Lift Chart'):
    return PlotUtils.plot_lift_list(pred_model, w_pred_list, w_act_list,
                                    weight_list, tgt_nme, n_bins, fig_nme)

# 定义双提纯曲线绘制函数


def plot_dlift_list(pred_model_1, pred_model_2,
                    model_nme_1, model_nme_2,
                    tgt_nme,
                    w_list, w_act_list, n_bins=10,
                    fig_nme='Double Lift Chart'):
    return PlotUtils.plot_dlift_list(pred_model_1, pred_model_2,
                                     model_nme_1, model_nme_2,
                                     tgt_nme, w_list, w_act_list,
                                     n_bins, fig_nme)
