#!/usr/bin/env python3
"""
Dota 2 Map and Series Win Probability Calculator (Strengthened)

Core models:
- Direct map_p or rating-difference logistic (Elo-style) conversion to single-map win probability.
- Exact dynamic programming for Bo1 / Bo2 / Bo3 / Bo5 series score distributions
  (maps treated independent — practical approximation).
- Optional Poisson kill-lead path: model expected kills as Poisson, derive lead
  distribution, then map lead → win probability via calibrated logistic.
  Use as sensitivity / cross-check only.

Examples:
  # Direct map probability
  python3 calculate_dota2_probs.py --map_p 0.58 --format bo3

  # From rating difference (positive = team A stronger)
  python3 calculate_dota2_probs.py --rating_diff 150 --format bo3

  # Series with per-map probabilities
  python3 calculate_dota2_probs.py --map_ps "0.62,0.55,0.48" --format bo3

  # Poisson kill-lead strengthened path
  python3 calculate_dota2_probs.py --kills_a 29.0 --kills_b 24.5 --format map --poisson_lead

  # Bo5 with more sims
  python3 calculate_dota2_probs.py --map_p 0.57 --format bo5 --n_sims 100000
"""

from __future__ import annotations

import argparse
import math
import sys
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

import numpy as np
from scipy.stats import poisson


def logistic_from_rating_diff(rating_diff: float, scale: float = 400.0) -> float:
    """Classic Elo-style win probability. rating_diff > 0 means team A stronger."""
    return 1.0 / (1.0 + 10.0 ** (-rating_diff / scale))


def logistic_from_lead(lead: float, midpoint: float = 0.0, steepness: float = 0.12) -> float:
    """
    Empirical mapping from kill (or gold/10k) lead to win probability.
    Calibrated roughly on pro data: ~+8–10 kill lead or +8–10k net worth mid-game
    corresponds to ~70–75% win probability. Adjust steepness if needed.
    """
    return 1.0 / (1.0 + math.exp(-steepness * (lead - midpoint)))


def poisson_lead_win_prob(
    kills_a: float,
    kills_b: float,
    max_k: int = 60,
    steepness: float = 0.11,
) -> float:
    """
    Model team kills as independent Poisson, compute P(A wins) by summing
    over possible kill outcomes weighted by the lead→win logistic.
    Returns expected win probability for team A.
    """
    if kills_a <= 0 or kills_b <= 0:
        raise ValueError("kills_a and kills_b must be positive")
    p_win = 0.0
    # Truncate Poisson tails
    for ka in range(0, max_k + 1):
        pa = poisson.pmf(ka, kills_a)
        if pa < 1e-9:
            continue
        for kb in range(0, max_k + 1):
            pb = poisson.pmf(kb, kills_b)
            if pb < 1e-9:
                continue
            lead = ka - kb
            p_win += pa * pb * logistic_from_lead(lead, steepness=steepness)
    return float(p_win)


def series_dp(map_ps: List[float], best_of: int) -> Tuple[float, Dict[str, float]]:
    """
    Exact dynamic programming for series win probability and score distribution.
    map_ps: list of P(A wins map i) for i = 0 .. best_of-1 (or shorter; cycles last if needed).
    Returns (P(A wins series), score_dist) where score_diff keys are like "2-0", "2-1".
    """
    if best_of not in (1, 2, 3, 5):
        raise ValueError("best_of must be 1, 2, 3 or 5")
    wins_needed = (best_of // 2) + 1
    # state: (maps_played, wins_a, wins_b) → probability
    # We only need up to wins_needed
    from functools import lru_cache

    @lru_cache(maxsize=None)
    def prob(wa: int, wb: int, idx: int) -> float:
        if wa >= wins_needed:
            return 1.0
        if wb >= wins_needed:
            return 0.0
        if idx >= len(map_ps) and len(map_ps) > 0:
            # reuse last map_p if more maps than provided
            p = map_ps[-1]
        else:
            p = map_ps[idx] if idx < len(map_ps) else map_ps[-1]
        return p * prob(wa + 1, wb, idx + 1) + (1.0 - p) * prob(wa, wb + 1, idx + 1)

    series_p = prob(0, 0, 0)

    # Score distribution via enumeration of paths (small state space)
    score_dist: Dict[str, float] = defaultdict(float)
    # BFS / recursive count of terminating scores
    def collect(wa: int, wb: int, idx: int, path_p: float):
        if wa >= wins_needed or wb >= wins_needed:
            key = f"{wa}-{wb}"
            score_dist[key] += path_p
            return
        p = map_ps[idx] if idx < len(map_ps) else map_ps[-1]
        collect(wa + 1, wb, idx + 1, path_p * p)
        collect(wa, wb + 1, idx + 1, path_p * (1.0 - p))

    collect(0, 0, 0, 1.0)
    # Normalize just in case of floating error
    total = sum(score_dist.values())
    if total > 0:
        for k in score_dist:
            score_dist[k] /= total
    return series_p, dict(score_dist)


def print_results(
    map_p: float,
    format_str: str,
    series_p: Optional[float] = None,
    score_dist: Optional[Dict[str, float]] = None,
    extra: str = "",
):
    print("=" * 60)
    print("Dota 2 Win Probability Results")
    print("=" * 60)
    if extra:
        print(extra)
    print(f"Single-map P(A wins) : {map_p*100:.1f}%")
    print(f"Single-map P(B wins) : {(1-map_p)*100:.1f}%")
    if series_p is not None and score_dist is not None:
        print(f"\nSeries format       : {format_str.upper()}")
        print(f"Series P(A wins)    : {series_p*100:.1f}%")
        print(f"Series P(B wins)    : {(1-series_p)*100:.1f}%")
        print("\nScore distribution:")
        # Sort nicely: higher A scores first
        def sort_key(s: str):
            a, b = map(int, s.split("-"))
            return (-a, b)
        for score in sorted(score_dist.keys(), key=sort_key):
            print(f"  {score:>5} : {score_dist[score]*100:5.1f}%")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(
        description="Dota 2 map & series win probability calculator"
    )
    parser.add_argument("--map_p", type=float, help="Direct P(A wins a single map)")
    parser.add_argument(
        "--rating_diff",
        type=float,
        help="Rating difference (A - B). Converted via Elo logistic (scale=400)",
    )
    parser.add_argument(
        "--map_ps",
        type=str,
        help='Comma-separated per-map probs e.g. "0.62,0.55,0.48"',
    )
    parser.add_argument(
        "--kills_a", type=float, help="Expected kills for team A (Poisson path)"
    )
    parser.add_argument(
        "--kills_b", type=float, help="Expected kills for team B (Poisson path)"
    )
    parser.add_argument(
        "--poisson_lead",
        action="store_true",
        help="Use Poisson kill-lead model (requires --kills_a --kills_b)",
    )
    parser.add_argument(
        "--format",
        type=str,
        default="map",
        choices=["map", "bo1", "bo2", "bo3", "bo5"],
        help="Output format: map only or series BoN",
    )
    parser.add_argument(
        "--n_sims",
        type=int,
        default=50000,
        help="(Reserved for future Monte-Carlo extensions)",
    )
    parser.add_argument(
        "--radiant_adv",
        type=float,
        default=0.0,
        help="Additive Radiant side advantage to apply to map_p (e.g. 0.015)",
    )
    args = parser.parse_args()

    map_p: Optional[float] = None
    extra_notes = []

    if args.poisson_lead:
        if args.kills_a is None or args.kills_b is None:
            print("ERROR: --poisson_lead requires --kills_a and --kills_b", file=sys.stderr)
            sys.exit(1)
        map_p = poisson_lead_win_prob(args.kills_a, args.kills_b)
        extra_notes.append(
            f"Poisson kill-lead model: λA={args.kills_a:.1f}, λB={args.kills_b:.1f}"
        )
    elif args.map_ps:
        try:
            map_ps = [float(x.strip()) for x in args.map_ps.split(",")]
            if not map_ps:
                raise ValueError
            for p in map_ps:
                if not (0.0 < p < 1.0):
                    raise ValueError("each map_p must be in (0,1)")
        except Exception as e:
            print(f"ERROR: invalid --map_ps: {e}", file=sys.stderr)
            sys.exit(1)
        # For series we will use the list; for single map report average
        map_p = float(np.mean(map_ps))
        extra_notes.append(f"Per-map probabilities: {map_ps}")
    elif args.rating_diff is not None:
        map_p = logistic_from_rating_diff(args.rating_diff)
        extra_notes.append(f"From rating_diff={args.rating_diff:.1f} (Elo scale=400)")
    elif args.map_p is not None:
        map_p = args.map_p
        if not (0.0 < map_p < 1.0):
            print("ERROR: --map_p must be in (0,1)", file=sys.stderr)
            sys.exit(1)
    else:
        print(
            "ERROR: provide one of --map_p, --rating_diff, --map_ps or --poisson_lead",
            file=sys.stderr,
        )
        sys.exit(1)

    # Optional Radiant advantage
    if args.radiant_adv != 0.0:
        map_p = min(max(map_p + args.radiant_adv, 0.01), 0.99)
        extra_notes.append(f"Applied Radiant advantage {args.radiant_adv:+.3f}")

    fmt = args.format.lower()
    if fmt == "map" or fmt == "bo1":
        print_results(map_p, "map", extra=" | ".join(extra_notes))
        return

    # Series
    best_of = {"bo2": 2, "bo3": 3, "bo5": 5}[fmt]
    if args.map_ps:
        map_ps_list = [float(x.strip()) for x in args.map_ps.split(",")]
        # Pad or trim to reasonable length
        while len(map_ps_list) < best_of:
            map_ps_list.append(map_ps_list[-1])
        map_ps_list = map_ps_list[:best_of]
    else:
        map_ps_list = [map_p] * best_of

    series_p, score_dist = series_dp(map_ps_list, best_of)
    print_results(
        map_p if not args.map_ps else float(np.mean(map_ps_list)),
        fmt,
        series_p=series_p,
        score_dist=score_dist,
        extra=" | ".join(extra_notes),
    )


if __name__ == "__main__":
    main()