#!/usr/bin/env python3
"""
Losing streaks on the 50/50 when a player goes after N featured units in a row
(pity and guarantee carry over between banners, as the developers state for the pity count).

    python3 lab/streaks.py lab/game_rules.json --games zzz hsr wuwa --targets 10 --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


def run(rules, players, targets, rng):
    p = pity_curve(rules)
    hard, f = p.size, rules["featured_rate"]["value"]
    pulls = np.zeros(players, dtype=np.int64)
    losses = np.zeros(players, dtype=np.int64)
    streak = np.zeros(players, dtype=np.int64)
    longest = np.zeros(players, dtype=np.int64)
    got = np.zeros(players, dtype=np.int64)
    guar = np.zeros(players, dtype=bool)
    while (got < targets).any():
        idx = np.flatnonzero(got < targets)
        first = (rng.random((idx.size, hard)) < p).argmax(axis=1) + 1
        pulls[idx] += first
        flip = ~guar[idx]                                   # a real 50/50 only when not guaranteed
        won = guar[idx] | (rng.random(idx.size) < f)
        lost = flip & ~won
        losses[idx[lost]] += 1
        streak[idx[lost]] += 1
        streak[idx[flip & won]] = 0                          # only a won flip ends a losing streak
        longest[idx] = np.maximum(longest[idx], streak[idx])
        got[idx[won]] += 1
        guar[idx] = lost
    return pulls, losses, longest


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("rules")
    ap.add_argument("--games", nargs="+", required=True)
    ap.add_argument("--targets", type=int, default=10)
    ap.add_argument("--players", type=int, default=1_000_000)
    ap.add_argument("--seed", type=int, default=20261020)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    all_rules = json.loads(Path(a.rules).read_text())["games"]
    rng = np.random.default_rng(a.seed)
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    res = {"players": a.players, "targets": a.targets, "seed": a.seed, "games": {}}
    for g in a.games:
        r = all_rules[g]
        pulls, losses, longest = run(r, a.players, a.targets, rng)
        np.save(out / f"streak_longest_{g}.npy", longest.astype(np.int8))
        np.save(out / f"streak_losses_{g}.npy", losses.astype(np.int8))
        usd = r["usd_per_pull"]["value"]
        res["games"][g] = {
            "name": r["name"], "featured_rate": r["featured_rate"]["value"],
            "losses_mean": round(float(losses.mean()), 2),
            "losses_dist": {int(k): round(float((losses == k).mean()), 4) for k in range(0, int(losses.max()) + 1)},
            "longest_streak_at_least": {k: round(float((longest >= k).mean()), 4) for k in range(1, 7)},
            "pulls_p50": int(np.percentile(pulls, 50)), "pulls_p90": int(np.percentile(pulls, 90)),
            "pulls_per_unit_p50": round(np.percentile(pulls, 50) / a.targets, 1),
            "usd_p50": round(np.percentile(pulls, 50) * usd), "usd_p90": round(np.percentile(pulls, 90) * usd),
        }
        print(g, json.dumps(res["games"][g]))
    (out / "streaks.json").write_text(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()