#!/usr/bin/env python3
"""How the starting point (pity already built up, guarantee or not) changes the cost of one featured unit.

    python3 lab/start_states.py lab/game_rules.json genshin --out releases/<EP>/data
"""
import argparse
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from pull_sim import pity_curve  # noqa: E402

STATES = [("Fresh start", 0, False), ("Fresh start, guaranteed", 0, True),
          ("60 pity, not guaranteed", 60, False), ("60 pity, guaranteed", 60, True)]


def simulate(rules, start_pity, guaranteed, players, rng):
    p = pity_curve(rules)
    hard, f = p.size, rules["featured_rate"]["value"]
    total = np.zeros(players, dtype=np.int64)
    done = np.zeros(players, dtype=bool)
    guar = np.full(players, guaranteed)
    offset = start_pity  # only the first cycle starts part-way up the pity curve
    while not done.all():
        idx = np.flatnonzero(~done)
        q = p[offset:]
        u = rng.random((idx.size, q.size))
        first = (u < q).argmax(axis=1) + 1
        total[idx] += first
        won = guar[idx] | (rng.random(idx.size) < f)
        done[idx[won]] = True
        guar[idx[~won]] = True
        offset = 0
    return total


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("rules")
    ap.add_argument("game")
    ap.add_argument("--players", type=int, default=1_000_000)
    ap.add_argument("--seed", type=int, default=20261019)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    rules = json.loads(Path(a.rules).read_text())["games"][a.game]
    usd = rules["usd_per_pull"]["value"]
    rng = np.random.default_rng(a.seed)
    rows = []
    for label, pity, guar in STATES:
        t = simulate(rules, pity, guar, a.players, rng)
        rows.append({"state": label, "start_pity": pity, "guaranteed": guar,
                     "p50": int(np.percentile(t, 50)), "p90": int(np.percentile(t, 90)), "max": int(t.max()),
                     "p50_usd": round(np.percentile(t, 50) * usd, 2), "p90_usd": round(np.percentile(t, 90) * usd, 2)})
        print(rows[-1])
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    (out / f"start_states_{a.game}.json").write_text(json.dumps(rows, indent=2))


if __name__ == "__main__":
    main()