from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset

from data.schema import OBJECT_TO_ID, PREDICATE_TO_ID, SUBJECT_TO_ID, Fact


def fact_to_ids(value: dict[str, str]) -> tuple[int, int, int]:
    fact = Fact.from_dict(value)
    return (
        SUBJECT_TO_ID[fact.subject],
        PREDICATE_TO_ID[fact.predicate],
        OBJECT_TO_ID[fact.object],
    )


class EvidenceDataset(Dataset):
    def __init__(self, manifest: Path, image_size: int = 224):
        self.root = manifest.parent
        self.image_size = image_size
        self.records = [json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines()]

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> dict[str, Any]:
        record = self.records[index]
        image = Image.open(self.root / record["image"]).convert("RGB")
        image = image.resize((self.image_size, self.image_size))
        image_tensor = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1).float() / 255.0
        return {
            "image": image_tensor,
            "rules_text": record["rules_text"],
            "visual_facts": torch.tensor(
                [fact_to_ids(value) for value in record["visual_facts"]], dtype=torch.long
            ),
            "text_facts": torch.tensor(
                [fact_to_ids(value) for value in record["text_facts"]], dtype=torch.long
            ),
            "scene_id": record["scene_id"],
        }


def collate_evidence(batch: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "images": torch.stack([item["image"] for item in batch]),
        "rules_text": [item["rules_text"] for item in batch],
        "visual_facts": [item["visual_facts"] for item in batch],
        "text_facts": [item["text_facts"] for item in batch],
        "scene_id": [item["scene_id"] for item in batch],
    }