from __future__ import annotations

import copy
import hashlib
import math
import os
import time
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

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 sklearn.neighbors import NearestNeighbors
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 Dataset, TensorDataset

from .utils import DistributedUtils, EPS, IOUtils, TorchTrainerMixin

try:
    from torch_geometric.nn import knn_graph
    from torch_geometric.utils import add_self_loops, to_undirected
    _PYG_AVAILABLE = True
except Exception:
    knn_graph = None  # type: ignore
    add_self_loops = None  # type: ignore
    to_undirected = None  # type: ignore
    _PYG_AVAILABLE = False

try:
    import pynndescent
    _PYNN_AVAILABLE = True
except Exception:
    pynndescent = None  # type: ignore
    _PYNN_AVAILABLE = False

_GNN_MPS_WARNED = False

# =============================================================================
# ResNet model and sklearn-style wrapper
# =============================================================================

# ResNet model definition
# Residual block: two linear layers + ReLU + residual connection
# ResBlock inherits nn.Module
class ResBlock(nn.Module):
    def __init__(self, dim: int, dropout: float = 0.1,
                 use_layernorm: bool = False, residual_scale: float = 0.1,
                 stochastic_depth: float = 0.0
                 ):
        super().__init__()
        self.use_layernorm = use_layernorm

        if use_layernorm:
            Norm = nn.LayerNorm      # Normalize the last dimension
        else:
            def Norm(d): return nn.BatchNorm1d(d)  # Keep a switch to try BN

        self.norm1 = Norm(dim)
        self.fc1 = nn.Linear(dim, dim, bias=True)
        self.act = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
        # Enable post-second-layer norm if needed: self.norm2 = Norm(dim)
        self.fc2 = nn.Linear(dim, dim, bias=True)

        # Residual scaling to stabilize early training
        self.res_scale = nn.Parameter(
            torch.tensor(residual_scale, dtype=torch.float32)
        )
        self.stochastic_depth = max(0.0, float(stochastic_depth))

    def _drop_path(self, x: torch.Tensor) -> torch.Tensor:
        if self.stochastic_depth <= 0.0 or not self.training:
            return x
        keep_prob = 1.0 - self.stochastic_depth
        if keep_prob <= 0.0:
            return torch.zeros_like(x)
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = keep_prob + torch.rand(
            shape, dtype=x.dtype, device=x.device)
        binary_tensor = torch.floor(random_tensor)
        return x * binary_tensor / keep_prob

    def forward(self, x):
        # Pre-activation structure
        out = self.norm1(x)
        out = self.fc1(out)
        out = self.act(out)
        out = self.dropout(out)
        # If a second norm is enabled: out = self.norm2(out)
        out = self.fc2(out)
        # Apply residual scaling then add
        out = self.res_scale * out
        out = self._drop_path(out)
        return x + out

# ResNetSequential defines the full network


class ResNetSequential(nn.Module):
    # Input shape: (batch, input_dim)
    # Network: FC + norm + ReLU, stack residual blocks, output Softplus

    def __init__(self, input_dim: int, hidden_dim: int = 64, block_num: int = 2,
                 use_layernorm: bool = True, dropout: float = 0.1,
                 residual_scale: float = 0.1, stochastic_depth: float = 0.0,
                 task_type: str = 'regression'):
        super(ResNetSequential, self).__init__()

        self.net = nn.Sequential()
        self.net.add_module('fc1', nn.Linear(input_dim, hidden_dim))

        # Optional explicit normalization after the first layer:
        # For LayerNorm:
        #     self.net.add_module('norm1', nn.LayerNorm(hidden_dim))
        # Or BatchNorm:
        #     self.net.add_module('norm1', nn.BatchNorm1d(hidden_dim))

        # If desired, insert ReLU before residual blocks:
        # self.net.add_module('relu1', nn.ReLU(inplace=True))

        # Residual blocks
        drop_path_rate = max(0.0, float(stochastic_depth))
        for i in range(block_num):
            if block_num > 1:
                block_drop = drop_path_rate * (i / (block_num - 1))
            else:
                block_drop = drop_path_rate
            self.net.add_module(
                f'ResBlk_{i+1}',
                ResBlock(
                    hidden_dim,
                    dropout=dropout,
                    use_layernorm=use_layernorm,
                    residual_scale=residual_scale,
                    stochastic_depth=block_drop)
            )

        self.net.add_module('fc_out', nn.Linear(hidden_dim, 1))

        if task_type == 'classification':
            self.net.add_module('softplus', nn.Identity())
        else:
            self.net.add_module('softplus', nn.Softplus())

    def forward(self, x):
        if self.training and not hasattr(self, '_printed_device'):
            print(f">>> ResNetSequential executing on device: {x.device}")
            self._printed_device = True
        return self.net(x)

# Define the ResNet sklearn-style wrapper.


class ResNetSklearn(TorchTrainerMixin, nn.Module):
    def __init__(self, model_nme: str, input_dim: int, hidden_dim: int = 64,
                 block_num: int = 2, batch_num: int = 100, epochs: int = 100,
                 task_type: str = 'regression',
                 tweedie_power: float = 1.5, learning_rate: float = 0.01, patience: int = 10,
                 use_layernorm: bool = True, dropout: float = 0.1,
                 residual_scale: float = 0.1,
                 stochastic_depth: float = 0.0,
                 weight_decay: float = 1e-4,
                 use_data_parallel: bool = True,
                 use_ddp: bool = False):
        super(ResNetSklearn, self).__init__()

        self.use_ddp = use_ddp
        self.is_ddp_enabled, self.local_rank, self.rank, self.world_size = (
            False, 0, 0, 1)

        if self.use_ddp:
            self.is_ddp_enabled, self.local_rank, self.rank, self.world_size = DistributedUtils.setup_ddp()

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.block_num = block_num
        self.batch_num = batch_num
        self.epochs = epochs
        self.task_type = task_type
        self.model_nme = model_nme
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.patience = patience
        self.use_layernorm = use_layernorm
        self.dropout = dropout
        self.residual_scale = residual_scale
        self.stochastic_depth = max(0.0, float(stochastic_depth))
        self.loss_curve_path: Optional[str] = None
        self.training_history: Dict[str, List[float]] = {
            "train": [], "val": []}
        self.use_data_parallel = bool(use_data_parallel)

        # Device selection: cuda > mps > cpu
        if self.is_ddp_enabled:
            self.device = torch.device(f'cuda:{self.local_rank}')
        elif torch.cuda.is_available():
            self.device = torch.device('cuda')
        elif torch.backends.mps.is_available():
            self.device = torch.device('mps')
        else:
            self.device = torch.device('cpu')

        # Tweedie power (unused for classification)
        if self.task_type == 'classification':
            self.tw_power = None
        elif 'f' in self.model_nme:
            self.tw_power = 1
        elif 's' in self.model_nme:
            self.tw_power = 2
        else:
            self.tw_power = tweedie_power

        # Build network (construct on CPU first)
        core = ResNetSequential(
            self.input_dim,
            self.hidden_dim,
            self.block_num,
            use_layernorm=self.use_layernorm,
            dropout=self.dropout,
            residual_scale=self.residual_scale,
            stochastic_depth=self.stochastic_depth,
            task_type=self.task_type
        )

        # ===== Multi-GPU: DataParallel vs DistributedDataParallel =====
        if self.is_ddp_enabled:
            core = core.to(self.device)
            core = DDP(core, device_ids=[
                       self.local_rank], output_device=self.local_rank)
            self.use_data_parallel = False
        elif use_data_parallel and (self.device.type == 'cuda') and (torch.cuda.device_count() > 1):
            if self.use_ddp and not self.is_ddp_enabled:
                print(
                    ">>> DDP requested but not initialized; falling back to DataParallel.")
            core = nn.DataParallel(core, device_ids=list(
                range(torch.cuda.device_count())))
            # DataParallel scatters inputs, but the primary device remains cuda:0.
            self.device = torch.device('cuda')
            self.use_data_parallel = True
        else:
            self.use_data_parallel = False

        self.resnet = core.to(self.device)

    # ================ Internal helpers ================
    @staticmethod
    def _validate_vector(arr, name: str, n_rows: int) -> None:
        if arr is None:
            return
        if isinstance(arr, pd.DataFrame):
            if arr.shape[1] != 1:
                raise ValueError(f"{name} must be 1d (single column).")
            length = len(arr)
        else:
            arr_np = np.asarray(arr)
            if arr_np.ndim == 0:
                raise ValueError(f"{name} must be 1d.")
            if arr_np.ndim > 2 or (arr_np.ndim == 2 and arr_np.shape[1] != 1):
                raise ValueError(f"{name} must be 1d or Nx1.")
            length = arr_np.shape[0]
        if length != n_rows:
            raise ValueError(
                f"{name} length {length} does not match X length {n_rows}."
            )

    def _validate_inputs(self, X, y, w, label: str) -> None:
        if X is None:
            raise ValueError(f"{label} X cannot be None.")
        n_rows = len(X)
        if y is None:
            raise ValueError(f"{label} y cannot be None.")
        self._validate_vector(y, f"{label} y", n_rows)
        self._validate_vector(w, f"{label} w", n_rows)

    def _build_train_val_tensors(self, X_train, y_train, w_train, X_val, y_val, w_val):
        self._validate_inputs(X_train, y_train, w_train, "train")
        if X_val is not None or y_val is not None or w_val is not None:
            if X_val is None or y_val is None:
                raise ValueError("validation X and y must both be provided.")
            self._validate_inputs(X_val, y_val, w_val, "val")

        def _to_numpy(arr):
            if hasattr(arr, "to_numpy"):
                return arr.to_numpy(dtype=np.float32, copy=False)
            return np.asarray(arr, dtype=np.float32)

        X_tensor = torch.as_tensor(_to_numpy(X_train))
        y_tensor = torch.as_tensor(_to_numpy(y_train)).view(-1, 1)
        w_tensor = (
            torch.as_tensor(_to_numpy(w_train)).view(-1, 1)
            if w_train is not None else torch.ones_like(y_tensor)
        )

        has_val = X_val is not None and y_val is not None
        if has_val:
            X_val_tensor = torch.as_tensor(_to_numpy(X_val))
            y_val_tensor = torch.as_tensor(_to_numpy(y_val)).view(-1, 1)
            w_val_tensor = (
                torch.as_tensor(_to_numpy(w_val)).view(-1, 1)
                if w_val is not None else torch.ones_like(y_val_tensor)
            )
        else:
            X_val_tensor = y_val_tensor = w_val_tensor = None
        return X_tensor, y_tensor, w_tensor, X_val_tensor, y_val_tensor, w_val_tensor, has_val

    def forward(self, x):
        # Handle SHAP NumPy input.
        if isinstance(x, np.ndarray):
            x_tensor = torch.as_tensor(x, dtype=torch.float32)
        else:
            x_tensor = x

        x_tensor = x_tensor.to(self.device)
        y_pred = self.resnet(x_tensor)
        return y_pred

    # ---------------- Training ----------------

    def fit(self, X_train, y_train, w_train=None,
            X_val=None, y_val=None, w_val=None, trial=None):

        X_tensor, y_tensor, w_tensor, X_val_tensor, y_val_tensor, w_val_tensor, has_val = \
            self._build_train_val_tensors(
                X_train, y_train, w_train, X_val, y_val, w_val)

        dataset = TensorDataset(X_tensor, y_tensor, w_tensor)
        dataloader, accum_steps = self._build_dataloader(
            dataset,
            N=X_tensor.shape[0],
            base_bs_gpu=(2048, 1024, 512),
            base_bs_cpu=(256, 128),
            min_bs=64,
            target_effective_cuda=2048,
            target_effective_cpu=1024
        )

        # Set sampler epoch at the start of each epoch to keep shuffling deterministic.
        if self.is_ddp_enabled and hasattr(dataloader.sampler, 'set_epoch'):
            self.dataloader_sampler = dataloader.sampler
        else:
            self.dataloader_sampler = None

        # === 4. Optimizer and AMP ===
        self.optimizer = torch.optim.Adam(
            self.resnet.parameters(),
            lr=self.learning_rate,
            weight_decay=float(self.weight_decay),
        )
        self.scaler = GradScaler(enabled=(self.device.type == 'cuda'))

        X_val_dev = y_val_dev = w_val_dev = None
        val_dataloader = None
        if has_val:
            # Build validation DataLoader.
            val_dataset = TensorDataset(
                X_val_tensor, y_val_tensor, w_val_tensor)
            # No backward pass in validation; batch size can be larger for throughput.
            val_dataloader = self._build_val_dataloader(
                val_dataset, dataloader, accum_steps)
            # Validation usually does not need a DDP sampler because we validate on the main process
            # or aggregate results. For simplicity, keep validation on a single GPU or the main process.

        is_data_parallel = isinstance(self.resnet, nn.DataParallel)

        def forward_fn(batch):
            X_batch, y_batch, w_batch = batch

            if not is_data_parallel:
                X_batch = X_batch.to(self.device, non_blocking=True)
            # Keep targets and weights on the main device for loss computation.
            y_batch = y_batch.to(self.device, non_blocking=True)
            w_batch = w_batch.to(self.device, non_blocking=True)

            y_pred = self.resnet(X_batch)
            return y_pred, y_batch, w_batch

        def val_forward_fn():
            total_loss = 0.0
            total_weight = 0.0
            for batch in val_dataloader:
                X_b, y_b, w_b = batch
                if not is_data_parallel:
                    X_b = X_b.to(self.device, non_blocking=True)
                y_b = y_b.to(self.device, non_blocking=True)
                w_b = w_b.to(self.device, non_blocking=True)

                y_pred = self.resnet(X_b)

                # Manually compute weighted loss for accurate aggregation.
                losses = self._compute_losses(
                    y_pred, y_b, apply_softplus=False)

                batch_weight_sum = torch.clamp(w_b.sum(), min=EPS)
                batch_weighted_loss_sum = (losses * w_b.view(-1)).sum()

                total_loss += batch_weighted_loss_sum.item()
                total_weight += batch_weight_sum.item()

            return total_loss / max(total_weight, EPS)

        clip_fn = None
        if self.device.type == 'cuda':
            def clip_fn(): return (self.scaler.unscale_(self.optimizer),
                                   clip_grad_norm_(self.resnet.parameters(), max_norm=1.0))

        # Under DDP, only the main process prints logs and saves models.
        if self.is_ddp_enabled and not DistributedUtils.is_main_process():
            # Non-main processes skip validation callback logging (handled inside _train_model).
            pass

        best_state, history = self._train_model(
            self.resnet,
            dataloader,
            accum_steps,
            self.optimizer,
            self.scaler,
            forward_fn,
            val_forward_fn if has_val else None,
            apply_softplus=False,
            clip_fn=clip_fn,
            trial=trial,
            loss_curve_path=getattr(self, "loss_curve_path", None)
        )

        if has_val and best_state is not None:
            self.resnet.load_state_dict(best_state)
        self.training_history = history

    # ---------------- Prediction ----------------

    def predict(self, X_test):
        self.resnet.eval()
        if isinstance(X_test, pd.DataFrame):
            X_np = X_test.to_numpy(dtype=np.float32, copy=False)
        else:
            X_np = np.asarray(X_test, dtype=np.float32)

        inference_cm = getattr(torch, "inference_mode", torch.no_grad)
        with inference_cm():
            y_pred = self(X_np).cpu().numpy()

        if self.task_type == 'classification':
            y_pred = 1 / (1 + np.exp(-y_pred))  # Sigmoid converts logits to probabilities.
        else:
            y_pred = np.clip(y_pred, 1e-6, None)
        return y_pred.flatten()

    # ---------------- Set Params ----------------

    def set_params(self, params):
        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Parameter {key} not found in model.")
        return self


# =============================================================================
# FT-Transformer model and sklearn-style wrapper.
# =============================================================================
# Define FT-Transformer model structure.


class FeatureTokenizer(nn.Module):
    """Map numeric/categorical/geo tokens into transformer input tokens."""

    def __init__(
        self,
        num_numeric: int,
        cat_cardinalities,
        d_model: int,
        num_geo: int = 0,
        num_numeric_tokens: int = 1,
    ):
        super().__init__()

        self.num_numeric = num_numeric
        self.num_geo = num_geo
        self.has_geo = num_geo > 0

        if num_numeric > 0:
            if int(num_numeric_tokens) <= 0:
                raise ValueError("num_numeric_tokens must be >= 1 when numeric features exist.")
            self.num_numeric_tokens = int(num_numeric_tokens)
            self.has_numeric = True
            self.num_linear = nn.Linear(num_numeric, d_model * self.num_numeric_tokens)
        else:
            self.num_numeric_tokens = 0
            self.has_numeric = False

        self.embeddings = nn.ModuleList([
            nn.Embedding(card, d_model) for card in cat_cardinalities
        ])

        if self.has_geo:
            # Map geo tokens with a linear layer to avoid one-hot on raw strings; upstream is encoded/normalized.
            self.geo_linear = nn.Linear(num_geo, d_model)

    def forward(self, X_num, X_cat, X_geo=None):
        tokens = []

        if self.has_numeric:
            batch_size = X_num.shape[0]
            num_token = self.num_linear(X_num)
            num_token = num_token.view(batch_size, self.num_numeric_tokens, -1)
            tokens.append(num_token)

        for i, emb in enumerate(self.embeddings):
            tok = emb(X_cat[:, i])
            tokens.append(tok.unsqueeze(1))

        if self.has_geo:
            if X_geo is None:
                raise RuntimeError("Geo tokens are enabled but X_geo was not provided.")
            geo_token = self.geo_linear(X_geo)
            tokens.append(geo_token.unsqueeze(1))

        x = torch.cat(tokens, dim=1)
        return x

# Encoder layer with residual scaling.


class ScaledTransformerEncoderLayer(nn.Module):
    def __init__(self, d_model: int, nhead: int, dim_feedforward: int = 2048,
                 dropout: float = 0.1, residual_scale_attn: float = 1.0,
                 residual_scale_ffn: float = 1.0, norm_first: bool = True,
                 ):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=nhead,
            dropout=dropout,
            batch_first=True
        )

        # Feed-forward network.
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        # Normalization and dropout.
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = nn.GELU()
        # If you prefer ReLU, set: self.activation = nn.ReLU()
        self.norm_first = norm_first

        # Residual scaling coefficients.
        self.res_scale_attn = residual_scale_attn
        self.res_scale_ffn = residual_scale_ffn

    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # Input tensor shape: (batch, seq_len, d_model).
        x = src

        if self.norm_first:
            # Pre-norm before attention.
            x = x + self._sa_block(self.norm1(x), src_mask,
                                   src_key_padding_mask)
            x = x + self._ff_block(self.norm2(x))
        else:
            # Post-norm (usually disabled).
            x = self.norm1(
                x + self._sa_block(x, src_mask, src_key_padding_mask))
            x = self.norm2(x + self._ff_block(x))

        return x

    def _sa_block(self, x, attn_mask, key_padding_mask):
        # Self-attention with residual scaling.
        attn_out, _ = self.self_attn(
            x, x, x,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=False
        )
        return self.res_scale_attn * self.dropout1(attn_out)

    def _ff_block(self, x):
        # Feed-forward block with residual scaling.
        x2 = self.linear2(self.dropout(self.activation(self.linear1(x))))
        return self.res_scale_ffn * self.dropout2(x2)

# FT-Transformer core model.


class FTTransformerCore(nn.Module):
    # Minimal FT-Transformer built from:
    #   1) FeatureTokenizer: convert numeric/categorical features to tokens;
    #   2) TransformerEncoder: model feature interactions;
    #   3) Pooling + MLP + Softplus: positive outputs for Tweedie/Gamma tasks.

    def __init__(self, num_numeric: int, cat_cardinalities, d_model: int = 64,
                 n_heads: int = 8, n_layers: int = 4, dropout: float = 0.1,
                 task_type: str = 'regression', num_geo: int = 0,
                 num_numeric_tokens: int = 1
                 ):
        super().__init__()

        self.num_numeric = int(num_numeric)
        self.cat_cardinalities = list(cat_cardinalities or [])

        self.tokenizer = FeatureTokenizer(
            num_numeric=num_numeric,
            cat_cardinalities=cat_cardinalities,
            d_model=d_model,
            num_geo=num_geo,
            num_numeric_tokens=num_numeric_tokens
        )
        scale = 1.0 / math.sqrt(n_layers)  # Recommended default.
        encoder_layer = ScaledTransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=dropout,
            residual_scale_attn=scale,
            residual_scale_ffn=scale,
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=n_layers
        )
        self.n_layers = n_layers

        layers = [
            # If you need a deeper head, enable the sample layers below:
            # nn.LayerNorm(d_model),  # Extra normalization
            # nn.Linear(d_model, d_model),  # Extra fully connected layer
            # nn.GELU(),  # Activation
            nn.Linear(d_model, 1),
        ]

        if task_type == 'classification':
            # Classification outputs logits for BCEWithLogitsLoss.
            layers.append(nn.Identity())
        else:
            # Regression keeps positive outputs for Tweedie/Gamma.
            layers.append(nn.Softplus())

        self.head = nn.Sequential(*layers)

        # ---- Self-supervised reconstruction head (masked modeling) ----
        self.num_recon_head = nn.Linear(
            d_model, self.num_numeric) if self.num_numeric > 0 else None
        self.cat_recon_heads = nn.ModuleList([
            nn.Linear(d_model, int(card)) for card in self.cat_cardinalities
        ])

    def forward(
            self,
            X_num,
            X_cat,
            X_geo=None,
            return_embedding: bool = False,
            return_reconstruction: bool = False):

        # Inputs:
        #   X_num -> float32 tensor with shape (batch, num_numeric_features)
        #   X_cat -> long tensor with shape (batch, num_categorical_features)
        #   X_geo -> float32 tensor with shape (batch, geo_token_dim)

        if self.training and not hasattr(self, '_printed_device'):
            print(f">>> FTTransformerCore executing on device: {X_num.device}")
            self._printed_device = True

        # => tensor shape (batch, token_num, d_model)
        tokens = self.tokenizer(X_num, X_cat, X_geo)
        # => tensor shape (batch, token_num, d_model)
        x = self.encoder(tokens)

        # Mean-pool tokens, then send to the head.
        x = x.mean(dim=1)                      # => tensor shape (batch, d_model)

        if return_reconstruction:
            num_pred, cat_logits = self.reconstruct(x)
            cat_logits_out = tuple(
                cat_logits) if cat_logits is not None else tuple()
            if return_embedding:
                return x, num_pred, cat_logits_out
            return num_pred, cat_logits_out

        if return_embedding:
            return x

        # => tensor shape (batch, 1); Softplus keeps it positive.
        out = self.head(x)
        return out

    def reconstruct(self, embedding: torch.Tensor) -> Tuple[Optional[torch.Tensor], List[torch.Tensor]]:
        """Reconstruct numeric/categorical inputs from pooled embedding (batch, d_model)."""
        num_pred = self.num_recon_head(
            embedding) if self.num_recon_head is not None else None
        cat_logits = [head(embedding) for head in self.cat_recon_heads]
        return num_pred, cat_logits

# TabularDataset.


class TabularDataset(Dataset):
    def __init__(self, X_num, X_cat, X_geo, y, w):

        # Input tensors:
        #   X_num: torch.float32, shape=(N, num_numeric_features)
        #   X_cat: torch.long,   shape=(N, num_categorical_features)
        #   X_geo: torch.float32, shape=(N, geo_token_dim), can be empty
        #   y:     torch.float32, shape=(N, 1)
        #   w:     torch.float32, shape=(N, 1)

        self.X_num = X_num
        self.X_cat = X_cat
        self.X_geo = X_geo
        self.y = y
        self.w = w

    def __len__(self):
        return self.y.shape[0]

    def __getitem__(self, idx):
        return (
            self.X_num[idx],
            self.X_cat[idx],
            self.X_geo[idx],
            self.y[idx],
            self.w[idx],
        )


class MaskedTabularDataset(Dataset):
    def __init__(self,
                 X_num_masked: torch.Tensor,
                 X_cat_masked: torch.Tensor,
                 X_geo: torch.Tensor,
                 X_num_true: Optional[torch.Tensor],
                 num_mask: Optional[torch.Tensor],
                 X_cat_true: Optional[torch.Tensor],
                 cat_mask: Optional[torch.Tensor]):
        self.X_num_masked = X_num_masked
        self.X_cat_masked = X_cat_masked
        self.X_geo = X_geo
        self.X_num_true = X_num_true
        self.num_mask = num_mask
        self.X_cat_true = X_cat_true
        self.cat_mask = cat_mask

    def __len__(self):
        return self.X_num_masked.shape[0]

    def __getitem__(self, idx):
        return (
            self.X_num_masked[idx],
            self.X_cat_masked[idx],
            self.X_geo[idx],
            None if self.X_num_true is None else self.X_num_true[idx],
            None if self.num_mask is None else self.num_mask[idx],
            None if self.X_cat_true is None else self.X_cat_true[idx],
            None if self.cat_mask is None else self.cat_mask[idx],
        )

# Scikit-Learn style wrapper for FTTransformer.


class FTTransformerSklearn(TorchTrainerMixin, nn.Module):

    # sklearn-style wrapper:
    #   - num_cols: numeric feature column names
    #   - cat_cols: categorical feature column names (label-encoded to [0, n_classes-1])

    @staticmethod
    def resolve_numeric_token_count(num_cols, cat_cols, requested: Optional[int]) -> int:
        num_cols_count = len(num_cols or [])
        if num_cols_count == 0:
            return 0
        if requested is not None:
            count = int(requested)
            if count <= 0:
                raise ValueError("num_numeric_tokens must be >= 1 when numeric features exist.")
            return count
        return max(1, num_cols_count)

    def __init__(self, model_nme: str, num_cols, cat_cols, d_model: int = 64, n_heads: int = 8,
                 n_layers: int = 4, dropout: float = 0.1, batch_num: int = 100, epochs: int = 100,
                 task_type: str = 'regression',
                 tweedie_power: float = 1.5, learning_rate: float = 1e-3, patience: int = 10,
                 weight_decay: float = 0.0,
                 use_data_parallel: bool = True,
                 use_ddp: bool = False,
                 num_numeric_tokens: Optional[int] = None
                 ):
        super().__init__()

        self.use_ddp = use_ddp
        self.is_ddp_enabled, self.local_rank, self.rank, self.world_size = (
            False, 0, 0, 1)
        if self.use_ddp:
            self.is_ddp_enabled, self.local_rank, self.rank, self.world_size = DistributedUtils.setup_ddp()

        self.model_nme = model_nme
        self.num_cols = list(num_cols)
        self.cat_cols = list(cat_cols)
        self.num_numeric_tokens = self.resolve_numeric_token_count(
            self.num_cols,
            self.cat_cols,
            num_numeric_tokens,
        )
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.dropout = dropout
        self.batch_num = batch_num
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.task_type = task_type
        self.patience = patience
        if self.task_type == 'classification':
            self.tw_power = None  # No Tweedie power for classification.
        elif 'f' in self.model_nme:
            self.tw_power = 1.0
        elif 's' in self.model_nme:
            self.tw_power = 2.0
        else:
            self.tw_power = tweedie_power

        if self.is_ddp_enabled:
            self.device = torch.device(f"cuda:{self.local_rank}")
        elif torch.cuda.is_available():
            self.device = torch.device("cuda")
        elif torch.backends.mps.is_available():
            self.device = torch.device("mps")
        else:
            self.device = torch.device("cpu")
        self.cat_cardinalities = None
        self.cat_categories = {}
        self.cat_maps: Dict[str, Dict[Any, int]] = {}
        self.cat_str_maps: Dict[str, Dict[str, int]] = {}
        self._num_mean = None
        self._num_std = None
        self.ft = None
        self.use_data_parallel = bool(use_data_parallel)
        self.num_geo = 0
        self._geo_params: Dict[str, Any] = {}
        self.loss_curve_path: Optional[str] = None
        self.training_history: Dict[str, List[float]] = {
            "train": [], "val": []}

    def _build_model(self, X_train):
        num_numeric = len(self.num_cols)
        cat_cardinalities = []

        if num_numeric > 0:
            num_arr = X_train[self.num_cols].to_numpy(
                dtype=np.float32, copy=False)
            num_arr = np.nan_to_num(num_arr, nan=0.0, posinf=0.0, neginf=0.0)
            mean = num_arr.mean(axis=0).astype(np.float32, copy=False)
            std = num_arr.std(axis=0).astype(np.float32, copy=False)
            std = np.where(std < 1e-6, 1.0, std).astype(np.float32, copy=False)
            self._num_mean = mean
            self._num_std = std
        else:
            self._num_mean = None
            self._num_std = None

        self.cat_maps = {}
        self.cat_str_maps = {}
        for col in self.cat_cols:
            cats = X_train[col].astype('category')
            categories = cats.cat.categories
            self.cat_categories[col] = categories           # Store full category list from training.
            self.cat_maps[col] = {cat: i for i, cat in enumerate(categories)}
            if categories.dtype == object or pd.api.types.is_string_dtype(categories.dtype):
                self.cat_str_maps[col] = {str(cat): i for i, cat in enumerate(categories)}

            card = len(categories) + 1                      # Reserve one extra class for unknown/missing.
            cat_cardinalities.append(card)

        self.cat_cardinalities = cat_cardinalities

        core = FTTransformerCore(
            num_numeric=num_numeric,
            cat_cardinalities=cat_cardinalities,
            d_model=self.d_model,
            n_heads=self.n_heads,
            n_layers=self.n_layers,
            dropout=self.dropout,
            task_type=self.task_type,
            num_geo=self.num_geo,
            num_numeric_tokens=self.num_numeric_tokens
        )
        use_dp = self.use_data_parallel and (self.device.type == "cuda") and (torch.cuda.device_count() > 1)
        if self.is_ddp_enabled:
            core = core.to(self.device)
            core = DDP(core, device_ids=[
                       self.local_rank], output_device=self.local_rank, find_unused_parameters=True)
            self.use_data_parallel = False
        elif use_dp:
            if self.use_ddp and not self.is_ddp_enabled:
                print(
                    ">>> DDP requested but not initialized; falling back to DataParallel.")
            core = nn.DataParallel(core, device_ids=list(
                range(torch.cuda.device_count())))
            self.device = torch.device("cuda")
            self.use_data_parallel = True
        else:
            self.use_data_parallel = False
        self.ft = core.to(self.device)

    def _encode_cats(self, X):
        # Input DataFrame must include all categorical feature columns.
        # Return int64 array with shape (N, num_categorical_features).

        if not self.cat_cols:
            return np.zeros((len(X), 0), dtype='int64')

        n_rows = len(X)
        n_cols = len(self.cat_cols)
        X_cat_np = np.empty((n_rows, n_cols), dtype='int64')
        for idx, col in enumerate(self.cat_cols):
            categories = self.cat_categories[col]
            mapping = self.cat_maps.get(col)
            if mapping is None:
                mapping = {cat: i for i, cat in enumerate(categories)}
                self.cat_maps[col] = mapping
            unknown_idx = len(categories)
            series = X[col]
            codes = series.map(mapping)
            unmapped = series.notna() & codes.isna()
            if unmapped.any():
                try:
                    series_cast = series.astype(categories.dtype)
                except Exception:
                    series_cast = None
                if series_cast is not None:
                    codes = series_cast.map(mapping)
                    unmapped = series_cast.notna() & codes.isna()
            if unmapped.any():
                str_map = self.cat_str_maps.get(col)
                if str_map is None:
                    str_map = {str(cat): i for i, cat in enumerate(categories)}
                    self.cat_str_maps[col] = str_map
                codes = series.astype(str).map(str_map)
            if pd.api.types.is_categorical_dtype(codes):
                codes = codes.astype("float")
            codes = codes.fillna(unknown_idx).astype(
                "int64", copy=False).to_numpy()
            X_cat_np[:, idx] = codes
        return X_cat_np

    def _build_train_tensors(self, X_train, y_train, w_train, geo_train=None):
        return self._tensorize_split(X_train, y_train, w_train, geo_tokens=geo_train)

    def _build_val_tensors(self, X_val, y_val, w_val, geo_val=None):
        return self._tensorize_split(X_val, y_val, w_val, geo_tokens=geo_val, allow_none=True)

    @staticmethod
    def _validate_vector(arr, name: str, n_rows: int) -> None:
        if arr is None:
            return
        if isinstance(arr, pd.DataFrame):
            if arr.shape[1] != 1:
                raise ValueError(f"{name} must be 1d (single column).")
            length = len(arr)
        else:
            arr_np = np.asarray(arr)
            if arr_np.ndim == 0:
                raise ValueError(f"{name} must be 1d.")
            if arr_np.ndim > 2 or (arr_np.ndim == 2 and arr_np.shape[1] != 1):
                raise ValueError(f"{name} must be 1d or Nx1.")
            length = arr_np.shape[0]
        if length != n_rows:
            raise ValueError(
                f"{name} length {length} does not match X length {n_rows}."
            )

    def _tensorize_split(self, X, y, w, geo_tokens=None, allow_none: bool = False):
        if X is None:
            if allow_none:
                return None, None, None, None, None, False
            raise ValueError("Input features X must not be None.")
        if not isinstance(X, pd.DataFrame):
            raise ValueError("X must be a pandas DataFrame.")
        missing_cols = [
            col for col in (self.num_cols + self.cat_cols) if col not in X.columns
        ]
        if missing_cols:
            raise ValueError(f"X is missing required columns: {missing_cols}")
        n_rows = len(X)
        if y is not None:
            self._validate_vector(y, "y", n_rows)
        if w is not None:
            self._validate_vector(w, "w", n_rows)

        num_np = X[self.num_cols].to_numpy(dtype=np.float32, copy=False)
        if not num_np.flags["OWNDATA"]:
            num_np = num_np.copy()
        num_np = np.nan_to_num(num_np, nan=0.0,
                               posinf=0.0, neginf=0.0, copy=False)
        if self._num_mean is not None and self._num_std is not None and num_np.size:
            num_np = (num_np - self._num_mean) / self._num_std
        X_num = torch.as_tensor(num_np)
        if self.cat_cols:
            X_cat = torch.as_tensor(self._encode_cats(X), dtype=torch.long)
        else:
            X_cat = torch.zeros((X_num.shape[0], 0), dtype=torch.long)

        if geo_tokens is not None:
            geo_np = np.asarray(geo_tokens, dtype=np.float32)
            if geo_np.shape[0] != n_rows:
                raise ValueError(
                    "geo_tokens length does not match X rows.")
            if geo_np.ndim == 1:
                geo_np = geo_np.reshape(-1, 1)
        elif self.num_geo > 0:
            raise RuntimeError("geo_tokens must not be empty; prepare geo tokens first.")
        else:
            geo_np = np.zeros((X_num.shape[0], 0), dtype=np.float32)
        X_geo = torch.as_tensor(geo_np)

        y_tensor = torch.as_tensor(
            y.to_numpy(dtype=np.float32, copy=False) if hasattr(
                y, "to_numpy") else np.asarray(y, dtype=np.float32)
        ).view(-1, 1) if y is not None else None
        if y_tensor is None:
            w_tensor = None
        elif w is not None:
            w_tensor = torch.as_tensor(
                w.to_numpy(dtype=np.float32, copy=False) if hasattr(
                    w, "to_numpy") else np.asarray(w, dtype=np.float32)
            ).view(-1, 1)
        else:
            w_tensor = torch.ones_like(y_tensor)
        return X_num, X_cat, X_geo, y_tensor, w_tensor, y is not None

    def fit(self, X_train, y_train, w_train=None,
            X_val=None, y_val=None, w_val=None, trial=None,
            geo_train=None, geo_val=None):

        # Build the underlying model on first fit.
        self.num_geo = geo_train.shape[1] if geo_train is not None else 0
        if self.ft is None:
            self._build_model(X_train)

        X_num_train, X_cat_train, X_geo_train, y_tensor, w_tensor, _ = self._build_train_tensors(
            X_train, y_train, w_train, geo_train=geo_train)
        X_num_val, X_cat_val, X_geo_val, y_val_tensor, w_val_tensor, has_val = self._build_val_tensors(
            X_val, y_val, w_val, geo_val=geo_val)

        # --- Build DataLoader ---
        dataset = TabularDataset(
            X_num_train, X_cat_train, X_geo_train, y_tensor, w_tensor
        )

        dataloader, accum_steps = self._build_dataloader(
            dataset,
            N=X_num_train.shape[0],
            base_bs_gpu=(2048, 1024, 512),
            base_bs_cpu=(256, 128),
            min_bs=64,
            target_effective_cuda=2048,
            target_effective_cpu=1024
        )

        if self.is_ddp_enabled and hasattr(dataloader.sampler, 'set_epoch'):
            self.dataloader_sampler = dataloader.sampler
        else:
            self.dataloader_sampler = None

        optimizer = torch.optim.Adam(
            self.ft.parameters(),
            lr=self.learning_rate,
            weight_decay=float(getattr(self, "weight_decay", 0.0)),
        )
        scaler = GradScaler(enabled=(self.device.type == 'cuda'))

        X_num_val_dev = X_cat_val_dev = y_val_dev = w_val_dev = None
        val_dataloader = None
        if has_val:
            val_dataset = TabularDataset(
                X_num_val, X_cat_val, X_geo_val, y_val_tensor, w_val_tensor
            )
            val_dataloader = self._build_val_dataloader(
                val_dataset, dataloader, accum_steps)

        is_data_parallel = isinstance(self.ft, nn.DataParallel)

        def forward_fn(batch):
            X_num_b, X_cat_b, X_geo_b, y_b, w_b = batch

            if not is_data_parallel:
                X_num_b = X_num_b.to(self.device, non_blocking=True)
                X_cat_b = X_cat_b.to(self.device, non_blocking=True)
                X_geo_b = X_geo_b.to(self.device, non_blocking=True)
            y_b = y_b.to(self.device, non_blocking=True)
            w_b = w_b.to(self.device, non_blocking=True)

            y_pred = self.ft(X_num_b, X_cat_b, X_geo_b)
            return y_pred, y_b, w_b

        def val_forward_fn():
            total_loss = 0.0
            total_weight = 0.0
            for batch in val_dataloader:
                X_num_b, X_cat_b, X_geo_b, y_b, w_b = batch
                if not is_data_parallel:
                    X_num_b = X_num_b.to(self.device, non_blocking=True)
                    X_cat_b = X_cat_b.to(self.device, non_blocking=True)
                    X_geo_b = X_geo_b.to(self.device, non_blocking=True)
                y_b = y_b.to(self.device, non_blocking=True)
                w_b = w_b.to(self.device, non_blocking=True)

                y_pred = self.ft(X_num_b, X_cat_b, X_geo_b)

                # Manually compute validation loss.
                losses = self._compute_losses(
                    y_pred, y_b, apply_softplus=False)

                batch_weight_sum = torch.clamp(w_b.sum(), min=EPS)
                batch_weighted_loss_sum = (losses * w_b.view(-1)).sum()

                total_loss += batch_weighted_loss_sum.item()
                total_weight += batch_weight_sum.item()

            return total_loss / max(total_weight, EPS)

        clip_fn = None
        if self.device.type == 'cuda':
            def clip_fn(): return (scaler.unscale_(optimizer),
                                   clip_grad_norm_(self.ft.parameters(), max_norm=1.0))

        best_state, history = self._train_model(
            self.ft,
            dataloader,
            accum_steps,
            optimizer,
            scaler,
            forward_fn,
            val_forward_fn if has_val else None,
            apply_softplus=False,
            clip_fn=clip_fn,
            trial=trial,
            loss_curve_path=getattr(self, "loss_curve_path", None)
        )

        if has_val and best_state is not None:
            self.ft.load_state_dict(best_state)
        self.training_history = history

    def fit_unsupervised(self,
                         X_train,
                         X_val=None,
                         trial: Optional[optuna.trial.Trial] = None,
                         geo_train=None,
                         geo_val=None,
                         mask_prob_num: float = 0.15,
                         mask_prob_cat: float = 0.15,
                         num_loss_weight: float = 1.0,
                         cat_loss_weight: float = 1.0) -> float:
        """Self-supervised pretraining via masked reconstruction (supports raw string categories)."""
        self.num_geo = geo_train.shape[1] if geo_train is not None else 0
        if self.ft is None:
            self._build_model(X_train)

        X_num, X_cat, X_geo, _, _, _ = self._tensorize_split(
            X_train, None, None, geo_tokens=geo_train, allow_none=True)
        has_val = X_val is not None
        if has_val:
            X_num_val, X_cat_val, X_geo_val, _, _, _ = self._tensorize_split(
                X_val, None, None, geo_tokens=geo_val, allow_none=True)
        else:
            X_num_val = X_cat_val = X_geo_val = None

        N = int(X_num.shape[0])
        num_dim = int(X_num.shape[1])
        cat_dim = int(X_cat.shape[1])
        device_type = self._device_type()

        gen = torch.Generator()
        gen.manual_seed(13 + int(getattr(self, "rank", 0)))

        base_model = self.ft.module if hasattr(self.ft, "module") else self.ft
        cardinals = getattr(base_model, "cat_cardinalities", None) or []
        unknown_idx = torch.tensor(
            [int(c) - 1 for c in cardinals], dtype=torch.long).view(1, -1)

        means = None
        if num_dim > 0:
            # Keep masked fill values on the same scale as model inputs (may be normalized in _tensorize_split).
            means = X_num.to(dtype=torch.float32).mean(dim=0, keepdim=True)

        def _mask_inputs(X_num_in: torch.Tensor,
                         X_cat_in: torch.Tensor,
                         generator: torch.Generator):
            n_rows = int(X_num_in.shape[0])
            num_mask_local = None
            cat_mask_local = None
            X_num_masked_local = X_num_in
            X_cat_masked_local = X_cat_in
            if num_dim > 0:
                num_mask_local = (torch.rand(
                    (n_rows, num_dim), generator=generator) < float(mask_prob_num))
                X_num_masked_local = X_num_in.clone()
                if num_mask_local.any():
                    X_num_masked_local[num_mask_local] = means.expand_as(
                        X_num_masked_local)[num_mask_local]
            if cat_dim > 0:
                cat_mask_local = (torch.rand(
                    (n_rows, cat_dim), generator=generator) < float(mask_prob_cat))
                X_cat_masked_local = X_cat_in.clone()
                if cat_mask_local.any():
                    X_cat_masked_local[cat_mask_local] = unknown_idx.expand_as(
                        X_cat_masked_local)[cat_mask_local]
            return X_num_masked_local, X_cat_masked_local, num_mask_local, cat_mask_local

        X_num_true = X_num if num_dim > 0 else None
        X_cat_true = X_cat if cat_dim > 0 else None
        X_num_masked, X_cat_masked, num_mask, cat_mask = _mask_inputs(
            X_num, X_cat, gen)

        dataset = MaskedTabularDataset(
            X_num_masked, X_cat_masked, X_geo,
            X_num_true, num_mask,
            X_cat_true, cat_mask
        )
        dataloader, accum_steps = self._build_dataloader(
            dataset,
            N=N,
            base_bs_gpu=(2048, 1024, 512),
            base_bs_cpu=(256, 128),
            min_bs=64,
            target_effective_cuda=2048,
            target_effective_cpu=1024
        )
        if self.is_ddp_enabled and hasattr(dataloader.sampler, 'set_epoch'):
            self.dataloader_sampler = dataloader.sampler
        else:
            self.dataloader_sampler = None

        optimizer = torch.optim.Adam(
            self.ft.parameters(),
            lr=self.learning_rate,
            weight_decay=float(getattr(self, "weight_decay", 0.0)),
        )
        scaler = GradScaler(enabled=(device_type == 'cuda'))

        def _batch_recon_loss(num_pred, cat_logits, num_true_b, num_mask_b, cat_true_b, cat_mask_b, device):
            loss = torch.zeros((), device=device, dtype=torch.float32)

            if num_pred is not None and num_true_b is not None and num_mask_b is not None:
                num_mask_b = num_mask_b.to(dtype=torch.bool)
                if num_mask_b.any():
                    diff = num_pred - num_true_b
                    mse = diff * diff
                    loss = loss + float(num_loss_weight) * \
                        mse[num_mask_b].mean()

            if cat_logits and cat_true_b is not None and cat_mask_b is not None:
                cat_mask_b = cat_mask_b.to(dtype=torch.bool)
                cat_losses: List[torch.Tensor] = []
                for j, logits in enumerate(cat_logits):
                    mask_j = cat_mask_b[:, j]
                    if not mask_j.any():
                        continue
                    targets = cat_true_b[:, j]
                    cat_losses.append(
                        F.cross_entropy(logits, targets, reduction='none')[
                            mask_j].mean()
                    )
                if cat_losses:
                    loss = loss + float(cat_loss_weight) * \
                        torch.stack(cat_losses).mean()
            return loss

        train_history: List[float] = []
        val_history: List[float] = []
        best_loss = float("inf")
        best_state = None
        patience_counter = 0
        is_ddp_model = isinstance(self.ft, DDP)

        clip_fn = None
        if self.device.type == 'cuda':
            def clip_fn(): return (scaler.unscale_(optimizer),
                                   clip_grad_norm_(self.ft.parameters(), max_norm=1.0))

        for epoch in range(1, int(self.epochs) + 1):
            if self.dataloader_sampler is not None:
                self.dataloader_sampler.set_epoch(epoch)

            self.ft.train()
            optimizer.zero_grad()
            epoch_loss_sum = 0.0
            epoch_count = 0.0

            for step, batch in enumerate(dataloader):
                is_update_step = ((step + 1) % accum_steps == 0) or \
                    ((step + 1) == len(dataloader))
                sync_cm = self.ft.no_sync if (
                    is_ddp_model and not is_update_step) else nullcontext
                with sync_cm():
                    with autocast(enabled=(device_type == 'cuda')):
                        X_num_b, X_cat_b, X_geo_b, num_true_b, num_mask_b, cat_true_b, cat_mask_b = batch
                        X_num_b = X_num_b.to(self.device, non_blocking=True)
                        X_cat_b = X_cat_b.to(self.device, non_blocking=True)
                        X_geo_b = X_geo_b.to(self.device, non_blocking=True)
                        num_true_b = None if num_true_b is None else num_true_b.to(
                            self.device, non_blocking=True)
                        num_mask_b = None if num_mask_b is None else num_mask_b.to(
                            self.device, non_blocking=True)
                        cat_true_b = None if cat_true_b is None else cat_true_b.to(
                            self.device, non_blocking=True)
                        cat_mask_b = None if cat_mask_b is None else cat_mask_b.to(
                            self.device, non_blocking=True)

                        num_pred, cat_logits = self.ft(
                            X_num_b, X_cat_b, X_geo_b, return_reconstruction=True)
                        batch_loss = _batch_recon_loss(
                            num_pred, cat_logits, num_true_b, num_mask_b, cat_true_b, cat_mask_b, device=X_num_b.device)
                        local_bad = 0 if bool(torch.isfinite(batch_loss)) else 1
                        global_bad = local_bad
                        if dist.is_initialized():
                            bad = torch.tensor(
                                [local_bad],
                                device=batch_loss.device,
                                dtype=torch.int32,
                            )
                            dist.all_reduce(bad, op=dist.ReduceOp.MAX)
                            global_bad = int(bad.item())

                        if global_bad:
                            msg = (
                                f"[FTTransformerSklearn.fit_unsupervised] non-finite loss "
                                f"(epoch={epoch}, step={step}, loss={batch_loss.detach().item()})"
                            )
                            should_log = (not dist.is_initialized()
                                          or DistributedUtils.is_main_process())
                            if should_log:
                                print(msg, flush=True)
                                print(
                                    f"  X_num: finite={bool(torch.isfinite(X_num_b).all())} "
                                    f"min={float(X_num_b.min().detach().cpu()) if X_num_b.numel() else 0.0:.3g} "
                                    f"max={float(X_num_b.max().detach().cpu()) if X_num_b.numel() else 0.0:.3g}",
                                    flush=True,
                                )
                                if X_geo_b is not None:
                                    print(
                                        f"  X_geo: finite={bool(torch.isfinite(X_geo_b).all())} "
                                        f"min={float(X_geo_b.min().detach().cpu()) if X_geo_b.numel() else 0.0:.3g} "
                                        f"max={float(X_geo_b.max().detach().cpu()) if X_geo_b.numel() else 0.0:.3g}",
                                        flush=True,
                                    )
                            if trial is not None:
                                raise optuna.TrialPruned(msg)
                            raise RuntimeError(msg)
                        loss_for_backward = batch_loss / float(accum_steps)
                    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()

                epoch_loss_sum += float(batch_loss.detach().item()) * \
                    float(X_num_b.shape[0])
                epoch_count += float(X_num_b.shape[0])

            train_history.append(epoch_loss_sum / max(epoch_count, 1.0))

            if has_val and X_num_val is not None and X_cat_val is not None and X_geo_val is not None:
                should_compute_val = (not dist.is_initialized()
                                      or DistributedUtils.is_main_process())
                loss_tensor_device = self.device if device_type == 'cuda' else torch.device(
                    "cpu")
                val_loss_tensor = torch.zeros(1, device=loss_tensor_device)

                if should_compute_val:
                    self.ft.eval()
                    with torch.no_grad(), autocast(enabled=(device_type == 'cuda')):
                        val_bs = min(
                            int(dataloader.batch_size * max(1, accum_steps)), int(X_num_val.shape[0]))
                        total_val = 0.0
                        total_n = 0.0
                        for start in range(0, int(X_num_val.shape[0]), max(1, val_bs)):
                            end = min(
                                int(X_num_val.shape[0]), start + max(1, val_bs))
                            X_num_v_true_cpu = X_num_val[start:end]
                            X_cat_v_true_cpu = X_cat_val[start:end]
                            X_geo_v = X_geo_val[start:end].to(
                                self.device, non_blocking=True)
                            gen_val = torch.Generator()
                            gen_val.manual_seed(10_000 + epoch + start)
                            X_num_v_cpu, X_cat_v_cpu, val_num_mask, val_cat_mask = _mask_inputs(
                                X_num_v_true_cpu, X_cat_v_true_cpu, gen_val)
                            X_num_v_true = X_num_v_true_cpu.to(
                                self.device, non_blocking=True)
                            X_cat_v_true = X_cat_v_true_cpu.to(
                                self.device, non_blocking=True)
                            X_num_v = X_num_v_cpu.to(
                                self.device, non_blocking=True)
                            X_cat_v = X_cat_v_cpu.to(
                                self.device, non_blocking=True)
                            val_num_mask = None if val_num_mask is None else val_num_mask.to(
                                self.device, non_blocking=True)
                            val_cat_mask = None if val_cat_mask is None else val_cat_mask.to(
                                self.device, non_blocking=True)
                            num_pred_v, cat_logits_v = self.ft(
                                X_num_v, X_cat_v, X_geo_v, return_reconstruction=True)
                            loss_v = _batch_recon_loss(
                                num_pred_v, cat_logits_v,
                                X_num_v_true if X_num_v_true.numel() else None, val_num_mask,
                                X_cat_v_true if X_cat_v_true.numel() else None, val_cat_mask,
                                device=X_num_v.device
                            )
                            if not torch.isfinite(loss_v):
                                total_val = float("inf")
                                total_n = 1.0
                                break
                            total_val += float(loss_v.detach().item()
                                               ) * float(end - start)
                            total_n += float(end - start)
                    val_loss_tensor[0] = total_val / max(total_n, 1.0)

                if dist.is_initialized():
                    dist.broadcast(val_loss_tensor, src=0)
                val_loss_value = float(val_loss_tensor.item())
                prune_now = False
                prune_msg = None
                if not np.isfinite(val_loss_value):
                    prune_now = True
                    prune_msg = (
                        f"[FTTransformerSklearn.fit_unsupervised] non-finite val loss "
                        f"(epoch={epoch}, val_loss={val_loss_value})"
                    )
                val_history.append(val_loss_value)

                if val_loss_value < best_loss:
                    best_loss = val_loss_value
                    best_state = {
                        k: (v.clone() if isinstance(
                            v, torch.Tensor) else copy.deepcopy(v))
                        for k, v in self.ft.state_dict().items()
                    }
                    patience_counter = 0
                else:
                    patience_counter += 1
                    if best_state is not None and patience_counter >= int(self.patience):
                        break

                if trial is not None and (not dist.is_initialized() or DistributedUtils.is_main_process()):
                    trial.report(val_loss_value, epoch)
                    if trial.should_prune():
                        prune_now = True

                if dist.is_initialized():
                    flag = torch.tensor(
                        [1 if prune_now else 0],
                        device=loss_tensor_device,
                        dtype=torch.int32,
                    )
                    dist.broadcast(flag, src=0)
                    prune_now = bool(flag.item())

                if prune_now:
                    if prune_msg:
                        raise optuna.TrialPruned(prune_msg)
                    raise optuna.TrialPruned()

        self.training_history = {"train": train_history, "val": val_history}
        self._plot_loss_curve(self.training_history, getattr(
            self, "loss_curve_path", None))
        if has_val and best_state is not None:
            self.ft.load_state_dict(best_state)
        return float(best_loss if has_val else (train_history[-1] if train_history else 0.0))

    def predict(self, X_test, geo_tokens=None, batch_size: Optional[int] = None, return_embedding: bool = False):
        # X_test must include all numeric/categorical columns; geo_tokens is optional.

        self.ft.eval()
        X_num, X_cat, X_geo, _, _, _ = self._tensorize_split(
            X_test, None, None, geo_tokens=geo_tokens, allow_none=True)

        num_rows = X_num.shape[0]
        if num_rows == 0:
            return np.empty(0, dtype=np.float32)

        device = self.device if isinstance(
            self.device, torch.device) else torch.device(self.device)

        def resolve_batch_size(n_rows: int) -> int:
            if batch_size is not None:
                return max(1, min(int(batch_size), n_rows))
            # Estimate a safe batch size based on model size to avoid attention OOM.
            token_cnt = self.num_numeric_tokens + len(self.cat_cols)
            if self.num_geo > 0:
                token_cnt += 1
            approx_units = max(1, token_cnt * max(1, self.d_model))
            if device.type == 'cuda':
                if approx_units >= 8192:
                    base = 512
                elif approx_units >= 4096:
                    base = 1024
                else:
                    base = 2048
            else:
                base = 512
            return max(1, min(base, n_rows))

        eff_batch = resolve_batch_size(num_rows)
        preds: List[torch.Tensor] = []

        inference_cm = getattr(torch, "inference_mode", torch.no_grad)
        with inference_cm():
            for start in range(0, num_rows, eff_batch):
                end = min(num_rows, start + eff_batch)
                X_num_b = X_num[start:end].to(device, non_blocking=True)
                X_cat_b = X_cat[start:end].to(device, non_blocking=True)
                X_geo_b = X_geo[start:end].to(device, non_blocking=True)
                pred_chunk = self.ft(
                    X_num_b, X_cat_b, X_geo_b, return_embedding=return_embedding)
                preds.append(pred_chunk.cpu())

        y_pred = torch.cat(preds, dim=0).numpy()

        if return_embedding:
            return y_pred

        if self.task_type == 'classification':
            # Convert logits to probabilities.
            y_pred = 1 / (1 + np.exp(-y_pred))
        else:
            # Model already has softplus; optionally apply log-exp smoothing: y_pred = log(1 + exp(y_pred)).
            y_pred = np.clip(y_pred, 1e-6, None)
        return y_pred.ravel()

    def set_params(self, params: dict):

        # Keep sklearn-style behavior.
        # Note: changing structural params (e.g., d_model/n_heads) requires refit to take effect.

        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Parameter {key} not found in model.")
        return self


# =============================================================================
# Simplified GNN implementation.
# =============================================================================


class SimpleGraphLayer(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.1):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)
        self.activation = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

    def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
        # Message passing with normalized sparse adjacency: A_hat * X * W.
        h = torch.sparse.mm(adj, x)
        h = self.linear(h)
        h = self.activation(h)
        return self.dropout(h)


class SimpleGNN(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 64, num_layers: int = 2,
                 dropout: float = 0.1, task_type: str = 'regression'):
        super().__init__()
        layers = []
        dim_in = input_dim
        for _ in range(max(1, num_layers)):
            layers.append(SimpleGraphLayer(
                dim_in, hidden_dim, dropout=dropout))
            dim_in = hidden_dim
        self.layers = nn.ModuleList(layers)
        self.output = nn.Linear(hidden_dim, 1)
        if task_type == 'classification':
            self.output_act = nn.Identity()
        else:
            self.output_act = nn.Softplus()
        self.task_type = task_type
        # Keep adjacency as a buffer for DataParallel copies.
        self.register_buffer("adj_buffer", torch.empty(0))

    def forward(self, x: torch.Tensor, adj: Optional[torch.Tensor] = None) -> torch.Tensor:
        adj_used = adj if adj is not None else getattr(
            self, "adj_buffer", None)
        if adj_used is None or adj_used.numel() == 0:
            raise RuntimeError("Adjacency is not set for GNN forward.")
        h = x
        for layer in self.layers:
            h = layer(h, adj_used)
        h = torch.sparse.mm(adj_used, h)
        out = self.output(h)
        return self.output_act(out)


class GraphNeuralNetSklearn(TorchTrainerMixin, nn.Module):
    def __init__(self, model_nme: str, input_dim: int, hidden_dim: int = 64,
                 num_layers: int = 2, k_neighbors: int = 10, dropout: float = 0.1,
                 learning_rate: float = 1e-3, epochs: int = 100, patience: int = 10,
                 task_type: str = 'regression', tweedie_power: float = 1.5,
                 weight_decay: float = 0.0,
                 use_data_parallel: bool = False, use_ddp: bool = False,
                 use_approx_knn: bool = True, approx_knn_threshold: int = 50000,
                 graph_cache_path: Optional[str] = None,
                 max_gpu_knn_nodes: Optional[int] = None,
                 knn_gpu_mem_ratio: float = 0.9,
                 knn_gpu_mem_overhead: float = 2.0,
                 knn_cpu_jobs: Optional[int] = -1) -> None:
        super().__init__()
        self.model_nme = model_nme
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.k_neighbors = max(1, k_neighbors)
        self.dropout = dropout
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.epochs = epochs
        self.patience = patience
        self.task_type = task_type
        self.use_approx_knn = use_approx_knn
        self.approx_knn_threshold = approx_knn_threshold
        self.graph_cache_path = Path(
            graph_cache_path) if graph_cache_path else None
        self.max_gpu_knn_nodes = max_gpu_knn_nodes
        self.knn_gpu_mem_ratio = max(0.0, min(1.0, knn_gpu_mem_ratio))
        self.knn_gpu_mem_overhead = max(1.0, knn_gpu_mem_overhead)
        self.knn_cpu_jobs = knn_cpu_jobs
        self._knn_warning_emitted = False
        self._adj_cache_meta: Optional[Dict[str, Any]] = None
        self._adj_cache_key: Optional[Tuple[Any, ...]] = None
        self._adj_cache_tensor: Optional[torch.Tensor] = None

        if self.task_type == 'classification':
            self.tw_power = None
        elif 'f' in self.model_nme:
            self.tw_power = 1.0
        elif 's' in self.model_nme:
            self.tw_power = 2.0
        else:
            self.tw_power = tweedie_power

        self.ddp_enabled = False
        self.local_rank = int(os.environ.get("LOCAL_RANK", 0))
        self.data_parallel_enabled = False
        self._ddp_disabled = False

        if use_ddp:
            world_size = int(os.environ.get("WORLD_SIZE", "1"))
            if world_size > 1:
                print(
                    "[GNN] DDP training is not supported; falling back to single process.",
                    flush=True,
                )
                self._ddp_disabled = True
                use_ddp = False

        # DDP only works with CUDA; fall back to single process if init fails.
        if use_ddp and torch.cuda.is_available():
            ddp_ok, local_rank, _, _ = DistributedUtils.setup_ddp()
            if ddp_ok:
                self.ddp_enabled = True
                self.local_rank = local_rank
                self.device = torch.device(f'cuda:{local_rank}')
            else:
                self.device = torch.device('cuda')
        elif torch.cuda.is_available():
            if self._ddp_disabled:
                self.device = torch.device(f'cuda:{self.local_rank}')
            else:
                self.device = torch.device('cuda')
        elif torch.backends.mps.is_available():
            self.device = torch.device('cpu')
            global _GNN_MPS_WARNED
            if not _GNN_MPS_WARNED:
                print(
                    "[GNN] MPS backend does not support sparse ops; falling back to CPU.",
                    flush=True,
                )
                _GNN_MPS_WARNED = True
        else:
            self.device = torch.device('cpu')
        self.use_pyg_knn = self.device.type == 'cuda' and _PYG_AVAILABLE

        self.gnn = SimpleGNN(
            input_dim=self.input_dim,
            hidden_dim=self.hidden_dim,
            num_layers=self.num_layers,
            dropout=self.dropout,
            task_type=self.task_type
        ).to(self.device)

        # DataParallel copies the full graph to each GPU and splits features; good for medium graphs.
        if (not self.ddp_enabled) and use_data_parallel and (self.device.type == 'cuda') and (torch.cuda.device_count() > 1):
            self.data_parallel_enabled = True
            self.gnn = nn.DataParallel(
                self.gnn, device_ids=list(range(torch.cuda.device_count())))
            self.device = torch.device('cuda')

        if self.ddp_enabled:
            self.gnn = DDP(
                self.gnn,
                device_ids=[self.local_rank],
                output_device=self.local_rank,
                find_unused_parameters=False
            )

    @staticmethod
    def _validate_vector(arr, name: str, n_rows: int) -> None:
        if arr is None:
            return
        if isinstance(arr, pd.DataFrame):
            if arr.shape[1] != 1:
                raise ValueError(f"{name} must be 1d (single column).")
            length = len(arr)
        else:
            arr_np = np.asarray(arr)
            if arr_np.ndim == 0:
                raise ValueError(f"{name} must be 1d.")
            if arr_np.ndim > 2 or (arr_np.ndim == 2 and arr_np.shape[1] != 1):
                raise ValueError(f"{name} must be 1d or Nx1.")
            length = arr_np.shape[0]
        if length != n_rows:
            raise ValueError(
                f"{name} length {length} does not match X length {n_rows}."
            )

    def _unwrap_gnn(self) -> nn.Module:
        if isinstance(self.gnn, (DDP, nn.DataParallel)):
            return self.gnn.module
        return self.gnn

    def _set_adj_buffer(self, adj: torch.Tensor) -> None:
        base = self._unwrap_gnn()
        if hasattr(base, "adj_buffer"):
            base.adj_buffer = adj
        else:
            base.register_buffer("adj_buffer", adj)

    def _graph_cache_meta(self, X_df: pd.DataFrame) -> Dict[str, Any]:
        row_hash = pd.util.hash_pandas_object(X_df, index=False).values
        idx_hash = pd.util.hash_pandas_object(X_df.index, index=False).values
        col_sig = ",".join(map(str, X_df.columns))
        hasher = hashlib.sha256()
        hasher.update(row_hash.tobytes())
        hasher.update(idx_hash.tobytes())
        hasher.update(col_sig.encode("utf-8", errors="ignore"))
        knn_config = {
            "k_neighbors": int(self.k_neighbors),
            "use_approx_knn": bool(self.use_approx_knn),
            "approx_knn_threshold": int(self.approx_knn_threshold),
            "use_pyg_knn": bool(self.use_pyg_knn),
            "pynndescent_available": bool(_PYNN_AVAILABLE),
            "max_gpu_knn_nodes": (
                None if self.max_gpu_knn_nodes is None else int(self.max_gpu_knn_nodes)
            ),
            "knn_gpu_mem_ratio": float(self.knn_gpu_mem_ratio),
            "knn_gpu_mem_overhead": float(self.knn_gpu_mem_overhead),
        }
        return {
            "n_samples": int(X_df.shape[0]),
            "n_features": int(X_df.shape[1]),
            "hash": hasher.hexdigest(),
            "knn_config": knn_config,
        }

    def _graph_cache_key(self, X_df: pd.DataFrame) -> Tuple[Any, ...]:
        return (
            id(X_df),
            id(getattr(X_df, "_mgr", None)),
            id(X_df.index),
            X_df.shape,
            tuple(map(str, X_df.columns)),
            X_df.attrs.get("graph_cache_key"),
        )

    def invalidate_graph_cache(self) -> None:
        self._adj_cache_meta = None
        self._adj_cache_key = None
        self._adj_cache_tensor = None

    def _load_cached_adj(self,
                         X_df: pd.DataFrame,
                         meta_expected: Optional[Dict[str, Any]] = None) -> Optional[torch.Tensor]:
        if self.graph_cache_path and self.graph_cache_path.exists():
            if meta_expected is None:
                meta_expected = self._graph_cache_meta(X_df)
            try:
                payload = torch.load(self.graph_cache_path,
                                     map_location=self.device)
            except Exception as exc:
                print(
                    f"[GNN] Failed to load cached graph from {self.graph_cache_path}: {exc}")
                return None
            if isinstance(payload, dict) and "adj" in payload:
                meta_cached = payload.get("meta")
                if meta_cached == meta_expected:
                    return payload["adj"].to(self.device)
                print(
                    f"[GNN] Cached graph metadata mismatch; rebuilding: {self.graph_cache_path}")
                return None
            if isinstance(payload, torch.Tensor):
                print(
                    f"[GNN] Cached graph missing metadata; rebuilding: {self.graph_cache_path}")
                return None
            print(
                f"[GNN] Invalid cached graph format; rebuilding: {self.graph_cache_path}")
        return None

    def _build_edge_index_cpu(self, X_np: np.ndarray) -> torch.Tensor:
        n_samples = X_np.shape[0]
        k = min(self.k_neighbors, max(1, n_samples - 1))
        n_neighbors = min(k + 1, n_samples)
        use_approx = (self.use_approx_knn or n_samples >=
                      self.approx_knn_threshold) and _PYNN_AVAILABLE
        indices = None
        if use_approx:
            try:
                nn_index = pynndescent.NNDescent(
                    X_np,
                    n_neighbors=n_neighbors,
                    random_state=0
                )
                indices, _ = nn_index.neighbor_graph
            except Exception as exc:
                print(
                    f"[GNN] Approximate kNN failed ({exc}); falling back to exact search.")
                use_approx = False

        if indices is None:
            nbrs = NearestNeighbors(
                n_neighbors=n_neighbors,
                algorithm="auto",
                n_jobs=self.knn_cpu_jobs,
            )
            nbrs.fit(X_np)
            _, indices = nbrs.kneighbors(X_np)

        indices = np.asarray(indices)
        rows = np.repeat(np.arange(n_samples), n_neighbors).astype(
            np.int64, copy=False)
        cols = indices.reshape(-1).astype(np.int64, copy=False)
        mask = rows != cols
        rows = rows[mask]
        cols = cols[mask]
        rows_base = rows
        cols_base = cols
        self_loops = np.arange(n_samples, dtype=np.int64)
        rows = np.concatenate([rows_base, cols_base, self_loops])
        cols = np.concatenate([cols_base, rows_base, self_loops])

        edge_index_np = np.stack([rows, cols], axis=0)
        edge_index = torch.as_tensor(edge_index_np, device=self.device)
        return edge_index

    def _build_edge_index_gpu(self, X_tensor: torch.Tensor) -> torch.Tensor:
        if not self.use_pyg_knn or knn_graph is None or add_self_loops is None or to_undirected is None:
            # Defensive: check use_pyg_knn before calling.
            raise RuntimeError(
                "GPU graph builder requested but PyG is unavailable.")

        n_samples = X_tensor.size(0)
        k = min(self.k_neighbors, max(1, n_samples - 1))

        # knn_graph runs on GPU to avoid CPU graph construction bottlenecks.
        edge_index = knn_graph(
            X_tensor,
            k=k,
            loop=False
        )
        edge_index = to_undirected(edge_index, num_nodes=n_samples)
        edge_index, _ = add_self_loops(edge_index, num_nodes=n_samples)
        return edge_index

    def _log_knn_fallback(self, reason: str) -> None:
        if self._knn_warning_emitted:
            return
        if (not self.ddp_enabled) or self.local_rank == 0:
            print(f"[GNN] Falling back to CPU kNN builder: {reason}")
        self._knn_warning_emitted = True

    def _should_use_gpu_knn(self, n_samples: int, X_tensor: torch.Tensor) -> bool:
        if not self.use_pyg_knn:
            return False

        reason = None
        if self.max_gpu_knn_nodes is not None and n_samples > self.max_gpu_knn_nodes:
            reason = f"node count {n_samples} exceeds max_gpu_knn_nodes={self.max_gpu_knn_nodes}"
        elif self.device.type == 'cuda' and torch.cuda.is_available():
            try:
                device_index = self.device.index
                if device_index is None:
                    device_index = torch.cuda.current_device()
                free_mem, total_mem = torch.cuda.mem_get_info(device_index)
                feature_bytes = X_tensor.element_size() * X_tensor.nelement()
                required = int(feature_bytes * self.knn_gpu_mem_overhead)
                budget = int(free_mem * self.knn_gpu_mem_ratio)
                if required > budget:
                    required_gb = required / (1024 ** 3)
                    budget_gb = budget / (1024 ** 3)
                    reason = (f"requires ~{required_gb:.2f} GiB temporary GPU memory "
                              f"but only {budget_gb:.2f} GiB free on cuda:{device_index}")
            except Exception:
                # On older versions or some environments, mem_get_info may be unavailable; default to trying GPU.
                reason = None

        if reason:
            self._log_knn_fallback(reason)
            return False
        return True

    def _normalized_adj(self, edge_index: torch.Tensor, num_nodes: int) -> torch.Tensor:
        values = torch.ones(edge_index.shape[1], device=self.device)
        adj = torch.sparse_coo_tensor(
            edge_index.to(self.device), values, (num_nodes, num_nodes))
        adj = adj.coalesce()

        deg = torch.sparse.sum(adj, dim=1).to_dense()
        deg_inv_sqrt = torch.pow(deg + 1e-8, -0.5)
        row, col = adj.indices()
        norm_values = deg_inv_sqrt[row] * adj.values() * deg_inv_sqrt[col]
        adj_norm = torch.sparse_coo_tensor(
            adj.indices(), norm_values, size=adj.shape)
        return adj_norm

    def _tensorize_split(self, X, y, w, allow_none: bool = False):
        if X is None and allow_none:
            return None, None, None
        if not isinstance(X, pd.DataFrame):
            raise ValueError("X must be a pandas DataFrame for GNN.")
        n_rows = len(X)
        if y is not None:
            self._validate_vector(y, "y", n_rows)
        if w is not None:
            self._validate_vector(w, "w", n_rows)
        X_np = X.to_numpy(dtype=np.float32, copy=False) if hasattr(
            X, "to_numpy") else np.asarray(X, dtype=np.float32)
        X_tensor = torch.as_tensor(
            X_np, dtype=torch.float32, device=self.device)
        if y is None:
            y_tensor = None
        else:
            y_np = y.to_numpy(dtype=np.float32, copy=False) if hasattr(
                y, "to_numpy") else np.asarray(y, dtype=np.float32)
            y_tensor = torch.as_tensor(
                y_np, dtype=torch.float32, device=self.device).view(-1, 1)
        if w is None:
            w_tensor = torch.ones(
                (len(X), 1), dtype=torch.float32, device=self.device)
        else:
            w_np = w.to_numpy(dtype=np.float32, copy=False) if hasattr(
                w, "to_numpy") else np.asarray(w, dtype=np.float32)
            w_tensor = torch.as_tensor(
                w_np, dtype=torch.float32, device=self.device).view(-1, 1)
        return X_tensor, y_tensor, w_tensor

    def _build_graph_from_df(self, X_df: pd.DataFrame, X_tensor: Optional[torch.Tensor] = None) -> torch.Tensor:
        if not isinstance(X_df, pd.DataFrame):
            raise ValueError("X must be a pandas DataFrame for graph building.")
        meta_expected = None
        cache_key = None
        if self.graph_cache_path:
            meta_expected = self._graph_cache_meta(X_df)
            if self._adj_cache_meta == meta_expected and self._adj_cache_tensor is not None:
                cached = self._adj_cache_tensor
                if cached.device != self.device:
                    cached = cached.to(self.device)
                    self._adj_cache_tensor = cached
                return cached
        else:
            cache_key = self._graph_cache_key(X_df)
            if self._adj_cache_key == cache_key and self._adj_cache_tensor is not None:
                cached = self._adj_cache_tensor
                if cached.device != self.device:
                    cached = cached.to(self.device)
                    self._adj_cache_tensor = cached
                return cached
        X_np = None
        if X_tensor is None:
            X_np = X_df.to_numpy(dtype=np.float32, copy=False)
            X_tensor = torch.as_tensor(
                X_np, dtype=torch.float32, device=self.device)
        if self.graph_cache_path:
            cached = self._load_cached_adj(X_df, meta_expected=meta_expected)
            if cached is not None:
                self._adj_cache_meta = meta_expected
                self._adj_cache_key = None
                self._adj_cache_tensor = cached
                return cached
        use_gpu_knn = self._should_use_gpu_knn(X_df.shape[0], X_tensor)
        if use_gpu_knn:
            edge_index = self._build_edge_index_gpu(X_tensor)
        else:
            if X_np is None:
                X_np = X_df.to_numpy(dtype=np.float32, copy=False)
            edge_index = self._build_edge_index_cpu(X_np)
        adj_norm = self._normalized_adj(edge_index, X_df.shape[0])
        if self.graph_cache_path:
            try:
                IOUtils.ensure_parent_dir(str(self.graph_cache_path))
                torch.save({"adj": adj_norm.cpu(), "meta": meta_expected}, self.graph_cache_path)
            except Exception as exc:
                print(
                    f"[GNN] Failed to cache graph to {self.graph_cache_path}: {exc}")
            self._adj_cache_meta = meta_expected
            self._adj_cache_key = None
        else:
            self._adj_cache_meta = None
            self._adj_cache_key = cache_key
        self._adj_cache_tensor = adj_norm
        return adj_norm

    def fit(self, X_train, y_train, w_train=None,
            X_val=None, y_val=None, w_val=None,
            trial: Optional[optuna.trial.Trial] = None):

        X_train_tensor, y_train_tensor, w_train_tensor = self._tensorize_split(
            X_train, y_train, w_train, allow_none=False)
        has_val = X_val is not None and y_val is not None
        if has_val:
            X_val_tensor, y_val_tensor, w_val_tensor = self._tensorize_split(
                X_val, y_val, w_val, allow_none=False)
        else:
            X_val_tensor = y_val_tensor = w_val_tensor = None

        adj_train = self._build_graph_from_df(X_train, X_train_tensor)
        adj_val = self._build_graph_from_df(
            X_val, X_val_tensor) if has_val else None
        # DataParallel needs adjacency cached on the model to avoid scatter.
        self._set_adj_buffer(adj_train)

        base_gnn = self._unwrap_gnn()
        optimizer = torch.optim.Adam(
            base_gnn.parameters(),
            lr=self.learning_rate,
            weight_decay=float(getattr(self, "weight_decay", 0.0)),
        )
        scaler = GradScaler(enabled=(self.device.type == 'cuda'))

        best_loss = float('inf')
        best_state = None
        patience_counter = 0
        best_epoch = None

        for epoch in range(1, self.epochs + 1):
            epoch_start_ts = time.time()
            self.gnn.train()
            optimizer.zero_grad()
            with autocast(enabled=(self.device.type == 'cuda')):
                if self.data_parallel_enabled:
                    y_pred = self.gnn(X_train_tensor)
                else:
                    y_pred = self.gnn(X_train_tensor, adj_train)
                loss = self._compute_weighted_loss(
                    y_pred, y_train_tensor, w_train_tensor, apply_softplus=False)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            clip_grad_norm_(self.gnn.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()

            val_loss = None
            if has_val:
                self.gnn.eval()
                if self.data_parallel_enabled and adj_val is not None:
                    self._set_adj_buffer(adj_val)
                with torch.no_grad(), autocast(enabled=(self.device.type == 'cuda')):
                    if self.data_parallel_enabled:
                        y_val_pred = self.gnn(X_val_tensor)
                    else:
                        y_val_pred = self.gnn(X_val_tensor, adj_val)
                    val_loss = self._compute_weighted_loss(
                        y_val_pred, y_val_tensor, w_val_tensor, apply_softplus=False)
                if self.data_parallel_enabled:
                    # Restore training adjacency.
                    self._set_adj_buffer(adj_train)

                is_best = val_loss is not None and val_loss < best_loss
                best_loss, best_state, patience_counter, stop_training = self._early_stop_update(
                    val_loss, best_loss, best_state, patience_counter, base_gnn,
                    ignore_keys=["adj_buffer"])
                if is_best:
                    best_epoch = epoch

                prune_now = False
                if trial is not None:
                    trial.report(val_loss, epoch)
                    if trial.should_prune():
                        prune_now = True

                if dist.is_initialized():
                    flag = torch.tensor(
                        [1 if prune_now else 0],
                        device=self.device,
                        dtype=torch.int32,
                    )
                    dist.broadcast(flag, src=0)
                    prune_now = bool(flag.item())

                if prune_now:
                    raise optuna.TrialPruned()
                if stop_training:
                    break

            should_log = (not dist.is_initialized()
                          or DistributedUtils.is_main_process())
            if should_log:
                elapsed = int(time.time() - epoch_start_ts)
                if val_loss is None:
                    print(
                        f"[GNN] Epoch {epoch}/{self.epochs} loss={float(loss):.6f} elapsed={elapsed}s",
                        flush=True,
                    )
                else:
                    print(
                        f"[GNN] Epoch {epoch}/{self.epochs} loss={float(loss):.6f} "
                        f"val_loss={float(val_loss):.6f} elapsed={elapsed}s",
                        flush=True,
                    )

        if best_state is not None:
            base_gnn.load_state_dict(best_state, strict=False)
        self.best_epoch = int(best_epoch or self.epochs)

    def predict(self, X: pd.DataFrame) -> np.ndarray:
        self.gnn.eval()
        X_tensor, _, _ = self._tensorize_split(
            X, None, None, allow_none=False)
        adj = self._build_graph_from_df(X, X_tensor)
        if self.data_parallel_enabled:
            self._set_adj_buffer(adj)
        inference_cm = getattr(torch, "inference_mode", torch.no_grad)
        with inference_cm():
            if self.data_parallel_enabled:
                y_pred = self.gnn(X_tensor).cpu().numpy()
            else:
                y_pred = self.gnn(X_tensor, adj).cpu().numpy()
        if self.task_type == 'classification':
            y_pred = 1 / (1 + np.exp(-y_pred))
        else:
            y_pred = np.clip(y_pred, 1e-6, None)
        return y_pred.ravel()

    def encode(self, X: pd.DataFrame) -> np.ndarray:
        """Return per-sample node embeddings (hidden representations)."""
        base = self._unwrap_gnn()
        base.eval()
        X_tensor, _, _ = self._tensorize_split(X, None, None, allow_none=False)
        adj = self._build_graph_from_df(X, X_tensor)
        if self.data_parallel_enabled:
            self._set_adj_buffer(adj)
        inference_cm = getattr(torch, "inference_mode", torch.no_grad)
        with inference_cm():
            h = X_tensor
            layers = getattr(base, "layers", None)
            if layers is None:
                raise RuntimeError("GNN base module does not expose layers.")
            for layer in layers:
                h = layer(h, adj)
            h = torch.sparse.mm(adj, h)
        return h.detach().cpu().numpy()

    def set_params(self, params: Dict[str, Any]):
        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Parameter {key} not found in GNN model.")
        # Rebuild the backbone after structural parameter changes.
        self.gnn = SimpleGNN(
            input_dim=self.input_dim,
            hidden_dim=self.hidden_dim,
            num_layers=self.num_layers,
            dropout=self.dropout,
            task_type=self.task_type
        ).to(self.device)
        return self
