#!/usr/bin/env python3
"""Create a node/chunk plan for full test-set CPU generation."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import pandas as pd


ROOT = Path("/workspace/mp_20_pxrdnet")
MODEL_PATH = ROOT / "hydra/singlerun/2026-07-20/pxrdgen_raw512_run02"
TEST_PATH = ROOT / "data/mp_20_pxrdgen_xrd90_0p1_raw512/test.csv"
DEFAULT_OUTPUT_ROOT = ROOT / "paper_results_pxrdgen_match_only/cpu_multinode/full_test"


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--chunk-size", type=int, required=True)
    parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
    parser.add_argument("--total-materials", type=int, default=0)
    args = parser.parse_args()

    if args.chunk_size <= 0:
        raise ValueError("--chunk-size must be positive.")
    total = args.total_materials or int(len(pd.read_pickle(TEST_PATH)))
    output_root = Path(args.output_root)
    output_root.mkdir(parents=True, exist_ok=True)
    node_count = math.ceil(total / args.chunk_size)

    rows = []
    for node_index in range(node_count):
        first_idx = node_index * args.chunk_size
        num_materials = min(args.chunk_size, total - first_idx)
        rows.append(
            {
                "node_index": node_index,
                "first_idx": first_idx,
                "num_materials": num_materials,
                "output_dir": str(output_root / f"node_{node_index:05d}"),
            }
        )

    payload = {
        "model_path": str(MODEL_PATH),
        "test_path": str(TEST_PATH),
        "total_materials": total,
        "chunk_size": args.chunk_size,
        "node_count": node_count,
        "output_root": str(output_root),
        "rows": rows,
    }
    (output_root / "plan.json").write_text(json.dumps(payload, indent=2) + "\n")
    with (output_root / "plan.tsv").open("w") as f:
        f.write("node_index\tfirst_idx\tnum_materials\toutput_dir\n")
        for row in rows:
            f.write(
                f"{row['node_index']}\t{row['first_idx']}\t{row['num_materials']}\t{row['output_dir']}\n"
            )
    print(json.dumps({k: payload[k] for k in ["total_materials", "chunk_size", "node_count", "output_root"]}, indent=2))
    return 0


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