import random
import pandas as pd
from typing import Optional

class InningOutcomeDistribution:
    def __init__(self, team_abbr: str, data_source: str,
                 start_date: Optional[str] = None, end_date: Optional[str] = None):
        cols = ['events', 'home_team', 'away_team', 'inning_topbot', 'game_date']
        df = pd.read_parquet(data_source, columns=cols, engine='pyarrow')

        team_abbr = team_abbr.upper()
        df['game_date'] = pd.to_datetime(df['game_date'])

        if start_date:
            df = df[df['game_date'] >= pd.to_datetime(start_date)]
        if end_date:
            df = df[df['game_date'] <= pd.to_datetime(end_date)]

        team_df = df[
            (((df['home_team'] == team_abbr) & (df['inning_topbot'] == 'Bot')) |
             ((df['away_team'] == team_abbr) & (df['inning_topbot'] == 'Top'))) &
            df['events'].notna()
        ]

        if team_df.empty:
            available = sorted(set(df['home_team'].dropna()) | set(df['away_team'].dropna()))
            raise ValueError(f"No data for '{team_abbr}'. Available teams: {available}")

        total = len(team_df)
        counts = team_df['events'].value_counts()

        self.team = team_abbr
        self.probabilities = {
            '1B':   counts.get('single', 0) / total,
            '2B':   counts.get('double', 0) / total,
            '3B':   counts.get('triple', 0) / total,
            'HR':   counts.get('home_run', 0) / total,
            'WALK': (counts.get('walk', 0) + counts.get('hit_by_pitch', 0)) / total,
        }
        self.probabilities['OUT'] = 1.0 - sum(self.probabilities.values())

        self._outcomes = list(self.probabilities.keys())
        self._weights  = list(self.probabilities.values())

    def sample(self) -> str:
        return random.choices(self._outcomes, weights=self._weights, k=1)[0]

    def __repr__(self):
        return f"<OutcomeDist {self.team}: HR={self.probabilities['HR']:.4f}>"


# Lookup table: (bases_tuple, bases_advanced) -> (new_bases_tuple, runs_scored)
# bases_advanced: 1=single, 2=double, 3=triple
_HIT_LUT: dict = {
    ((0,0,0), 1): ((1,0,0), 0), ((0,0,0), 2): ((0,1,0), 0), ((0,0,0), 3): ((0,0,1), 0),
    ((1,0,0), 1): ((1,1,0), 0), ((1,0,0), 2): ((0,1,1), 0), ((1,0,0), 3): ((0,0,1), 1),
    ((0,1,0), 1): ((1,0,1), 0), ((0,1,0), 2): ((0,1,0), 1), ((0,1,0), 3): ((0,0,1), 1),
    ((0,0,1), 1): ((1,0,0), 1), ((0,0,1), 2): ((0,1,0), 1), ((0,0,1), 3): ((0,0,1), 1),
    ((1,1,0), 1): ((1,1,1), 0), ((1,1,0), 2): ((0,1,1), 1), ((1,1,0), 3): ((0,0,1), 2),
    ((1,0,1), 1): ((1,1,0), 1), ((1,0,1), 2): ((0,1,1), 1), ((1,0,1), 3): ((0,0,1), 2),
    ((0,1,1), 1): ((1,0,1), 1), ((0,1,1), 2): ((0,1,0), 2), ((0,1,1), 3): ((0,0,1), 2),
    ((1,1,1), 1): ((1,1,1), 1), ((1,1,1), 2): ((0,1,1), 2), ((1,1,1), 3): ((0,0,1), 3),
}


def _transition(outs: int, bases: tuple, outcome: str):
    if outcome == 'OUT':
        return outs + 1, bases, 0
    if outcome == 'HR':
        return outs, (0, 0, 0), sum(bases) + 1
    if outcome == 'WALK':
        b = list(bases)
        runs = 0
        if not b[0]:   b[0] = 1
        elif not b[1]: b[1] = 1
        elif not b[2]: b[2] = 1
        else:          runs = 1
        return outs, tuple(b), runs
    # 1B / 2B / 3B
    new_bases, runs = _HIT_LUT[(bases, {'1B': 1, '2B': 2, '3B': 3}[outcome])]
    return outs, new_bases, runs


def simulate_inning(dist: InningOutcomeDistribution) -> int:
    outs, bases, runs = 0, (0, 0, 0), 0
    while outs < 3:
        outs, bases, r = _transition(outs, bases, dist.sample())
        runs += r
    return runs


def simulate_game(home_dist: InningOutcomeDistribution,
                  away_dist: InningOutcomeDistribution,
                  n: int) -> dict:
    results: dict = {}
    for _ in range(n):
        home_score = away_score = 0
        inning = 0
        while True:
            inning += 1
            home_score += simulate_inning(home_dist)
            away_score += simulate_inning(away_dist)
            if inning >= 9 and home_score != away_score:
                break
            if inning >= 20:
                away_score += 1
                break
        key = (home_score, away_score)
        results[key] = results.get(key, 0) + 1
    return results