"""Measure route multiplicity and temperature effects from exact graph laws."""

import argparse
from pathlib import Path
import pandas as pd
from dooable.graph import toy_graph
from dooable.exact import solve, tilted_reference, endpoint_distribution, expected_cost


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--output", default="results/exact_sweep")
    args = p.parse_args()
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    rows = []
    for multiplicity in [1, 2, 4, 8, 16]:
        graph = toy_graph(multiplicity)
        rewards = {"A": 0.0, "B": 0.0}
        for temperature in [0.1, 0.3, 0.7, 1.0, 2.0]:
            for name, policy in [
                ("exact", solve(graph, rewards, temperature).forward),
                ("reference_tilt", tilted_reference(graph, rewards)),
            ]:
                rows.append(
                    {
                        "multiplicity": multiplicity,
                        "temperature": temperature,
                        "method": name,
                        "probability_A": endpoint_distribution(graph, policy)["A"],
                        "mean_cost": expected_cost(graph, policy),
                    }
                )
    pd.DataFrame(rows).to_csv(out / "measurements.csv", index=False)
    print(f"Saved {len(rows)} exact comparisons to {out}")


if __name__ == "__main__":
    main()