from pathlib import Path

import torch
from torch.utils.data import Dataset

from ..base import BatchMatrixData, MatrixDomain
from .storage import load_domain, read_meta


class OfflineMatrixDataset(Dataset):

    def __init__(self, dataset_dir: Path, domains: list[MatrixDomain] | None = None):
        meta = read_meta(dataset_dir)

        if domains is None:
            domains = [MatrixDomain(d) for d in meta["domains"]]

        self.entries = []
        for domain in domains:
            data = load_domain(dataset_dir, domain)
            num_samples = len(data["n"])
            for i in range(num_samples):
                self.entries.append({
                    "indices": data["indices"][i],
                    "values": data["values"][i],
                    "diagonals": data["diagonals"][i],
                    "context_pairs": data["context_pairs"][i],
                    "n": data["n"][i],
                    "domain": MatrixDomain(data["domain"]),
                })

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

    def __getitem__(self, idx: int) -> BatchMatrixData:
        e = self.entries[idx]
        return BatchMatrixData(
            indices=e["indices"],
            values=e["values"],
            diagonals=e["diagonals"],
            context_pairs=e["context_pairs"],
            n=e["n"],
            batch_size=1,
            domain=e["domain"],
        )
