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

try:  # matplotlib is optional; avoid hard import failures in headless/minimal envs
    import matplotlib
    if os.name != "nt" and not os.environ.get("DISPLAY") and not os.environ.get("MPLBACKEND"):
        matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    _MPL_IMPORT_ERROR: Optional[BaseException] = None
except Exception as exc:  # pragma: no cover - optional dependency
    matplotlib = None  # type: ignore[assignment]
    plt = None  # type: ignore[assignment]
    _MPL_IMPORT_ERROR = exc
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

# 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
# Limit CUDA allocator split size to reduce fragmentation and OOM risk.
# Override via PYTORCH_CUDA_ALLOC_CONF if needed.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:256")


# Constants and utility helpers
# =============================================================================
torch.backends.cudnn.benchmark = True
EPS = 1e-8


def _plot_skip(label: str) -> None:
    if _MPL_IMPORT_ERROR is not None:
        print(f"[Plot] Skip {label}: matplotlib unavailable ({_MPL_IMPORT_ERROR}).", flush=True)
    else:
        print(f"[Plot] Skip {label}: matplotlib unavailable.", flush=True)


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:
    # File and path utilities.

    @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]:
        # Filter index-like columns such as "Unnamed: 0" from pandas I/O.
        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]:
        """Load parameter dict from JSON/CSV/TSV files.

        - JSON: accept dict or {"best_params": {...}} wrapper
        - CSV/TSV: read the first row as params
        """
        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:
        # Create parent directories when missing.
        directory = os.path.dirname(file_path)
        if directory:
            os.makedirs(directory, exist_ok=True)


class TrainingUtils:
    # Small helpers used during training.

    @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):
        # Clamp predictions to positive values for stability.
        pred_clamped = torch.clamp(pred, min=eps)
        if p == 1:
            term1 = target * torch.log(target / pred_clamped + eps)  # Poisson
            term2 = -target + pred_clamped
            term3 = 0
        elif p == 0:
            term1 = 0.5 * torch.pow(target - pred_clamped, 2)  # Gaussian
            term2 = 0
            term3 = 0
        elif p == 2:
            term1 = torch.log(pred_clamped / target + eps)  # Gamma
            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 negative log-likelihood (constant omitted)
            2 * (term1 - term2 + term3),
            nan=eps,
            posinf=max_clip,
            neginf=-max_clip
        )

    @staticmethod
    def free_cuda() -> None:
        print(">>> Moving all models to CPU...")
        for obj in gc.get_objects():
            try:
                if hasattr(obj, "to") and callable(obj.to):
                    obj.to("cpu")
            except Exception:
                pass

        print(">>> Releasing tensor/optimizer/DataLoader references...")
        gc.collect()

        print(">>> Clearing CUDA cache...")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            print(">>> CUDA memory released.")
        else:
            print(">>> CUDA not available; cleanup skipped.")


class DistributedUtils:
    _cached_state: Optional[tuple] = None

    @staticmethod
    def setup_ddp():
        """Initialize the DDP process group for distributed training."""
        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():
        """Destroy the DDP process group and clear cached state."""
        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:
    # Plot helpers shared across models.

    @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
        if plt is None:
            _plot_skip("lift plot")
            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
        if plt is None:
            _plot_skip("double lift plot")
            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]]:
    """Infer factor_nmes/cate_list when feature names are not provided.

    Rules:
      - factor_nmes: start from shared train/test columns, exclude target/weight/(optional binary target).
      - cate_list: object/category/bool plus low-cardinality integer columns.
      - Always intersect with shared train/test columns to avoid mismatches.
    """
    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


# Backward-compatible functional wrappers
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)


# Tweedie deviance loss for PyTorch.
# Reference: 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 memory release helper.
def free_cuda():
    TrainingUtils.free_cuda()


class TorchTrainerMixin:
    # Shared helpers for Torch tabular trainers.

    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
            # In multi-GPU, increase min batch size so each GPU gets enough data.
            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

        # Recompute batch_size to respect the adjusted 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)))
        # Limit gradient accumulation to avoid scaling beyond actual batches.
        desired_accum = max(1, target_effective_bs // max(1, batch_size))
        accum_steps = max(1, min(desired_accum, steps_per_epoch))

        # Linux (posix) uses fork; Windows (nt) uses spawn with higher overhead.
        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 handles shuffling.
        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:
            if plt is None:
                _plot_skip("loss curve")
                return
            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}")


# =============================================================================
# Plotting helpers
# =============================================================================

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

# Lift curve plotting wrapper


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)

# Double lift curve plotting wrapper


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)
