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 (
    FACT_TYPE,
    OBJECT_TO_ID,
    PREDICATE_TO_ID,
    RULE_TYPE,
    SUBJECT_TO_ID,
    Fact,
    Rule,
)


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


def rule_to_ids(value: dict[str, str]) -> tuple[int, int, int, int, int, int]:
    rule = Rule.from_dict(value)
    return (
        RULE_TYPE,
        SUBJECT_TO_ID[rule.variable],
        PREDICATE_TO_ID[rule.antecedent_predicate],
        OBJECT_TO_ID[rule.antecedent_object],
        PREDICATE_TO_ID[rule.consequent_predicate],
        OBJECT_TO_ID[rule.consequent_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_targets": torch.tensor(
                [fact_to_ids(value) for value in record["visual_facts"]], dtype=torch.long
            ),
            "text_targets": torch.tensor(
                [fact_to_ids(value) for value in record["text_facts"]]
                + [rule_to_ids(value) for value in record["text_rules"]],
                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_targets": [item["visual_targets"] for item in batch],
        "text_targets": [item["text_targets"] for item in batch],
        "scene_id": [item["scene_id"] for item in batch],
    }