#!/usr/bin/env python3

from __future__ import annotations

import csv
import importlib.util
import json
import sys
import tempfile
import unittest
from pathlib import Path


ROOT = Path(__file__).parent


def load(name: str):
    path = ROOT / name
    spec = importlib.util.spec_from_file_location(path.stem, path)
    assert spec and spec.loader
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


class FidelityEnvelopeTest(unittest.TestCase):
    def test_materialize_allreduce(self) -> None:
        module = load("materialize_frontier_allreduce.py")
        rows = []
        for tp in (2, 4):
            for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192):
                rows.append(
                    {
                        "tensor_parallel_size": tp,
                        "num_tokens": tokens,
                        "hidden_dim": 2048,
                        "payload_bytes": tokens * 2048 * 2,
                        "critical_path_median_ms": tp + tokens / 1000,
                    }
                )
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            source = root / "allreduce.json"
            source.write_text(
                json.dumps(
                    {
                        "schema_version": "qwen30_vllm020_allreduce_frozen.v1",
                        "rows": rows,
                    }
                )
            )
            output = root / "all_reduce.csv"
            manifest = module.convert(source, output)
            self.assertEqual(manifest["rows"], 24)
            with output.open(newline="") as handle:
                converted = list(csv.DictReader(handle))
            self.assertEqual(converted[0]["num_workers"], "2")
            self.assertEqual(converted[0]["size"], "4096")
            self.assertEqual(converted[-1]["num_workers"], "4")
            self.assertEqual(converted[-1]["size"], str(8192 * 2048 * 2))
            self.assertEqual(
                converted[-1]["time_stats.all_reduce.median"],
                str(4 + 8192 / 1000),
            )


if __name__ == "__main__":
    unittest.main()