import json
from datetime import datetime, timezone
from pathlib import Path

import torch

from ..base import BatchMatrixData, MatrixDomain


def save_domain(output_dir: Path, domain: MatrixDomain, samples: list[BatchMatrixData]):
    output_dir.mkdir(parents=True, exist_ok=True)

    data = {
        "indices": [s.indices.cpu() for s in samples],
        "values": [s.values.cpu() for s in samples],
        "diagonals": [s.diagonals.cpu() for s in samples],
        "context_pairs": [s.context_pairs.cpu() for s in samples],
        "n": [s.n for s in samples],
        "domain": domain.value,
    }

    torch.save(data, output_dir / f"{domain.value}.pt")


def load_domain(dataset_dir: Path, domain: MatrixDomain) -> dict:
    path = dataset_dir / f"{domain.value}.pt"
    if not path.exists():
        raise FileNotFoundError(f"Dataset file not found: {path}")
    return torch.load(path, weights_only=False)


def write_meta(
    output_dir: Path,
    domains: list[MatrixDomain],
    num_matrices_per_domain: int,
    num_context_pairs: int,
    seed: int,
):
    meta = {
        "created_at": datetime.now(timezone.utc).isoformat(),
        "domains": [d.value for d in domains],
        "num_matrices_per_domain": num_matrices_per_domain,
        "num_context_pairs": num_context_pairs,
        "seed": seed,
    }
    with open(output_dir / "meta.json", "w") as f:
        json.dump(meta, f, indent=2)


def read_meta(dataset_dir: Path) -> dict:
    path = dataset_dir / "meta.json"
    if not path.exists():
        raise FileNotFoundError(f"Meta file not found: {path}")
    with open(path) as f:
        return json.load(f)
