modelling/loader.py

# loader.py
# -*- coding: utf-8 -*-

import ast
from io import BytesIO
from urllib.parse import urljoin

import pandas as pd
import requests
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

from utils import load_json


class CenterSquareCrop:
    """
    Crop image to a centered square without resizing.
    """

    def __call__(self, img: Image.Image):
        w, h = img.size

        if w == h:
            return img

        if w > h:
            left = (w - h) // 2
            right = left + h
            top = 0
            bottom = h
        else:
            top = (h - w) // 2
            bottom = top + w
            left = 0
            right = w
        return img.crop((left, top, right, bottom))


def build_image_transform(image_size: int):
    return transforms.Compose([
        CenterSquareCrop(),
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
    ])


def join_photo_root(photo_root: str, relative_path: str) -> str:
    """
    Join photo_root and relative path.

    Supports:
    - local filesystem roots
    - http / https roots
    """
    if photo_root.startswith("http://") or photo_root.startswith("https://"):  # noqa
        return urljoin(photo_root.rstrip("/") + "/", relative_path)

    return photo_root.rstrip("/") + "/" + relative_path.lstrip("/")


def parse_numeric_cell(value: str, n_in: int):
    """
    Convert numeric csv cell to list[float].

    Returns:
        values, is_valid

    Data assumption:
    - Empty value is always ""
    - Scalar numeric -> "12.3"
    - Vector numeric -> "[1.2,3.4,5.6]"
    """
    if value == "":
        return [0.0] * n_in, False

    if n_in == 1:
        return [float(value)], True

    vec = ast.literal_eval(value)
    if len(vec) != n_in:
        raise ValueError(f"Numeric vector length mismatch: expected {n_in}, got {len(vec)}")
    return [float(v) for v in vec], True


class SoilFormerDataset(Dataset):

    def __init__(
            self,
            csv_path: str,
            photo_map_path: str,
            cat_vocab_path: str,
            numeric_vocab_path: str,
            numeric_stats_path: str,
            photo_root: str,
            image_size: int = 512,
            id_column: str = "id",
    ):
        self.df = pd.read_csv(
            csv_path,
            keep_default_na=False,
            na_filter=False,
            low_memory=False,
        )

        self.photo_map = load_json(photo_map_path)
        self.cat_vocab = load_json(cat_vocab_path)
        self.numeric_vocab = load_json(numeric_vocab_path)

        self.photo_root = photo_root
        self.id_column = id_column
        self.image_size = int(image_size)
        self.image_transform = build_image_transform(self.image_size)

        # Keep json order exactly
        self.cat_columns = list(self.cat_vocab.keys())
        self.numeric_groups = self.numeric_vocab["groups"]
        self.numeric_stats_df = pd.read_csv(numeric_stats_path)
        self.numeric_stats_index = self.numeric_stats_df.set_index("column")

        # Numeric mean/std
        self.numeric_stats = {}
        for _, row in self.numeric_stats_df.iterrows():
            col = row["column"]
            mean = float(row["mean"])
            std = float(row["std"])
            if std == 0.0:
                std = 1.0
            self.numeric_stats[col] = (mean, std)

        # For active masking
        self.cat_mask_local_ids = torch.tensor(
            [int(self.cat_vocab[col]["mask_local_id"]) for col in self.cat_columns],
            dtype=torch.long,
        )

    def __len__(self):
        return len(self.df)

    def load_image(self, path: str):
        if path.startswith("http://") or path.startswith("https://"):  # noqa
            resp = requests.get(path, timeout=(3, 10))
            resp.raise_for_status()
            img = Image.open(BytesIO(resp.content)).convert("RGB")
        else:
            img = Image.open(path).convert("RGB")

        return self.image_transform(img)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        sample_id = row[self.id_column]

        # -----------------------
        # categorical features
        # -----------------------
        cat_ids = []
        cat_valids = []

        for col in self.cat_columns:
            spec = self.cat_vocab[col]
            label2id = spec["label2id"]
            mask_id = spec["mask_local_id"]

            value = row[col]

            if value == "":
                cat_ids.append(mask_id)
                cat_valids.append(False)
            else:
                if value not in label2id:
                    raise KeyError(f"Unknown categorical value: column={col}, value={value!r}")
                cat_ids.append(label2id[value])
                cat_valids.append(True)

        cat_ids = torch.tensor(cat_ids, dtype=torch.long)
        cat_valids = torch.tensor(cat_valids, dtype=torch.bool)

        # -----------------------
        # numeric features
        # -----------------------
        numeric_values_by_nin = {}
        numeric_valid_positions_by_nin = {}

        for group in self.numeric_groups:
            n_in = int(group["n_in"])
            features = group["feature_names"]

            values = []
            valids = []

            for feat in features:
                cell = row[feat]
                parsed, is_valid = parse_numeric_cell(cell, n_in)
                if is_valid:
                    mean, std = self.numeric_stats[feat]
                    parsed = [(v - mean) / std for v in parsed]
                values.append(parsed)
                valids.append(is_valid)

            numeric_values_by_nin[n_in] = torch.tensor(values, dtype=torch.float32)
            numeric_valid_positions_by_nin[n_in] = torch.tensor(valids, dtype=torch.bool)

        # -----------------------
        # vision
        # -----------------------
        try:
            relative_path = self.photo_map[sample_id]
            full_path = join_photo_root(self.photo_root, relative_path)
            image = self.load_image(full_path)
            vision_valid = True
        except Exception:  # noqa
            image = torch.zeros(3, self.image_size, self.image_size, dtype=torch.float32)
            vision_valid = False

        vision_valid = torch.tensor(vision_valid, dtype=torch.bool)

        return {
            "row_idx": torch.tensor(idx, dtype=torch.long),
            "sample_id": sample_id,
            "cat_local_ids": cat_ids,
            "cat_valid_positions": cat_valids,
            "numeric_values_by_nin": numeric_values_by_nin,
            "numeric_valid_positions_by_nin": numeric_valid_positions_by_nin,
            "pixel_values": image,
            "vision_valid_positions": vision_valid,
        }

    @staticmethod
    def collate_fn(batch):
        cat_ids = torch.stack([b["cat_local_ids"] for b in batch], dim=0)
        cat_valids = torch.stack([b["cat_valid_positions"] for b in batch], dim=0)

        group_keys = list(batch[0]["numeric_values_by_nin"].keys())

        numeric_values_by_nin = {}
        numeric_valid_positions_by_nin = {}

        for k in group_keys:
            numeric_values_by_nin[k] = torch.stack(
                [b["numeric_values_by_nin"][k] for b in batch],
                dim=0,
            )
            numeric_valid_positions_by_nin[k] = torch.stack(
                [b["numeric_valid_positions_by_nin"][k] for b in batch],
                dim=0,
            )

        pixel_values = torch.stack([b["pixel_values"] for b in batch], dim=0)
        vision_valid_positions = torch.stack([b["vision_valid_positions"] for b in batch], dim=0)
        row_idx = torch.stack([b["row_idx"] for b in batch], dim=0)
        sample_ids = [b["sample_id"] for b in batch]

        return {
            "row_idx": row_idx,
            "sample_id": sample_ids,
            "cat_local_ids": cat_ids,
            "numeric_values_by_nin": numeric_values_by_nin,
            "cat_valid_positions": cat_valids,
            "numeric_valid_positions_by_nin": numeric_valid_positions_by_nin,
            "pixel_values": pixel_values,
            "vision_valid_positions": vision_valid_positions,
        }

    def perform_active_mask(self, batch, cat_ratio=0.15, num_ratio=0.15, seed=None):
        """
        Apply active masking to categorical and numeric inputs.

        Conventions
        -----------
        Input batch must contain:
          - cat_local_ids: [B, M] LongTensor
          - cat_valid_positions: [B, M] Bool/0-1 tensor
          - numeric_values_by_nin: Dict[int, Tensor[B, V, n_in]]
          - numeric_valid_positions_by_nin: Dict[int, Tensor[B, V]]

        Output batch will additionally contain:
          - original_cat_local_ids
          - original_cat_valid_positions
          - original_numeric_values_by_nin
          - original_numeric_valid_positions_by_nin

          - masked_cat_local_ids
          - masked_cat_valid_positions
          - masked_numeric_values_by_nin
          - masked_numeric_valid_positions_by_nin

          - cat_loss_mask: [B, M] BoolTensor
          - numeric_loss_mask_by_nin: Dict[int, BoolTensor[B, V]]

        Semantics
        ---------
        - Only originally valid positions can be actively masked.
        - Masked categorical positions:
            local_id -> self.cat_mask_local_ids[col]
            valid    -> False
        - Masked numeric positions:
            values   -> 0
            valid    -> False
        - original_* fields always preserve the unmodified input batch content.
        """
        # --------------------------------------------------
        # Validate ratios
        # --------------------------------------------------
        if not (0.0 <= cat_ratio <= 1.0):
            raise ValueError(f"cat_ratio must be in [0, 1], got {cat_ratio}")
        if not (0.0 <= num_ratio <= 1.0):
            raise ValueError(f"num_ratio must be in [0, 1], got {num_ratio}")

        # --------------------------------------------------
        # Validate required keys
        # --------------------------------------------------
        required_keys = [
            "cat_local_ids",
            "cat_valid_positions",
            "numeric_values_by_nin",
            "numeric_valid_positions_by_nin",
        ]
        for k in required_keys:
            if k not in batch:
                raise KeyError(f"Missing key in batch: {k}")

        cat_local_ids = batch["cat_local_ids"]
        cat_valid_positions = batch["cat_valid_positions"]
        numeric_values_by_nin = batch["numeric_values_by_nin"]
        numeric_valid_positions_by_nin = batch["numeric_valid_positions_by_nin"]

        if cat_local_ids.dim() != 2:
            raise ValueError(f"cat_local_ids must be [B, M], got {tuple(cat_local_ids.shape)}")
        if cat_valid_positions.shape != cat_local_ids.shape:
            raise ValueError(
                f"cat_valid_positions must match cat_local_ids shape, got "
                f"{tuple(cat_valid_positions.shape)} vs {tuple(cat_local_ids.shape)}"
            )

        if not isinstance(numeric_values_by_nin, dict):
            raise ValueError("numeric_values_by_nin must be a dict")
        if not isinstance(numeric_valid_positions_by_nin, dict):
            raise ValueError("numeric_valid_positions_by_nin must be a dict")

        B, M = cat_local_ids.shape
        device = cat_local_ids.device

        if self.cat_mask_local_ids.dim() != 1 or self.cat_mask_local_ids.numel() != M:
            raise ValueError(
                f"self.cat_mask_local_ids must be [M] with M={M}, got {tuple(self.cat_mask_local_ids.shape)}"
            )
        cat_mask_local_ids = self.cat_mask_local_ids.to(device=device, dtype=cat_local_ids.dtype)

        # --------------------------------------------------
        # Random generator
        # --------------------------------------------------
        if device.type == "cuda":
            generator = torch.Generator(device=device)
        else:
            generator = torch.Generator()

        if seed is not None:
            generator.manual_seed(seed)

        # --------------------------------------------------
        # Start from shallow copy only
        # --------------------------------------------------
        masked_batch = dict(batch)

        # Preserve original aliases (do NOT deepcopy)
        masked_batch["original_cat_local_ids"] = batch["cat_local_ids"]
        masked_batch["original_cat_valid_positions"] = batch["cat_valid_positions"]
        masked_batch["original_numeric_values_by_nin"] = batch["numeric_values_by_nin"]
        masked_batch["original_numeric_valid_positions_by_nin"] = batch["numeric_valid_positions_by_nin"]

        # --------------------------------------------------
        # Fast path: no active masking at all
        # --------------------------------------------------
        if cat_ratio == 0.0 and num_ratio == 0.0:
            masked_batch["masked_cat_local_ids"] = batch["cat_local_ids"]
            masked_batch["masked_cat_valid_positions"] = batch["cat_valid_positions"]

            masked_batch["masked_numeric_values_by_nin"] = batch["numeric_values_by_nin"]
            masked_batch["masked_numeric_valid_positions_by_nin"] = batch["numeric_valid_positions_by_nin"]

            masked_batch["cat_loss_mask"] = torch.zeros(
                (B, M), dtype=torch.bool, device=device
            )
            masked_batch["numeric_loss_mask_by_nin"] = {
                n_in: torch.zeros_like(valid_positions, dtype=torch.bool)
                for n_in, valid_positions in numeric_valid_positions_by_nin.items()
            }
            return masked_batch

        # --------------------------------------------------
        # Categorical masking
        # --------------------------------------------------
        original_cat_valid_positions = cat_valid_positions.bool()

        masked_cat_local_ids = cat_local_ids.clone()
        masked_cat_valid_positions = original_cat_valid_positions.clone()
        cat_loss_mask = torch.zeros((B, M), dtype=torch.bool, device=device)

        if cat_ratio > 0.0:
            for b in range(B):
                valid_idx = torch.nonzero(original_cat_valid_positions[b], as_tuple=False).squeeze(1)
                n_valid = valid_idx.numel()
                if n_valid == 0:
                    continue

                k = int(round(n_valid * cat_ratio))
                if k <= 0:
                    continue
                if k > n_valid:
                    k = n_valid

                perm = valid_idx[
                    torch.randperm(n_valid, generator=generator, device=device)[:k]
                ]
                cat_loss_mask[b, perm] = True

            expanded_cat_mask_ids = cat_mask_local_ids.view(1, M).expand(B, M)
            masked_cat_local_ids[cat_loss_mask] = expanded_cat_mask_ids[cat_loss_mask]
            masked_cat_valid_positions = masked_cat_valid_positions & (~cat_loss_mask)

        masked_batch["masked_cat_local_ids"] = masked_cat_local_ids
        masked_batch["masked_cat_valid_positions"] = masked_cat_valid_positions
        masked_batch["cat_loss_mask"] = cat_loss_mask

        # --------------------------------------------------
        # Numeric masking
        # --------------------------------------------------
        masked_numeric_values_by_nin = {}
        masked_numeric_valid_positions_by_nin = {}
        numeric_loss_mask_by_nin = {}

        # keep deterministic ordering if caller passed mixed int-like keys
        for n_in in sorted(numeric_values_by_nin.keys(), key=int):
            values = numeric_values_by_nin[n_in]
            if n_in not in numeric_valid_positions_by_nin:
                raise KeyError(f"Missing numeric_valid_positions_by_nin[{n_in}]")

            valid_positions = numeric_valid_positions_by_nin[n_in]

            if values.dim() != 3:
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] must be [B, V, n_in], got {tuple(values.shape)}"
                )

            Bn, V, Nin = values.shape
            if Bn != B:
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] batch mismatch: got {Bn}, expected {B}"
                )
            if int(Nin) != int(n_in):
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] last dim mismatch: got {Nin}, expected {n_in}"
                )
            if valid_positions.shape != (B, V):
                raise ValueError(
                    f"numeric_valid_positions_by_nin[{n_in}] must be [B,V]=({B},{V}), "
                    f"got {tuple(valid_positions.shape)}"
                )

            original_valid = valid_positions.bool()

            # IMPORTANT: clone before modifying
            masked_values = values.clone()
            masked_valid_positions = original_valid.clone()
            num_loss_mask = torch.zeros((B, V), dtype=torch.bool, device=values.device)

            if num_ratio > 0.0:
                for b in range(B):
                    valid_idx = torch.nonzero(original_valid[b], as_tuple=False).squeeze(1)
                    n_valid = valid_idx.numel()
                    if n_valid == 0:
                        continue

                    k = int(round(n_valid * num_ratio))
                    if k <= 0:
                        continue
                    if k > n_valid:
                        k = n_valid

                    perm = valid_idx[
                        torch.randperm(n_valid, generator=generator, device=values.device)[:k]
                    ]
                    num_loss_mask[b, perm] = True

                # masked numeric columns become zero and invalid
                masked_values[num_loss_mask] = 0.0
                masked_valid_positions = masked_valid_positions & (~num_loss_mask)

            masked_numeric_values_by_nin[n_in] = masked_values
            masked_numeric_valid_positions_by_nin[n_in] = masked_valid_positions
            numeric_loss_mask_by_nin[n_in] = num_loss_mask

        masked_batch["masked_numeric_values_by_nin"] = masked_numeric_values_by_nin
        masked_batch["masked_numeric_valid_positions_by_nin"] = masked_numeric_valid_positions_by_nin
        masked_batch["numeric_loss_mask_by_nin"] = numeric_loss_mask_by_nin

        return masked_batch


    def perform_active_mask_single(self, batch, feature_name, assert_not_missing=True):
        """
        Actively mask exactly one feature specified by feature_name.
    
        Parameters
        ----------
        batch : dict
            Same input convention as perform_active_mask(...).
        feature_name : str
            Full feature name. Can be either categorical or numeric.
        assert_not_missing : bool
            If True, require the target feature to be originally valid for all samples
            in the batch. Otherwise raise ValueError.
            If False, only originally valid positions are masked; naturally missing
            positions remain missing and are not included in the loss mask.
    
        Returns
        -------
        masked_batch : dict
            Same output convention as perform_active_mask(...), except that exactly
            one feature is actively masked.
        """
    
        # --------------------------------------------------
        # Validate required keys
        # --------------------------------------------------
        required_keys = [
            "cat_local_ids",
            "cat_valid_positions",
            "numeric_values_by_nin",
            "numeric_valid_positions_by_nin",
        ]
        for k in required_keys:
            if k not in batch:
                raise KeyError(f"Missing key in batch: {k}")
    
        cat_local_ids = batch["cat_local_ids"]
        cat_valid_positions = batch["cat_valid_positions"]
        numeric_values_by_nin = batch["numeric_values_by_nin"]
        numeric_valid_positions_by_nin = batch["numeric_valid_positions_by_nin"]
    
        if cat_local_ids.dim() != 2:
            raise ValueError(f"cat_local_ids must be [B, M], got {tuple(cat_local_ids.shape)}")
        if cat_valid_positions.shape != cat_local_ids.shape:
            raise ValueError(
                f"cat_valid_positions must match cat_local_ids shape, got "
                f"{tuple(cat_valid_positions.shape)} vs {tuple(cat_local_ids.shape)}"
            )
    
        if not isinstance(numeric_values_by_nin, dict):
            raise ValueError("numeric_values_by_nin must be a dict")
        if not isinstance(numeric_valid_positions_by_nin, dict):
            raise ValueError("numeric_valid_positions_by_nin must be a dict")
    
        B, M = cat_local_ids.shape
        device = cat_local_ids.device
    
        if self.cat_mask_local_ids.dim() != 1 or self.cat_mask_local_ids.numel() != M:
            raise ValueError(
                f"self.cat_mask_local_ids must be [M] with M={M}, got {tuple(self.cat_mask_local_ids.shape)}"
            )
        cat_mask_local_ids = self.cat_mask_local_ids.to(device=device, dtype=cat_local_ids.dtype)
    
        # --------------------------------------------------
        # Resolve feature_name -> categorical col or numeric (n_in, v_idx)
        # --------------------------------------------------
        # Assumptions:
        #   - self.cat_vocab is the categorical vocab dict keyed by full feature name
        #   - self.numeric_vocab contains:
        #         numeric_vocab["ordered_feature_names"]
        #         numeric_vocab["features"][name]["n_in"]
        #         numeric_vocab["features"][name]["col_id"]
        #
        # If your actual attribute names differ, only this block needs adaptation.
        is_cat = False
        is_num = False
        cat_col = None
        num_n_in = None
        num_v_idx = None
    
        # categorical
        if hasattr(self, "cat_vocab") and feature_name in self.cat_vocab:
            is_cat = True
            cat_col = int(self.cat_vocab[feature_name]["col_id"])
    
        # numeric
        if hasattr(self, "numeric_vocab"):
            num_features = self.numeric_vocab.get("features", {})
            if feature_name in num_features:
                is_num = True
                meta = num_features[feature_name]
                num_n_in = int(meta["n_in"])
                num_v_idx = int(meta["col_id"])
    
        if is_cat and is_num:
            raise ValueError(f"Feature name appears in both categorical and numeric vocab: {feature_name}")
        if not is_cat and not is_num:
            raise KeyError(f"Unknown feature_name: {feature_name}")
    
        # --------------------------------------------------
        # Start from shallow copy only
        # --------------------------------------------------
        masked_batch = dict(batch)
    
        # Preserve original aliases (do NOT deepcopy)
        masked_batch["original_cat_local_ids"] = batch["cat_local_ids"]
        masked_batch["original_cat_valid_positions"] = batch["cat_valid_positions"]
        masked_batch["original_numeric_values_by_nin"] = batch["numeric_values_by_nin"]
        masked_batch["original_numeric_valid_positions_by_nin"] = batch["numeric_valid_positions_by_nin"]
    
        # --------------------------------------------------
        # Default: no masking anywhere
        # --------------------------------------------------
        masked_cat_local_ids = batch["cat_local_ids"].clone()
        masked_cat_valid_positions = batch["cat_valid_positions"].bool().clone()
        cat_loss_mask = torch.zeros((B, M), dtype=torch.bool, device=device)
    
        masked_numeric_values_by_nin = {}
        masked_numeric_valid_positions_by_nin = {}
        numeric_loss_mask_by_nin = {}
    
        for n_in in sorted(numeric_values_by_nin.keys(), key=int):
            values = numeric_values_by_nin[n_in]
            if n_in not in numeric_valid_positions_by_nin:
                raise KeyError(f"Missing numeric_valid_positions_by_nin[{n_in}]")
    
            valid_positions = numeric_valid_positions_by_nin[n_in]
    
            if values.dim() != 3:
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] must be [B, V, n_in], got {tuple(values.shape)}"
                )
    
            Bn, V, Nin = values.shape
            if Bn != B:
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] batch mismatch: got {Bn}, expected {B}"
                )
            if int(Nin) != int(n_in):
                raise ValueError(
                    f"numeric_values_by_nin[{n_in}] last dim mismatch: got {Nin}, expected {n_in}"
                )
            if valid_positions.shape != (B, V):
                raise ValueError(
                    f"numeric_valid_positions_by_nin[{n_in}] must be [B,V]=({B},{V}), "
                    f"got {tuple(valid_positions.shape)}"
                )
    
            masked_numeric_values_by_nin[n_in] = values.clone()
            masked_numeric_valid_positions_by_nin[n_in] = valid_positions.bool().clone()
            numeric_loss_mask_by_nin[n_in] = torch.zeros((B, V), dtype=torch.bool, device=values.device)
    
        # --------------------------------------------------
        # Apply single-feature masking
        # --------------------------------------------------
        if is_cat:
            original_valid = cat_valid_positions[:, cat_col].bool()  # [B]
    
            if assert_not_missing and not bool(original_valid.all().item()):
                n_bad = int((~original_valid).sum().item())
                raise ValueError(
                    f"Categorical feature '{feature_name}' has {n_bad} naturally missing samples in batch"
                )
    
            # only originally valid positions are actively masked
            cat_loss_mask[:, cat_col] = original_valid
    
            masked_cat_local_ids[cat_loss_mask] = cat_mask_local_ids.view(1, M).expand(B, M)[cat_loss_mask]
            masked_cat_valid_positions = masked_cat_valid_positions & (~cat_loss_mask)
    
        else:
            if num_n_in not in masked_numeric_values_by_nin:
                raise KeyError(f"numeric_values_by_nin does not contain n_in={num_n_in} for {feature_name}")
    
            values = masked_numeric_values_by_nin[num_n_in]
            valid_positions = masked_numeric_valid_positions_by_nin[num_n_in]
            num_loss_mask = numeric_loss_mask_by_nin[num_n_in]
    
            if num_v_idx >= values.shape[1]:
                raise IndexError(
                    f"Numeric feature '{feature_name}' resolved to v_idx={num_v_idx}, "
                    f"but numeric_values_by_nin[{num_n_in}] has V={values.shape[1]}"
                )
    
            original_valid = valid_positions[:, num_v_idx].bool()  # [B]
    
            if assert_not_missing and not bool(original_valid.all().item()):
                n_bad = int((~original_valid).sum().item())
                raise ValueError(
                    f"Numeric feature '{feature_name}' has {n_bad} naturally missing samples in batch"
                )
    
            # only originally valid positions are actively masked
            num_loss_mask[:, num_v_idx] = original_valid
    
            values[num_loss_mask] = 0.0
            valid_positions[:] = valid_positions & (~num_loss_mask)
    
        # --------------------------------------------------
        # Finalize outputs
        # --------------------------------------------------
        masked_batch["masked_cat_local_ids"] = masked_cat_local_ids
        masked_batch["masked_cat_valid_positions"] = masked_cat_valid_positions
        masked_batch["cat_loss_mask"] = cat_loss_mask
    
        masked_batch["masked_numeric_values_by_nin"] = masked_numeric_values_by_nin
        masked_batch["masked_numeric_valid_positions_by_nin"] = masked_numeric_valid_positions_by_nin
        masked_batch["numeric_loss_mask_by_nin"] = numeric_loss_mask_by_nin
    
        return masked_batch


def build_train_eval_dataloaders(
        dataset,
        train_ratio=0.8,
        seed=42,
        batch_size=32,
):
    n = len(dataset)

    n_train = int(n * train_ratio)
    n_eval = n - n_train

    split_generator = torch.Generator().manual_seed(seed)

    train_ds, eval_ds = torch.utils.data.random_split(
        dataset,
        [n_train, n_eval],
        generator=split_generator
    )

    train_generator = torch.Generator()

    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=dataset.collate_fn,
        generator=train_generator,
    )

    eval_loader = DataLoader(
        eval_ds,
        batch_size=batch_size,
        shuffle=False,
        collate_fn=dataset.collate_fn,
    )

    return train_loader, eval_loader, train_generator


def debug_print_first_sample(dataset, batch, batch_pos=0):
    """
    Inspect one sample in a batch.

    This debug function checks masked_* fields against the original csv row.
    Positions in loss_mask are allowed to mismatch.

    Args:
        dataset: SoilFormerDataset
        batch: collated + optionally masked batch
        batch_pos: index inside the batch (not dataset row index)
    """
    import math

    def numeric_list_close(a, b, atol=1e-6, rtol=1e-5):
        if len(a) != len(b):
            return False
        for x, y in zip(a, b):
            if not math.isclose(float(x), float(y), rel_tol=rtol, abs_tol=atol):
                return False
        return True

    def normalize_numeric_list(feat_name, vals, is_valid):
        if not is_valid:
            return [0.0] * len(vals)

        stat_row = dataset.numeric_stats_index.loc[feat_name]
        mean = float(stat_row["mean"])
        std = float(stat_row["std"])
        if std == 0.0:
            std = 1.0

        return [(float(v) - mean) / std for v in vals]

    if "row_idx" not in batch:
        raise KeyError("batch must contain 'row_idx' for debug_print_first_sample")
    if "sample_id" not in batch:
        raise KeyError("batch must contain 'sample_id' for debug_print_first_sample")

    row_idx = int(batch["row_idx"][batch_pos].item())
    row = dataset.df.iloc[row_idx]
    sample_id = batch["sample_id"][batch_pos]

    print("\n====================================================")
    print("DEBUG SAMPLE")
    print("====================================================")
    print("batch_pos :", batch_pos)
    print("row_idx   :", row_idx)
    print("sample_id :", sample_id)

    # ====================================================
    # categorical
    # ====================================================
    print("\n[CATEGORICAL FEATURES]")

    cat_ids = batch["masked_cat_local_ids"][batch_pos]
    cat_valids = batch["masked_cat_valid_positions"][batch_pos]
    cat_loss_mask = batch.get("cat_loss_mask", None)
    if cat_loss_mask is not None:
        cat_loss_mask = cat_loss_mask[batch_pos]

    for i, col in enumerate(dataset.cat_columns):
        raw = row[col]
        raw_str = str(raw)

        got_id = int(cat_ids[i].item())
        got_valid = bool(cat_valids[i].item())

        spec = dataset.cat_vocab[col]
        label2id = spec["label2id"]
        mask_id = int(spec["mask_local_id"])

        if raw == "":
            expected_id = mask_id
            expected_valid = False
        else:
            expected_id = int(label2id[raw])
            expected_valid = True

        is_loss_position = False
        if cat_loss_mask is not None:
            is_loss_position = bool(cat_loss_mask[i].item())

        if is_loss_position:
            ok = True
        else:
            ok = (got_id == expected_id) and (got_valid == expected_valid)

        print(
            f"{i:03d} | {col} | "
            f"raw={raw_str:<60} | "
            f"id={got_id:<6} | expected={expected_id:<6} | "
            f"valid={got_valid} | exp_valid={expected_valid} | "
            f"loss_mask={is_loss_position} | ok={ok}"
        )

        if not ok:
            raise AssertionError(
                f"\nCategorical mismatch\n"
                f"batch_pos={batch_pos}\n"
                f"row_idx={row_idx}\n"
                f"feature={col}\n"
                f"raw={raw}\n"
                f"id={got_id}, expected={expected_id}\n"
                f"valid={got_valid}, expected={expected_valid}"
            )

    # ====================================================
    # numeric
    # ====================================================
    print("\n[NUMERIC FEATURES]")

    numeric_loss_mask_by_nin = batch.get("numeric_loss_mask_by_nin", None)

    for group in dataset.numeric_groups:
        n_in = int(group["n_in"])
        features = group["feature_names"]

        values = batch["masked_numeric_values_by_nin"][n_in][batch_pos]
        valids = batch["masked_numeric_valid_positions_by_nin"][n_in][batch_pos]

        if numeric_loss_mask_by_nin is not None:
            loss_mask = numeric_loss_mask_by_nin[n_in][batch_pos]
        else:
            loss_mask = None

        print(f"\nGroup n_in={n_in}")

        for i, feat in enumerate(features):
            raw = row[feat]
            raw_str = str(raw)

            parsed, expected_valid = parse_numeric_cell(raw, n_in)
            expected_norm = normalize_numeric_list(feat, parsed, expected_valid)

            tensor_val = values[i].tolist()
            got_valid = bool(valids[i].item())

            is_loss_position = False
            if loss_mask is not None:
                is_loss_position = bool(loss_mask[i].item())

            if is_loss_position:
                ok = True
            else:
                value_ok = numeric_list_close(tensor_val, expected_norm)
                valid_ok = (got_valid == expected_valid)
                ok = value_ok and valid_ok

            print(
                f"{i:03d} | {feat} | "
                f"raw={raw_str:<60} | "
                f"tensor={tensor_val} | expected_norm={expected_norm} | "
                f"valid={got_valid} | exp_valid={expected_valid} | "
                f"loss_mask={is_loss_position} | ok={ok}"
            )

            if not ok:
                raise AssertionError(
                    f"\nNumeric mismatch\n"
                    f"batch_pos={batch_pos}\n"
                    f"row_idx={row_idx}\n"
                    f"feature={feat}\n"
                    f"raw={raw}\n"
                    f"tensor={tensor_val}\n"
                    f"expected={parsed}\n"
                    f"valid={got_valid}, expected={expected_valid}"
                )

    # ====================================================
    # vision
    # ====================================================
    print("\n[VISION]")

    try:
        relative_path = dataset.photo_map[sample_id]
        expected_path = join_photo_root(dataset.photo_root, relative_path)

        # Use the same logic as __getitem__: valid only if image can actually be loaded
        _ = dataset.load_image(expected_path)
        expected_valid = True

    except Exception:  # noqa
        expected_path = None
        expected_valid = False

    got_valid = bool(batch["vision_valid_positions"][batch_pos].item())
    img_shape = tuple(batch["pixel_values"][batch_pos].shape)

    print("expected_path :", expected_path)
    print("vision_valid  :", got_valid)
    print("image_shape   :", img_shape)

    if got_valid != expected_valid:
        raise AssertionError(
            f"\nVision validity mismatch\n"
            f"batch_pos={batch_pos}\n"
            f"row_idx={row_idx}\n"
            f"expected={expected_valid}, got={got_valid}"
        )

    print("\n====================================================")
    print("DEBUG CHECK PASSED")
    print("====================================================\n")


def main():
    dataset = SoilFormerDataset(
        csv_path="data/tabular_data.csv",
        photo_map_path="data/photo_map.json",
        cat_vocab_path="data/cat_vocab.json",
        numeric_vocab_path="data/numeric_vocab.json",
        numeric_stats_path="data/tabular_meta_numeric_stats.csv",
        photo_root="/Volumes/TOSHIBA EXT",
        image_size=512,
        id_column="id",
    )

    train_loader, eval_loader, train_generator = build_train_eval_dataloaders(dataset)

    print("Dataset size:", len(dataset))

    raw_batch = next(iter(eval_loader))
    batch = dataset.perform_active_mask(
        raw_batch,
        cat_ratio=0.15,
        num_ratio=0.15,
        seed=42,
    )

    print("\nBatch check")
    if "row_idx" in batch:
        print("row_idx:", batch["row_idx"].shape, batch["row_idx"].dtype)
    if "sample_id" in batch:
        print("sample_id:", len(batch["sample_id"]))

    print("original_cat_local_ids:", batch["original_cat_local_ids"].shape)
    print("masked_cat_local_ids:", batch["masked_cat_local_ids"].shape)
    print("original_cat_valid_positions:", batch["original_cat_valid_positions"].shape)
    print("masked_cat_valid_positions:", batch["masked_cat_valid_positions"].shape)
    print("cat_loss_mask:", batch["cat_loss_mask"].shape)

    for k, v in batch["original_numeric_values_by_nin"].items():
        print(f"original_numeric_values_by_nin[{k}]:", v.shape)

    for k, v in batch["masked_numeric_values_by_nin"].items():
        print(f"masked_numeric_values_by_nin[{k}]:", v.shape)

    for k, v in batch["original_numeric_valid_positions_by_nin"].items():
        print(f"original_numeric_valid_positions_by_nin[{k}]:", v.shape)

    for k, v in batch["masked_numeric_valid_positions_by_nin"].items():
        print(f"masked_numeric_valid_positions_by_nin[{k}]:", v.shape)

    for k, v in batch["numeric_loss_mask_by_nin"].items():
        print(f"numeric_loss_mask_by_nin[{k}]:", v.shape)

    print("pixel_values:", batch["pixel_values"].shape)
    print("vision_valid_positions:", batch["vision_valid_positions"].shape)

    print("\nTensor dtype check")
    print("masked cat ids dtype:", batch["masked_cat_local_ids"].dtype)
    print("masked numeric dtype:", next(iter(batch["masked_numeric_values_by_nin"].values())).dtype)
    print("image dtype:", batch["pixel_values"].dtype)

    print("\nLoader test finished successfully")

    debug_print_first_sample(dataset, batch, batch_pos=0)


if __name__ == "__main__":
    main()