datasets.py

import os
import pickle
import json
import requests
from collections import defaultdict, deque
from typing import Optional, List, Tuple, Dict, Any
from tqdm import tqdm
from io import BytesIO
from PIL import Image
from pathlib import Path
from urllib3 import disable_warnings
from urllib3.exceptions import InsecureRequestWarning

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder

import torch
from torch.utils.data import Dataset, Subset, DataLoader, TensorDataset
from torchvision import transforms
from huggingface_hub import hf_hub_download
from datasets import load_dataset

from experiments.datasets import (
    LiverTumorDataset,
    LiverTumorDatasetHf,
    AKIDataset,
    AKI_COLUMNS,
)


class ImageNetDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.img_dir = os.path.join(self.root_dir, "samples/")
        self.label_dir = os.path.join(self.root_dir, "imagenet_class_index.json")

        with open(self.label_dir) as json_data:
            self.idx_to_labels = json.load(json_data)

        self.img_names = os.listdir(self.img_dir)
        self.img_names.sort()

        self.transform = transform

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_names[idx])
        image = Image.open(img_path).convert("RGB")
        label = idx

        if self.transform:
            image = self.transform(image)

        return image, label

    def idx_to_label(self, idx):
        return self.idx_to_labels[str(idx)][1]


def get_imagenet_dataset(
    transform,
    subset_size: int = 100,  # ignored if indices is not None
    root_dir="./data/ImageNet",
    indices: Optional[List[int]] = None,
):
    os.chdir(Path(__file__).parent)  # ensure path
    dataset = ImageNetDataset(root_dir=root_dir, transform=transform)
    if indices is not None:
        return Subset(dataset, indices=indices)
    indices = list(range(len(dataset)))
    subset = Subset(dataset, indices=indices[:subset_size])
    return subset


class ImageNetValDataset(Dataset):
    def __init__(self, img_dir, label_file, class_index_file, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.image_files = sorted(
            [f for f in os.listdir(img_dir) if f.endswith(".JPEG")]
        )

        with open(label_file, "r") as f:
            self.labels = [line.strip() for line in f]

        with open(class_index_file, "r") as f:
            self.class_index = json.load(f)
        self.synset_to_idx = {v[0]: int(k) for k, v in self.class_index.items()}

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

    def __getitem__(self, idx):
        img_name = self.image_files[idx]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert("RGB")
        label_synset = self.labels[idx].split(" ")[-1]
        label = self.synset_to_idx[label_synset]

        if self.transform:
            image = self.transform(image)

        return image, label


def get_imagenet_val_dataset(transform, root_dir):
    img_dir = os.path.join(root_dir, "ImageNet1k", "val", "val")
    info_dir = os.path.join(root_dir, "ImageNet1k_info")
    val_label_file = os.path.join(info_dir, "ImageNet_val_label.txt")
    val_class_index_file = os.path.join(info_dir, "ImageNet_class_index.json")
    dataset = ImageNetValDataset(
        img_dir, val_label_file, val_class_index_file, transform
    )
    return dataset


class IMDBDataset(Dataset):
    def __init__(self, split="test"):
        super().__init__()
        # data_iter = IMDB(split=split)
        # self.annotations = [(line, label-1) for label, line in tqdm(data_iter)]

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

    def __getitem__(self, idx):
        return self.annotations[idx]


def get_imdb_dataset(split="test"):
    return IMDBDataset(split=split)


disable_warnings(InsecureRequestWarning)


class VQADataset(Dataset):
    def __init__(self):
        super().__init__()
        res = requests.get("https://visualqa.org/balanced_data.json")
        self.annotations = eval(res.text)

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

    def __getitem__(self, idx):
        data = self.annotations[idx]
        if isinstance(data["original_image"], str):
            print(f"Requesting {data['original_image']}...")
            res = requests.get(data["original_image"], verify=False)
            img = Image.open(BytesIO(res.content)).convert("RGB")
            data["original_image"] = img
        return data["original_image"], data["question"], data["original_answer"]


def get_vqa_dataset():
    return VQADataset()


def get_livertumor_dataset(
    transform,
    subset_size: int = 100,  # ignored if indices is not None
    root_dir="./data/LiverTumor",
    indices: Optional[List[int]] = None,
):
    dataset = LiverTumorDataset(data_dir=root_dir, transform=transform)
    if indices is not None:
        return Subset(dataset, indices=indices)
    indices = list(range(len(dataset)))
    subset = Subset(dataset, indices=indices[:subset_size])
    return subset


def get_livertumor_dataset_from_hf(
    transform,
    hf_repo_id: str = "seongun/liver-tumor-classification",
    indices: Optional[List[int]] = None,
    data_root: str = "./data",
    cache_dir: Optional[str] = None,
):
    """
    Downloads only necessary files (metadata + images for requested indices)
    from Hugging Face Hub using hf_hub_download and creates a PyTorch Dataset.

    Args:
        transform: Torchvision transforms to apply to the image-like data.
        hf_repo_id (str): Repository ID of the dataset on Hugging Face Hub.
        indices (Optional[List[int]]): Absolute indices to select/download.
        data_root (str): The root directory within the project to store datasets.
        cache_dir (Optional[str]): Path to HF cache (used for intermediate downloads).

    Returns:
        A PyTorch Dataset containing only the data for the requested indices.
    """
    if indices is None:
        print(
            "Warning: No indices provided. Attempting to load metadata only, but image loading might fail later or be inefficient."
        )

    dataset_local_dir = os.path.join(data_root, hf_repo_id.replace("/", "_"))
    os.makedirs(dataset_local_dir, exist_ok=True)

    print(f"Downloading metadata for '{hf_repo_id}' from Hugging Face Hub...")
    try:
        # 1. Download metadata.jsonl only
        metadata_local_path = hf_hub_download(
            repo_id=hf_repo_id,
            filename="metadata.jsonl",
            repo_type="dataset",
            local_dir=dataset_local_dir,
            local_dir_use_symlinks=True,
            cache_dir=cache_dir,
        )
        base_download_dir = dataset_local_dir
        print(f"Metadata available at: {metadata_local_path}")
        print(f"Base download/cache directory: {base_download_dir}")

    except Exception as e:
        print(f"Failed to download metadata.jsonl from Hugging Face Hub: {e}")
        raise e

    # 2. Read metadata and filter for requested indices
    filtered_metadata = []
    required_image_paths = set()  # Use set to avoid duplicate downloads
    all_metadata = []
    try:
        with open(metadata_local_path, "r") as f:
            all_metadata = [json.loads(line.strip()) for line in f]

        if indices is not None:
            num_total = len(all_metadata)
            for idx in indices:
                if 0 <= idx < num_total:
                    entry = all_metadata[idx]
                    filtered_metadata.append(entry)
                    required_image_paths.add(entry["sample_path"])
                    required_image_paths.add(entry["w_sample_path"])
                    required_image_paths.add(entry["mask_path"])
                else:
                    print(
                        f"Warning: Requested index {idx} is out of range (0-{num_total-1}). Skipping."
                    )
        else:
            print(
                "Warning: Loading without specific indices. Using all metadata entries."
            )
            filtered_metadata = (
                all_metadata  # Less efficient if not all images are needed later
            )

    except Exception as e:
        print(f"Error reading or processing metadata file {metadata_local_path}: {e}")
        raise e

    if not filtered_metadata:
        raise ValueError("No valid metadata found for the requested indices.")

    # 3. Download only the required images
    print(
        f"Downloading {len(required_image_paths)} required image files (if not cached)..."
    )
    for img_rel_path in tqdm(list(required_image_paths), desc="Downloading images"):
        try:
            # hf_hub_download will download to the cache or find existing file
            hf_hub_download(
                repo_id=hf_repo_id,
                filename=img_rel_path,
                repo_type="dataset",
                local_dir=dataset_local_dir,
                local_dir_use_symlinks=True,
                cache_dir=cache_dir,
            )
        except Exception as e:
            print(f"Warning: Failed to download image file {img_rel_path}: {e}")

    print("Required image files downloaded/cached.")

    # 4. Create and return the Dataset using filtered metadata and base download dir
    dataset = LiverTumorDatasetHf(
        metadata=filtered_metadata,
        base_download_dir=base_download_dir,
        transform=transform,
    )

    print(f"Created dataset with {len(dataset)} instances.")
    return dataset


def get_imagenet_sample_from_hf(
    transform,
    hf_repo_id: str = "geonhyeongkim/imagenet-samples-for-pnpxai-experiments",
    indices: Optional[List[int]] = None,
    data_root: str = "./data",
    cache_dir: Optional[str] = None,
):
    """
    Downloads only necessary files (metadata + images for requested indices)
    from Hugging Face Hub using hf_hub_download and creates a PyTorch Dataset.

    Args:
        transform: Torchvision transforms to apply to the image-like data.
        hf_repo_id (str): Repository ID of the dataset on Hugging Face Hub.
        indices (Optional[List[int]]): Absolute indices to select/download.
        data_root (str): The root directory within the project to store datasets.
        cache_dir (Optional[str]): Path to HF cache (used for intermediate downloads).

    Returns:
        A PyTorch Dataset containing only the data for the requested indices.
    """
    if indices is None:
        print(
            "Warning: No indices provided. Attempting to load metadata only, but image loading might fail later or be inefficient."
        )

    dataset_local_dir = os.path.join(data_root, hf_repo_id.replace("/", "_"))
    os.makedirs(dataset_local_dir, exist_ok=True)

    print(f"Downloading metadata for '{hf_repo_id}' from Hugging Face Hub...")
    try:
        # 1. Download metadata.jsonl only
        metadata_local_path = hf_hub_download(
            repo_id=hf_repo_id,
            filename="imagenet_class_index.json",
            repo_type="dataset",
            local_dir=dataset_local_dir,
            local_dir_use_symlinks=True,
            cache_dir=cache_dir,
        )
        base_download_dir = dataset_local_dir
        print(f"Metadata available at: {metadata_local_path}")
        print(f"Base download/cache directory: {base_download_dir}")

    except Exception as e:
        print(
            f"Failed to download imagenet_class_index.json from Hugging Face Hub: {e}"
        )
        raise e

    # 2. Read metadata and filter for requested indices
    filtered_metadata = {}
    required_image_paths = set()  # Use set to avoid duplicate downloads
    all_metadata = {}
    try:
        with open(metadata_local_path, "r") as f:
            all_metadata = json.load(f)

        if indices is not None:
            num_total = len(all_metadata)
            for idx in indices:
                if 0 <= idx < num_total:
                    metadata = all_metadata[str(idx)]
                    filtered_metadata[idx] = metadata
                    required_image_paths.add(f'samples/{"_".join(metadata)}.JPEG')
                else:
                    print(
                        f"Warning: Requested index {idx} is out of range (0-{num_total-1}). Skipping."
                    )
        else:
            print(
                "Warning: Loading without specific indices. Using all metadata entries."
            )
            filtered_metadata = {
                int(k): all_metadata[k] for k in all_metadata
            }  # Less efficient if not all images are needed later

    except Exception as e:
        print(f"Error reading or processing metadata file {metadata_local_path}: {e}")
        raise e

    if not filtered_metadata:
        raise ValueError("No valid metadata found for the requested indices.")

    # 3. Download only the required images
    print(
        f"Downloading {len(required_image_paths)} required image files (if not cached)..."
    )
    for img_rel_path in tqdm(list(required_image_paths), desc="Downloading images"):
        try:
            # hf_hub_download will download to the cache or find existing file
            hf_hub_download(
                repo_id=hf_repo_id,
                filename=img_rel_path,
                repo_type="dataset",
                local_dir=dataset_local_dir,
                local_dir_use_symlinks=True,
                cache_dir=cache_dir,
            )
        except Exception as e:
            print(f"Warning: Failed to download image file {img_rel_path}: {e}")

    print("Required image files downloaded/cached.")

    # 4. Create and return the Dataset using filtered metadata and base download dir
    fp_img = os.path.join(base_download_dir, list(required_image_paths)[0])
    img = transform(Image.open(fp_img).convert("RGB"))
    label = all_metadata[str(indices[0])][-1]
    return img, label


def get_aki_dataset(
    data_path: str = "data/mimiciii/formatted/data.csv",
    test_split: float = 0.2,
) -> AKIDataset:
    data = pd.read_csv(data_path)

    data = data.replace([np.inf, -np.inf], np.nan).dropna()
    data = data[AKI_COLUMNS]

    scaler = StandardScaler()
    scaler.fit(data.iloc[:, 2:])
    data.iloc[:, 2:] = scaler.transform(data.iloc[:, 2:])

    n_entries = len(data)
    df_test = data.iloc[-int(test_split * n_entries) :, :]

    x_data = df_test.drop(["AKI_STAGE_7DAY", "AKI"], axis=1).values
    y_data = df_test["AKI_STAGE_7DAY"].values

    dataset = AKIDataset(x_data, y_data)

    return dataset


def get_ecg_dataset_from_hf(repo_id: str = "enver1323/ucr-twoleadecg") -> TensorDataset:
    data = load_dataset(repo_id)["test"].with_format("numpy")
    x_data = np.stack(data['segment'])
    y_data = data['label']

    encoder = LabelEncoder()
    y_data = encoder.fit_transform(y_data)

    return TensorDataset(
        torch.from_numpy(x_data),
        torch.from_numpy(y_data)
    )


def get_winequality_dataset(data_dir: str = "data/wine_quality") -> Tuple:
    """Load Wine Quality dataset and feature metadata."""
    data_path = Path(data_dir)

    X_train = np.load(data_path / "X_train.npy")
    X_test = np.load(data_path / "X_test.npy")
    y_train = np.load(data_path / "y_train.npy")
    y_test = np.load(data_path / "y_test.npy")

    with open(data_path / "feature_metadata.pkl", "rb") as f:
        feature_metadata = pickle.load(f)

    raw_data = pd.read_csv(data_path / "raw_data.csv")

    return X_train, X_test, y_train, y_test, feature_metadata, raw_data


def winequality_transform(X: pd.DataFrame, feature_metadata: Dict[str, Any]) -> np.ndarray:
    """Transform raw data using feature metadata encoders."""
    input_data = []
    for k, v in feature_metadata.items():
        if np.isin('missing', X[[k]].values):
            X[[k]] = X[[k]].replace("missing", v['encoder'].categories_[0][-1])
        preprocessed = v['encoder'].transform(X[[k]].values)
        if v['type'] == 'categorical':
            preprocessed = preprocessed.toarray()
        input_data.append(preprocessed)

    input_array = np.concatenate(input_data, axis=1)
    return input_array


def winequality_invert_transform(input_array: np.ndarray, feature_metadata: Dict[str, Any]) -> pd.DataFrame:
    """Invert transformed data back to original feature space."""
    inverted_data = {}

    for col, meta in feature_metadata.items():
        if meta['type'] == 'categorical':
            start_idx, end_idx = meta['index'][0], meta['index'][-1] + 1
            cat_data = input_array[:, start_idx:end_idx]
            inverted_col = meta['encoder'].inverse_transform(cat_data)
            inverted_data[col] = inverted_col.flatten()
        else:
            idx = meta['index']
            num_data = input_array[:, idx].reshape(-1, 1)
            inverted_col = meta['encoder'].inverse_transform(num_data)
            inverted_data[col] = inverted_col.flatten()

    return pd.DataFrame(inverted_data)


def winequality_find_idx(a: list, b: list) -> list:
    """Find permutation index where a[idx] = b."""
    if sorted(a) != sorted(b):
        return None

    pos_map = defaultdict(deque)
    for i, val in enumerate(a):
        pos_map[val].append(i)

    idx = []
    for val in b:
        idx.append(pos_map[val].popleft())

    return idx