#!/usr/bin/env python3
"""Generate the scenario list for the TP×DP matrix experiment.

Reads matrix.json and prints TSV lines:

    mark  input_len  output_len  concurrency  num_prompts

mark is one of Y/P/N.  The caller (run_bench.sh) decides how to treat each.
"""
import argparse
import json
import math
import os
from pathlib import Path


def sample_concurrency(low: int, high: int, target: int) -> list[int]:
    """Return only the low and high concurrency values in [low, high].

    For the TP×DP matrix we only need the two endpoints of the concurrency
    range (e.g. 1 and 128 for ISL=1024). The `target` argument is kept for
    API compatibility but is ignored.
    """
    assert 1 <= low <= high, f"invalid concurrency range: {low}-{high}"
    if low == high:
        return [low]
    return [low, high]


def generate_scenarios(matrix_path: Path, mode: str, target_samples: int) -> list[dict]:
    with open(matrix_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    matrix = data["matrix"]
    concurrency_cfg = data["concurrency"]

    scenarios = []
    for isl_str in sorted(matrix.keys(), key=int):
        dsl_map = matrix[isl_str]
        low = concurrency_cfg[isl_str]["low"]
        high = concurrency_cfg[isl_str]["high"]
        concurrencies = sample_concurrency(low, high, target_samples)

        for dsl_str in sorted(dsl_map.keys(), key=int):
            mark = dsl_map[dsl_str]
            if mode == "Y" and mark != "Y":
                continue
            if mode == "Y+P" and mark not in ("Y", "P"):
                continue
            # mode == "all" keeps everything, including N.

            for conc in concurrencies:
                scenarios.append(
                    {
                        "mark": mark,
                        "input_len": int(isl_str),
                        "output_len": int(dsl_str),
                        "concurrency": conc,
                        "num_prompts": conc * 5,
                    }
                )

    return scenarios


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--matrix", type=Path, default=Path("matrix.json"))
    parser.add_argument("--mode", choices=["Y", "Y+P", "all"], default=None,
                        help="Scenario selection mode. Defaults to matrix.mode.")
    parser.add_argument("--target-samples", type=int, default=0,
                        help="Target number of concurrency samples. 0 = heuristic (6-8).")
    args = parser.parse_args()

    with open(args.matrix, "r", encoding="utf-8") as f:
        data = json.load(f)

    mode = args.mode if args.mode else data.get("mode", "Y+P")

    target_samples = args.target_samples
    if target_samples <= 0:
        env_samples = os.getenv("CONCURRENCY_SAMPLES", "0")
        try:
            target_samples = int(env_samples)
        except ValueError:
            target_samples = 0
    if target_samples <= 0:
        target_samples = 7

    scenarios = generate_scenarios(args.matrix, mode, target_samples)

    print("mark\tinput_len\toutput_len\tconcurrency\tnum_prompts")
    for s in scenarios:
        print(f"{s['mark']}\t{s['input_len']}\t{s['output_len']}\t{s['concurrency']}\t{s['num_prompts']}")


if __name__ == "__main__":
    main()