#!/usr/bin/env python3
"""Dump deterministic batch membership and aggregate sizes for one implementation."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

from validation.artifacts import load_artifact
from validation.forward import _graph


class _LegacyDataset(torch.utils.data.Dataset):
    def __init__(self, artifact):
        self.artifact = artifact

    def __len__(self):
        return self.artifact.event_count

    def __getitem__(self, index):
        graph = _graph(self.artifact, index)
        tracking = torch.tensor(
            [
                self.artifact.folds[index],
                self.artifact.weights[index],
                index,
            ],
            dtype=torch.float32,
        )
        return (
            graph,
            torch.tensor(self.artifact.labels[index]),
            tracking,
            torch.empty(0),
        )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--implementation", choices=("legacy", "rewrite"), required=True
    )
    parser.add_argument("--artifact", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--batch-size", type=int, default=8)
    args = parser.parse_args()
    artifact = load_artifact(args.artifact)
    if args.implementation == "legacy":
        from dgl.dataloading import GraphDataLoader

        loader = GraphDataLoader(
            _LegacyDataset(artifact),
            batch_size=args.batch_size,
            shuffle=False,
            drop_last=False,
            num_workers=0,
        )
    else:
        from gnn4colliders.data import (
            EventMetadata,
            GraphDataLoader,
            GraphDataset,
            GraphSample,
        )

        samples = []
        for index in range(artifact.event_count):
            samples.append(
                GraphSample(
                    _graph(artifact, index),
                    torch.tensor(artifact.labels[index]),
                    None,
                    EventMetadata(
                        int(artifact.folds[index]),
                        float(artifact.weights[index]),
                        str(artifact.sample_id[index]),
                        {"index": index},
                    ),
                )
            )
        loader = GraphDataLoader(GraphDataset(samples), args.batch_size, shuffle=False)
    batches = []
    for batch in loader:
        if args.implementation == "legacy":
            graph, labels, tracking, _ = batch
            indices = tracking[:, 2].to(torch.long).tolist()
            ids = [str(artifact.sample_id[index]) for index in indices]
            weights = tracking[:, 1].tolist()
        else:
            graph = batch.graph
            labels = batch.labels
            ids = list(batch.metadata.sample_id)
            weights = batch.metadata.weight.tolist()
        batches.append(
            {
                "sample_id": ids,
                "num_graphs": len(ids),
                "nodes": int(graph.num_nodes()),
                "edges": int(graph.num_edges()),
                "labels": labels.reshape(-1).tolist(),
                "weights": weights,
            }
        )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps({"batch_size": args.batch_size, "batches": batches}, indent=2) + "\n"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())