data/multi_dataset.py

"""Multi-dataset pipeline for ArcisVLM training.

Supports loading from HuggingFace Hub, local files, or generating dummy data.
All datasets are converted to a unified format:
  {"image": tensor, "instruction": str, "answer": str, "task_token": str}
"""

import torch
from torch.utils.data import Dataset
import random
from typing import Optional
from data.formatters import format_for_task, classify_dataset_task


class UnifiedVLMDataset(Dataset):
    """Wraps a list of (image, question, answer) samples into unified format.

    Args:
        samples: List of dicts with at minimum "image" (path or tensor),
                 "question" (str), "answer" (str)
        dataset_name: Name of source dataset (for task type classification)
        tokenizer: BPE tokenizer for encoding
        img_size: Target image resolution
        max_instruction_len: Max tokens for instruction
        max_answer_len: Max tokens for answer
        transform: Image transform (if None, uses default ImageNet normalization)
    """

    def __init__(self, samples: list, dataset_name: str, tokenizer=None,
                 img_size: int = 448, max_instruction_len: int = 256,
                 max_answer_len: int = 512, transform=None):
        self.samples = samples
        self.dataset_name = dataset_name
        self.task_type = classify_dataset_task(dataset_name)
        self.tokenizer = tokenizer
        self.img_size = img_size
        self.max_instruction_len = max_instruction_len
        self.max_answer_len = max_answer_len

        if transform is None:
            try:
                from torchvision import transforms
                self.transform = transforms.Compose([
                    transforms.Resize((img_size, img_size)),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225]
                    ),
                ])
            except ImportError:
                self.transform = None  # Will handle tensors directly
        else:
            self.transform = transform

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

    def __getitem__(self, idx):
        sample = self.samples[idx]

        # Format the instruction/answer
        question = sample.get("question", sample.get("instruction", ""))
        answer = sample.get("answer", sample.get("response", ""))
        task_type = sample.get("task_type", self.task_type)

        formatted = format_for_task(question, answer, task_type)

        # Handle image
        image = sample.get("image", None)
        if image is None:
            raise ValueError(f"Image is None for sample {idx}. Dataset must provide real images.")
        elif isinstance(image, str):
            # Load from path
            from PIL import Image
            try:
                img = Image.open(image).convert("RGB")
                image = self.transform(img)
            except Exception as e:
                raise FileNotFoundError(f"Image not found or corrupted: {image}. Original error: {e}")
        elif hasattr(image, 'convert'):
            # PIL Image
            image = self.transform(image.convert("RGB"))

        result = {
            "image": image,
            "instruction": formatted["instruction"],
            "answer": formatted["answer"],
            "task_token": formatted["task_token"],
        }

        # Tokenize if tokenizer is available
        if self.tokenizer is not None:
            inst_ids = self.tokenizer.encode(formatted["task_token"] + " " + formatted["instruction"])
            ans_ids = self.tokenizer.encode(formatted["answer"])

            # Pad/truncate
            inst_ids = self._pad_or_truncate(inst_ids, self.max_instruction_len)
            ans_ids = self._pad_or_truncate(ans_ids, self.max_answer_len)

            result["instruction_ids"] = torch.tensor(inst_ids, dtype=torch.long)
            result["answer_ids"] = torch.tensor(ans_ids, dtype=torch.long)
            result["instruction_mask"] = (result["instruction_ids"] != self.tokenizer.pad_id).long()
            result["answer_mask"] = (result["answer_ids"] != self.tokenizer.pad_id).long()

        return result

    def _pad_or_truncate(self, ids, max_len):
        if len(ids) > max_len:
            return ids[:max_len]
        return ids + [self.tokenizer.pad_id] * (max_len - len(ids))


class DatasetMixer(Dataset):
    """Combines multiple datasets with configurable sampling weights.

    Args:
        datasets: Dict of name -> Dataset
        weights: Dict of name -> float (sampling probability, will be normalized)
        total_samples: Total virtual dataset size
        seed: Random seed for reproducibility
    """

    def __init__(self, datasets: dict, weights: dict = None,
                 total_samples: int = None, seed: int = 42):
        self.datasets = datasets
        self.dataset_names = list(datasets.keys())

        # Normalize weights
        if weights is None:
            weights = {name: len(ds) for name, ds in datasets.items()}
        total_weight = sum(weights.values())
        self.weights = {name: w / total_weight for name, w in weights.items()}

        # Compute total size
        if total_samples is None:
            total_samples = sum(len(ds) for ds in datasets.values())
        self.total_samples = total_samples

        # Pre-compute sampling indices
        self.rng = random.Random(seed)
        self._build_index()

    def _build_index(self):
        """Pre-compute which dataset and index each virtual index maps to."""
        self.index_map = []
        cumulative_weights = []
        cumsum = 0.0
        for name in self.dataset_names:
            cumsum += self.weights[name]
            cumulative_weights.append(cumsum)

        for _ in range(self.total_samples):
            r = self.rng.random()
            for i, cw in enumerate(cumulative_weights):
                if r <= cw:
                    name = self.dataset_names[i]
                    ds = self.datasets[name]
                    idx = self.rng.randint(0, len(ds) - 1)
                    self.index_map.append((name, idx))
                    break

    def __len__(self):
        return self.total_samples

    def __getitem__(self, idx):
        name, ds_idx = self.index_map[idx]
        return self.datasets[name][ds_idx]


def build_dummy_dataset(dataset_name: str, num_samples: int = 100,
                        img_size: int = 448) -> list:
    """Build dummy samples FOR TESTING ONLY.

    WARNING: This function exists solely for unit tests. Production training
    code must NEVER call this. Use real data instead.

    Returns list of dicts with "image", "question", "answer" keys.
    """
    samples = []
    task_type = classify_dataset_task(dataset_name)

    dummy_qa = {
        "vqa": [("What color is the car?", "red"), ("How many people?", "3"), ("What is this?", "a dog")],
        "caption": [("", "A busy street with pedestrians and cars"), ("", "A warehouse with stacked pallets")],
        "detection": [("", "person, car, traffic light"), ("", "forklift, pallet, worker")],
        "alert": [("", "Person detected in restricted zone near heavy machinery"), ("", "No anomalies detected")],
        "mcq": [("What is shown? A) cat B) dog", "A"), ("Which vehicle? A) car B) truck", "B")],
        "ocr": [("What text is visible?", "EXIT"), ("Read the sign", "PARKING LOT B")],
        "conversation": [("What do you see?", "I see a park with trees and a playground.")],
        "counting": [("How many vehicles?", "5"), ("Count the people", "12")],
        "tracking": [("Describe the movement", "Person walked from left to right across the frame")],
    }

    qa_pairs = dummy_qa.get(task_type, dummy_qa["vqa"])

    for i in range(num_samples):
        q, a = qa_pairs[i % len(qa_pairs)]
        samples.append({
            "image": torch.randn(3, img_size, img_size),
            "question": q,
            "answer": a,
            "task_type": task_type,
        })

    return samples


def build_stage2_dataset(config: dict, tokenizer=None, use_dummy: bool = False) -> DatasetMixer:
    """Build the full Stage 2 instruction tuning dataset mix.

    Args:
        config: Full config dict (with data.train_datasets.stage2)
        tokenizer: BPE tokenizer
        use_dummy: If True, use dummy data instead of downloading

    Returns:
        DatasetMixer combining all Stage 2 datasets
    """
    stage2_config = config.get("data", {}).get("train_datasets", {}).get("stage2", [])
    img_size = config.get("vision", {}).get("img_size", 448)

    datasets = {}
    weights = {}

    for ds_config in stage2_config:
        name = ds_config["name"]
        num_samples = ds_config.get("samples", 1000)

        if use_dummy:
            raise RuntimeError(
                f"FATAL: use_dummy=True is no longer supported. Download real data for '{name}'.\n"
                "Run: python3 scripts/download_all_data.py --stage 2"
            )
        else:
            # Try HuggingFace download — crash if unavailable
            try:
                from datasets import load_dataset
                source = ds_config.get("source", name)
                hf_ds = load_dataset(source, split="train", streaming=True)
                samples = []
                for i, item in enumerate(hf_ds):
                    if i >= num_samples:
                        break
                    samples.append(_convert_hf_sample(item, name))
            except Exception as e:
                raise RuntimeError(
                    f"FATAL: Failed to load dataset '{name}' from HuggingFace.\n"
                    f"Error: {e}\n"
                    "Install datasets: pip install datasets\n"
                    "Or download data: python3 scripts/download_all_data.py --stage 2"
                )

        ds = UnifiedVLMDataset(
            samples=samples,
            dataset_name=name,
            tokenizer=tokenizer,
            img_size=img_size,
        )
        datasets[name] = ds
        weights[name] = num_samples  # Weight proportional to target size

    return DatasetMixer(datasets=datasets, weights=weights)


def build_stage3_dataset(config: dict, tokenizer=None, use_dummy: bool = False) -> DatasetMixer:
    """Build Stage 3 domain fine-tuning dataset."""
    stage3_config = config.get("data", {}).get("train_datasets", {}).get("stage3", [])
    img_size = config.get("vision", {}).get("img_size", 448)

    datasets = {}
    weights = {}

    for ds_config in stage3_config:
        name = ds_config["name"]
        num_samples = ds_config.get("samples", 1000)

        if use_dummy:
            raise RuntimeError(
                f"FATAL: use_dummy=True is no longer supported. Download real data for '{name}'.\n"
                "Run: python3 scripts/download_all_data.py --stage 3"
            )

        # Try HuggingFace download — crash if unavailable
        try:
            from datasets import load_dataset
            source = ds_config.get("source", name)
            hf_ds = load_dataset(source, split="train", streaming=True)
            samples = []
            for i, item in enumerate(hf_ds):
                if i >= num_samples:
                    break
                samples.append(_convert_hf_sample(item, name))
        except Exception as e:
            raise RuntimeError(
                f"FATAL: Failed to load dataset '{name}' for Stage 3.\n"
                f"Error: {e}\n"
                "Install datasets: pip install datasets\n"
                "Or download data: python3 scripts/download_all_data.py --stage 3"
            )

        ds = UnifiedVLMDataset(
            samples=samples,
            dataset_name=name,
            tokenizer=tokenizer,
            img_size=img_size,
        )
        datasets[name] = ds
        weights[name] = num_samples

    return DatasetMixer(datasets=datasets, weights=weights)


def _convert_hf_sample(item: dict, dataset_name: str) -> dict:
    """Convert a HuggingFace dataset sample to our unified format."""
    # Handle different HF dataset schemas
    image = item.get("image", item.get("img", None))
    question = item.get("question", item.get("text", item.get("caption", "")))
    answer = item.get("answer", item.get("multiple_choice_answer",
             item.get("answers", [{"answer": ""}])[0] if isinstance(item.get("answers"), list) else ""))

    if isinstance(answer, dict):
        answer = answer.get("answer", answer.get("text", str(answer)))
    if isinstance(answer, list):
        answer = answer[0] if answer else ""

    return {
        "image": image,
        "question": str(question),
        "answer": str(answer),
    }