utils.py

from typing import (
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Type,
    Union,
)

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal

import logging

import torch
from torchvision.datasets import ImageFolder
from torchvision.utils import make_grid
from torchvision.transforms.functional import to_pil_image
from PIL import ImageDraw, ImageFont
import textwrap
import random
from torchvision.datasets.folder import default_loader  # same one ImageFolder uses
from torch.utils.data import Sampler
import torch.distributed as dist

import os, sys, fcntl, subprocess, time

import hashlib

from contextlib import contextmanager


import numpy as np
import torchvision.transforms as T
from torchvision.transforms import functional as F


def normalize_to_0_1(x: torch.Tensor, dim=(1, 2, 3)) -> torch.Tensor:
    x_min = x.amin(dim=dim, keepdim=True)
    x_max = x.amax(dim=dim, keepdim=True)

    x_norm = (x - x_min) / (x_max - x_min + 1e-8)
    return x_norm


def randint_unique_once(
    low: int,
    high: int,
    size: Tuple[int, ...],
    cur_k: int,
    *,
    device=None,
    dtype=torch.long,
) -> Tuple[bool, torch.Tensor]:
    """
    Draw with torch.randint(..., size) but ensure the LAST dimension is unique.
    • Works for any batch shape (..., k)   (k ≤ high-low)
    • Does a single replacement pass; never uses randperm or re-draw loops.
    """
    size = tuple(size)
    k = size[-1]
    n = high - low
    assert k <= n, "Need range ≥ k to make elements unique"

    # 1) initial draw — duplicates possible
    out = torch.randint(low, high, size, device=device, dtype=dtype)

    if out.numel() == 0:  # edge-case: empty tensor
        return False, out

    # 2) flatten batch dims so we can iterate row-wise in Python
    batch_view = out.view(-1, k)  # shape: (B, k) where B = prod(size[:-1])

    has_dup = False

    for row in batch_view:  # simple Python loop over rows
        # ---- find duplicates (everything after the first occurrence) -------
        # make a Bool mask that's True for dup positions
        dup_mask = torch.zeros_like(row, dtype=torch.bool)
        seen = {}
        for idx, val in enumerate(row):
            if val.item() in seen:
                dup_mask[idx] = True  # second, third, … occurrence
                has_dup = has_dup or idx < cur_k
            else:
                seen[val.item()] = idx

        if not dup_mask.any():  # row already unique
            continue

        # ---- pool of free ints in [low, high) NOT yet in the row -----------
        taken = torch.zeros(n, dtype=torch.bool, device=device)
        taken[row - low] = True  # shift so 0…n-1
        free_vals = torch.arange(low, high, device=device)[~taken]

        # sample exactly as many as duplicates, WITH torch.randint
        num_to_fill = dup_mask.sum().item()
        rand_idx = torch.randint(0, free_vals.numel(), (num_to_fill,), device=device)
        row[dup_mask] = free_vals[rand_idx]  # in-place replace duplicates

    return has_dup, out


def randint_no_replacement(
    low: int,
    high: int,
    size: Tuple[int, ...],
    *,
    device=None,
    dtype=torch.long,
) -> torch.Tensor:
    """
    Return a tensor of `size` whose last dim contains UNIQUE ints in [low, high).

    Example:
        >>> torch.manual_seed(0)
        >>> randint_no_replacement(0, 10, (4, 3))
        tensor([[2, 0, 6],
                [4, 5, 1],
                [7, 8, 9],
                [3, 2, 5]])
    """
    *batch_shape, k = size
    n = high - low
    assert k <= n, "Cannot sample more unique numbers than the range"

    # Step 1: make random keys with shape (*batch, n)
    keys = torch.rand(*batch_shape, n, device=device)

    # Step 2: argsort gives a per-row permutation 0..n-1
    perm = keys.argsort(dim=-1)

    # Step 3: slice the first k columns and shift by `low`
    return perm[..., :k].to(dtype) + low


def randint_from_pool(low: int, high: int, size: int, exclude: List[int], device):
    pool = [i for i in range(low, high) if i not in exclude]
    indices = torch.randperm(len(pool))[:size]
    return torch.tensor([pool[i] for i in indices], dtype=torch.long, device=device)


def save_image_with_text(
    images,
    labels,
    filename,
    nrow=4,
    font_size=80,
    normalize=False,
    text_per_line=5,
):
    """
    images: Tensor of shape (B, C, H, W)
    labels: List of strings, one per image
    filename: Path to save the final image
    """
    B = images.size(0)
    nrow = nrow
    padding = 10
    # Create grid image
    grid = make_grid(images, nrow=nrow, padding=10, normalize=normalize)
    pil_img = to_pil_image(grid)

    # Draw text
    draw = ImageDraw.Draw(pil_img)

    # Try to load a TTF font; fallback if not available
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except IOError:
        font = ImageFont.load_default(size=font_size)

    # Size and layout info

    img_h = images.size(2)
    img_w = images.size(3)

    for idx, label in enumerate(labels):
        if label == "":
            continue

        # wrap text
        lines = textwrap.wrap(label, width=text_per_line)

        row = idx // nrow
        col = idx % nrow
        x = col * (img_w + padding) + padding
        y = row * (img_h + padding) + padding

        for line in lines:
            # Get size of text to draw background rectangle
            bbox = draw.textbbox((0, 0), line, font=font)
            text_w = bbox[2] - bbox[0]
            text_h = bbox[3] - bbox[1]
            rect_coords = [(x - 2, y - 2), (x + text_w + 2, y + text_h + 6)]

            # Draw black rectangle
            draw.rectangle(rect_coords, fill="black")

            # Draw text over it
            draw.text((x, y), line, font=font, fill="yellow")
            y += text_h + 10

    pil_img.save(filename)


class ImageFolderWithPaths(ImageFolder):
    def __init__(self, base: ImageFolder):
        # Manually copy relevant fields
        self.root = base.root
        self.loader = base.loader
        self.extensions = base.extensions
        self.classes = base.classes
        self.class_to_idx = base.class_to_idx
        self.samples = base.samples
        self.targets = base.targets
        self.transform = base.transform
        self.target_transform = base.target_transform
        self.imgs = base.imgs  # same as samples

    def __getitem__(self, index):
        sample, target = super(ImageFolder, self).__getitem__(index)
        path = self.samples[index][0]
        return sample, target, path, index


class ImageFolderWithFilenameFilter(ImageFolder):
    def __init__(self, base: ImageFolder, selected_filenames: List[str]):

        self.root = base.root
        self.loader = base.loader
        self.extensions = base.extensions
        self.classes = base.classes
        self.class_to_idx = base.class_to_idx
        self.samples = base.samples
        self.targets = base.targets
        self.transform = base.transform
        self.target_transform = base.target_transform
        self.imgs = base.imgs  # same as samples

        self.selected_filenames = set(selected_filenames)

        # Filter samples that match selected filenames
        self.filtered_samples = [
            (path, label, idx)
            for idx, (path, label) in enumerate(self.samples)
            if os.path.basename(path) in self.selected_filenames
        ]

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

    def __getitem__(self, idx):
        path, label, idx_in_orig_samples = self.filtered_samples[idx]
        img = self.loader(path)
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            label = self.target_transform(label)
        return img, label, path, idx_in_orig_samples


class KPerClassSampler(Sampler):
    """
    ex
    dataset = ImageFolder("/imagenet/val", transform=my_tfm)
    sampler = KPerClassSampler(dataset, k=4, shuffle=True)
    loader  = DataLoader(dataset,
                        batch_size=64,
                        sampler=sampler,
                        num_workers=8,
                        pin_memory=True)

    from timm.data.loader import create_loader, DistributedSamplerWrapper

    base_sampler = KPerClassSampler(dataset, k=4, shuffle=True)
    dist_sampler = DistributedSamplerWrapper(base_sampler,
                                            num_replicas=world_size,
                                            rank=rank,
                                            shuffle=True)

    loader = DataLoader(dataset, batch_size=64, sampler=dist_sampler, ...)
    """

    def __init__(self, img_folder: ImageFolder, k: int = 3, shuffle=True):
        self.k = k
        self.shuffle = shuffle

        # class_id → indices
        self.cls2idx = {}
        for idx, (_, c) in enumerate(img_folder.samples):
            self.cls2idx.setdefault(c, []).append(idx)

        self.num_samples = len(self.cls2idx) * k

    def __iter__(self):
        classes = list(self.cls2idx.keys())
        if self.shuffle:
            random.shuffle(classes)

        for c in classes:
            pool = self.cls2idx[c]
            # choose k indices **without replacement**
            chosen = (
                random.sample(pool, self.k)
                if len(pool) >= self.k
                else random.choices(pool, k=self.k)
            )
            # optionally shuffle the k within-class indices
            if self.shuffle:
                random.shuffle(chosen)
            yield from chosen

    def __len__(self):
        return self.num_samples


class KPerClassFolder:
    """
    Wrapper around ImageFolder that selects K fixed samples per class using
    md5 hashing of image paths for deterministic and cross-machine consistency.

    Returns (img, label, path, sample_index).
    """

    def __init__(
        self,
        base: ImageFolder,
        k: int = 1,
        max_samples: Optional[int] = None,
        selected_filenames: Optional[List[str]] = None,
        return_path_and_idx: bool = True,
        deterministic_transform: bool = False,
    ):
        self.root = base.root
        self.loader = base.loader
        self.extensions = base.extensions
        self.classes = base.classes
        self.class_to_idx = base.class_to_idx
        self.samples = base.samples
        self.targets = base.targets
        self.transform = base.transform
        self.target_transform = base.target_transform
        self.imgs = base.imgs  # same as samples

        self.return_path_and_idx = return_path_and_idx
        self.deterministic_transform = deterministic_transform

        self.k = k
        self.max_samples = max_samples
        self.selected_filenames = (
            set(selected_filenames) if selected_filenames else None
        )

        # Build class → list of (idx, path) map
        self.cls2idx: Dict[int, List[int]] = {}
        for idx, (path, cls) in enumerate(self.samples):
            fname = os.path.basename(path)
            if self.selected_filenames is None or fname in self.selected_filenames:
                self.cls2idx.setdefault(cls, []).append(idx)

        # Use hash-based sorting to ensure consistency
        self.fixed_cls2idx: Dict[int, List[int]] = {}
        for cls, idxs in self.cls2idx.items():
            paths = [self.samples[i][0] for i in idxs]
            top_k = min(len(idxs), self.k)
            self.fixed_cls2idx[cls] = self._select_k_by_hash(paths, idxs)[:top_k]

        self.classes_sorted = sorted(self.fixed_cls2idx)

    def _select_k_by_hash(self, paths: List[str], idxs: List[int]) -> List[int]:
        """Sort indices based on md5 hash of their path and select the top-k."""
        hashed = sorted(
            zip(paths, idxs),
            key=lambda x: hashlib.md5(x[0].encode("utf-8")).hexdigest(),
        )
        return [idx for _, idx in hashed[: self.k]]

    def __len__(self) -> int:
        total = sum(len(v) for v in self.fixed_cls2idx.values())
        return min(self.max_samples, total) if self.max_samples else total

    def set_epoch(self, epoch):
        if self.transform and hasattr(self.transform, "set_epoch"):
            self.transform.set_epoch(epoch)

    def __getitem__(self, global_idx: int) -> Tuple:
        if global_idx >= len(self):
            raise IndexError(
                f"Index {global_idx} out of range for dataset of length {len(self)}."
            )

        counter = 0
        for cls_id in self.classes_sorted:
            sample_list = self.fixed_cls2idx[cls_id]
            count = len(sample_list)

            if global_idx < counter + count:
                local_idx = global_idx - counter
                sample_idx = sample_list[local_idx]
                break

            counter += count

        path, target = self.samples[sample_idx]
        img = self.loader(path)
        if self.transform:
            img = (
                self.transform(img, sample_idx)
                if self.deterministic_transform
                else self.transform(img)
            )
        if self.target_transform:
            target = self.target_transform(target)

        if self.return_path_and_idx:
            return img, target, path, sample_idx

        return img, target


class IndexedImageFolder(ImageFolder):
    def __init__(self, base: ImageFolder, deterministic_transform: bool = False):
        self.root = base.root
        self.loader = base.loader
        self.extensions = base.extensions
        self.classes = base.classes
        self.class_to_idx = base.class_to_idx
        self.samples = base.samples
        self.targets = base.targets
        self.transform = base.transform
        self.target_transform = base.target_transform
        self.imgs = base.imgs  # same as samples

        self.deterministic_transform = deterministic_transform

    def set_epoch(self, epoch):
        if self.transform and hasattr(self.transform, "set_epoch"):
            self.transform.set_epoch(epoch)

    def __getitem__(self, index):
        path, target = self.samples[index]
        img = self.loader(path)
        if self.transform:
            img = (
                self.transform(img, index)
                if self.deterministic_transform
                else self.transform(img)
            )
        if self.target_transform:
            target = self.target_transform(target)

        return img, target


class DeterministicCompose:
    def __init__(self, transforms: Sequence[Callable]):
        self.transforms = transforms
        self.epoch = 0

    def set_epoch(self, epoch):
        self.epoch = epoch

    def __call__(self, img, index):
        seed = hash((self.epoch, index)) % (2**32)
        rng_python = random.Random(seed)
        rng_numpy = np.random.default_rng(seed)
        rng_torch = torch.Generator().manual_seed(seed)

        for t in self.transforms:
            if isinstance(t, T.RandomResizedCrop):
                i, j, h, w = T.RandomResizedCrop.get_params(
                    img, scale=t.scale, ratio=t.ratio, rng=rng_python
                )
                img = F.resized_crop(
                    img, i, j, h, w, t.size, t.interpolation, t.antialias
                )
            elif isinstance(t, T.RandomHorizontalFlip):
                if rng_python.random() < t.p:
                    img = F.hflip(img)
            elif isinstance(t, T.ColorJitter):
                (
                    fn_idx,
                    brightness_factor,
                    contrast_factor,
                    saturation_factor,
                    hue_factor,
                ) = T.ColorJitter.get_params(
                    t.brightness, t.contrast, t.saturation, t.hue
                )
                for fn_id in fn_idx:
                    if fn_id == 0 and brightness_factor is not None:
                        img = F.adjust_brightness(img, brightness_factor)
                    elif fn_id == 1 and contrast_factor is not None:
                        img = F.adjust_contrast(img, contrast_factor)
                    elif fn_id == 2 and saturation_factor is not None:
                        img = F.adjust_saturation(img, saturation_factor)
                    elif fn_id == 3 and hue_factor is not None:
                        img = F.adjust_hue(img, hue_factor)
            elif isinstance(t, T.ToTensor):
                img = F.to_tensor(img)
            elif isinstance(t, T.Normalize):
                img = F.normalize(img, mean=t.mean, std=t.std)
            else:
                raise NotImplementedError(
                    f"Transform {type(t)} not supported in DeterministicCompose."
                )
        return img


def setup_logging(log_file=None, log_level=logging.INFO):
    is_main_process = (
        not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
    )

    logging.basicConfig(
        level=log_level,
        format="%(asctime)s | %(levelname)s | %(message)s",
        handlers=[
            logging.StreamHandler(sys.stdout),
        ],
    )

    if is_main_process and log_file:
        file_handler = logging.FileHandler(log_file, mode="a")
        file_handler.setFormatter(
            logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")
        )
        logging.getLogger().addHandler(file_handler)


def cache_rng_state():
    return {
        "torch": torch.get_rng_state(),
        "cuda": torch.cuda.get_rng_state_all(),
        "numpy": np.random.get_state(),
        "python": random.getstate(),
    }


def restore_rng_state(state):
    torch.set_rng_state(state["torch"])
    torch.cuda.set_rng_state_all(state["cuda"])
    np.random.set_state(state["numpy"])
    random.setstate(state["python"])