"""
VQA Dataset Loader — for Flickr8k captions and VQAv2 question-answer pairs.

Handles image loading, resizing to 384×384, and tokenization.
"""

import os
import json
import torch
from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms


class CaptionDataset(Dataset):
    """
    Image captioning dataset (Flickr8k format).

    Each item: (image_tensor, caption_tokens, caption_padding_mask)
    Used for Stage 1 JEPA pretraining.
    """

    def __init__(
        self,
        image_dir: str,
        captions_file: str,
        tokenizer,
        img_size: int = 384,
        max_seq_len: int = 128,
    ):
        self.image_dir = image_dir
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len

        # Image 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]),
        ])

        # Load captions: expect format "image_file\tcaption" per line
        self.samples = []
        with open(captions_file) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                parts = line.split("\t", 1)
                if len(parts) == 2:
                    img_name = parts[0].split("#")[0]  # handle "img#0" format
                    caption = parts[1]
                    img_path = os.path.join(image_dir, img_name)
                    if os.path.exists(img_path):
                        self.samples.append((img_path, caption))

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

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

        # Load and transform image
        image = Image.open(img_path).convert("RGB")
        image = self.transform(image)

        # Tokenize caption
        token_ids = self.tokenizer.encode(caption)
        token_ids = token_ids[:self.max_seq_len]

        # Pad to max length
        padding = [self.tokenizer.pad_id] * (self.max_seq_len - len(token_ids))
        padding_mask = [True] * len(token_ids) + [False] * len(padding)
        token_ids = token_ids + padding

        return {
            "image": image,
            "caption_ids": torch.tensor(token_ids, dtype=torch.long),
            "caption_mask": torch.tensor(padding_mask, dtype=torch.bool),
        }


class VQADataset(Dataset):
    """
    Visual Question Answering dataset (VQAv2 format).

    Each item: (image_tensor, question_tokens, question_mask, answer_tokens, answer_mask)
    Used for Stage 2 supervised finetuning.
    """

    def __init__(
        self,
        image_dir: str,
        questions_file: str,
        annotations_file: str,
        tokenizer,
        img_size: int = 384,
        max_question_len: int = 64,
        max_answer_len: int = 32,
    ):
        self.image_dir = image_dir
        self.tokenizer = tokenizer
        self.max_question_len = max_question_len
        self.max_answer_len = max_answer_len

        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]),
        ])

        # Load questions and annotations
        with open(questions_file) as f:
            questions_data = json.load(f)
        with open(annotations_file) as f:
            annotations_data = json.load(f)

        # Build question_id → annotation map
        ann_map = {a["question_id"]: a for a in annotations_data["annotations"]}

        self.samples = []
        for q in questions_data["questions"]:
            qid = q["question_id"]
            if qid in ann_map:
                ann = ann_map[qid]
                # Use most frequent answer
                answer = ann.get("multiple_choice_answer", ann["answers"][0]["answer"])
                img_id = q["image_id"]
                img_name = f"COCO_val2014_{img_id:012d}.jpg"
                img_path = os.path.join(image_dir, img_name)
                self.samples.append({
                    "image_path": img_path,
                    "question": q["question"],
                    "answer": answer,
                })

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

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

        # Load image
        image = Image.open(sample["image_path"]).convert("RGB")
        image = self.transform(image)

        # Tokenize question
        q_ids = self.tokenizer.encode(sample["question"])[:self.max_question_len]
        q_pad = [self.tokenizer.pad_id] * (self.max_question_len - len(q_ids))
        q_mask = [True] * len(q_ids) + [False] * len(q_pad)
        q_ids = q_ids + q_pad

        # Tokenize answer
        a_ids = self.tokenizer.encode(sample["answer"])[:self.max_answer_len]
        a_pad = [self.tokenizer.pad_id] * (self.max_answer_len - len(a_ids))
        a_mask = [True] * len(a_ids) + [False] * len(a_pad)
        a_ids = a_ids + a_pad

        return {
            "image": image,
            "question_ids": torch.tensor(q_ids, dtype=torch.long),
            "question_mask": torch.tensor(q_mask, dtype=torch.bool),
            "answer_ids": torch.tensor(a_ids, dtype=torch.long),
            "answer_mask": torch.tensor(a_mask, dtype=torch.bool),
        }