#!/usr/bin/env python3
"""
Open Gaming Forensic Lab — gacha pull simulator (published with every episode that uses it).

    python3 lab/pull_sim.py lab/game_rules.json --players 1000000 --out releases/EP03-.../data

Model per game (all parameters live in game_rules.json with a source and official/community status):
  - base_rate      probability of a top-rarity drop on a normal pull
  - soft_pity_start, soft_pity_step   from pull `start`, the rate rises by `step` per pull (community-measured)
  - hard_pity      the pull that guarantees a top-rarity drop
  - featured_rate  chance a top-rarity drop is the featured unit (0.5 = "50/50")
  - guarantee      if you lose the 50/50, the next top-rarity drop is the featured unit

The simulator answers one question: how many pulls (and how much money at the best-value top-up)
does it take to obtain ONE featured top-rarity unit, starting from zero pity and no guarantee?
Results are a distribution, not a single number, because luck matters.
"""
import argparse
import json
from pathlib import Path

import numpy as np


def pity_curve(rules: dict) -> np.ndarray:
    """Probability of a top-rarity drop on pull k (index k-1) since the last one.

    Either a linear soft pity (soft_pity_start + soft_pity_step) or a stepped schedule
    soft_pity_schedule = [[first_pull, last_pull, step_per_pull], ...]; the hard-pity pull is 100%.
    """
    base, hard = rules["base_rate"]["value"], int(rules["hard_pity"]["value"])
    k = np.arange(1, hard + 1)
    p = np.full(hard, base, dtype=float)
    if rules.get("soft_pity_schedule"):
        rate = base
        for kk in range(1, hard + 1):
            for first, last, step in rules["soft_pity_schedule"]["value"]:
                if first <= kk <= last:
                    rate += step
            p[kk - 1] = min(1.0, rate)
    elif rules.get("soft_pity_start"):
        start, step = int(rules["soft_pity_start"]["value"]), rules["soft_pity_step"]["value"]
        p = np.where(k < start, base, np.minimum(1.0, base + step * (k - start + 1)))
    p[-1] = 1.0
    return p


def consolidated_rates(rules: dict) -> dict:
    """Exact (non-random) consolidated rates implied by the model, to check against official figures."""
    p = pity_curve(rules)
    survive = np.concatenate([[1.0], np.cumprod(1 - p)[:-1]])
    expected_pulls = float((survive).sum())          # E[pulls per top-rarity drop]
    f = rules["featured_rate"]["value"]
    per_featured = expected_pulls * (f + 2 * (1 - f))  # win: 1 drop, lose: 2 drops
    return {"top_rarity_rate": round(1 / expected_pulls, 5), "featured_rate": round(1 / per_featured, 5),
            "pulls_per_top_rarity": round(expected_pulls, 2), "pulls_per_featured": round(per_featured, 2)}


def pulls_to_featured(rules: dict, players: int, rng: np.random.Generator) -> np.ndarray:
    p = pity_curve(rules)
    hard = p.size
    featured = rules["featured_rate"]["value"]

    total = np.zeros(players, dtype=np.int64)
    done = np.zeros(players, dtype=bool)
    guaranteed = np.zeros(players, dtype=bool)
    while not done.all():
        idx = np.flatnonzero(~done)
        # sample pulls until the next top-rarity drop for every unfinished player
        u = rng.random((idx.size, hard))
        hit = u < p
        first = hit.argmax(axis=1) + 1          # hard pity column is always a hit
        total[idx] += first
        won = guaranteed[idx] | (rng.random(idx.size) < featured)
        done[idx[won]] = True
        guaranteed[idx[~won]] = True             # lost the 50/50 -> next one is guaranteed
    return total


def summarize(pulls: np.ndarray, rules: dict) -> dict:
    price = rules.get("usd_per_pull", {}).get("value")
    q = {f"p{int(x)}": int(np.percentile(pulls, x)) for x in (10, 25, 50, 75, 90, 99)}
    out = {"mean_pulls": round(float(pulls.mean()), 2), "max_pulls": int(pulls.max()), **q,
           "share_needing_hard_guarantee": round(float((pulls > int(rules["hard_pity"]["value"])).mean()), 4)}
    if price:
        out["usd_per_pull"] = price
        out["mean_usd"] = round(float(pulls.mean()) * price, 2)
        out["p50_usd"] = round(q["p50"] * price, 2)
        out["p90_usd"] = round(q["p90"] * price, 2)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("rules")
    ap.add_argument("--players", type=int, default=1_000_000)
    ap.add_argument("--seed", type=int, default=20261018)
    ap.add_argument("--games", nargs="*")
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    rules = json.loads(Path(a.rules).read_text())
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(a.seed)
    results = {"players": a.players, "seed": a.seed, "games": {}}
    for game, r in rules["games"].items():
        if a.games and game not in a.games:
            continue
        pulls = pulls_to_featured(r, a.players, rng)
        np.save(out / f"{game}_pulls.npy", pulls.astype(np.int16))
        results["games"][game] = {**summarize(pulls, r), "model_check": consolidated_rates(r),
                                  "official": {k: r[k]["value"] for k in ("official_top_rarity_rate", "official_featured_rate") if k in r}}
        print(game, results["games"][game])
    (out / "results.json").write_text(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()