"""Evaluate new public-property weights on one stored reaction graph."""

import argparse, json
from pathlib import Path
import pandas as pd
from dooable.graph import Graph
from dooable.properties import property_rewards
from dooable.exact import solve, endpoint_distribution, expected_cost
from dooable.learning import train


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--graph", required=True)
    parser.add_argument("--models", required=True)
    parser.add_argument("--output", default="results/preferences")
    parser.add_argument("--neural-steps", type=int, default=0)
    args = parser.parse_args()
    g = Graph.load(args.graph)
    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    rows = []
    previous = None
    for beta in [1.0, 2.0, 5.0]:
        for w in [0.0, 0.25, 0.5, 0.75, 1.0]:
            rewards, scores = property_rewards(
                g, args.models, weights=(w, 1 - w), concentration=beta
            )
            directory = out / f"beta{beta:g}_weight{w:g}"
            directory.mkdir(exist_ok=True)
            (directory / "rewards.json").write_text(json.dumps(rewards, indent=2))
            policies = {"exact": solve(g, rewards, 0.7).forward}
            if args.neural_steps:
                model, _ = train(
                    g,
                    rewards,
                    0.7,
                    steps=args.neural_steps,
                    output=directory,
                    initialize=previous,
                )
                policies["dooable"] = model.probabilities()
                previous = directory
            for name, p in policies.items():
                masses = endpoint_distribution(g, p)
                s = scores.set_index("smiles")
                a = pd.Series(masses).reindex(s.index)
                rows.append(
                    {
                        "method": name,
                        "bace_weight": w,
                        "concentration": beta,
                        "mean_bace_utility": float(a @ s.bace_utility),
                        "mean_caco2_utility": float(a @ s.caco2_utility),
                        "mean_cost": expected_cost(g, p),
                    }
                )
    pd.DataFrame(rows).to_csv(out / "measurements.csv", index=False)
    print(f"Saved {len(rows)} preference measurements to {out}")


if __name__ == "__main__":
    main()