import random
from typing import Literal


BAT = "🦇"
SNAKE = "🐍"

Choice = Literal["🗿", "🗞", "🔪"]
ROCK: Choice = "🗿"
PAPER: Choice = "🗞"
SCISSORS: Choice = "🔪"

CHOICES: tuple[Choice, ...] = ROCK, PAPER, SCISSORS

def beats(x: Choice, y: Choice):
    return (x, y) in {(PAPER, ROCK), (SCISSORS, PAPER), (ROCK, SCISSORS)}


def main():
    counter = 0

    player_score = 0
    adversary_score = 0

    adversary = Adversary()

    while True:
        print(f"== Round {counter + 1} ==")
        print(f"SCORE: {player_score}:{adversary_score}")
        print()
        print(f"{SNAKE}: Hsss!")
        player_choice: Choice = input_choice(f"{BAT}", {"R": ROCK, "P": PAPER, "S": SCISSORS})
        adversary_choice: Choice = adversary.pick()

        if beats(player_choice, adversary_choice):
            player_score += 1
            msg = "VICTORY!"
        elif beats(adversary_choice, player_choice):
            adversary_score += 1
            msg = "DEFEAT!"
        else:
            msg = "DRAW!"

        print()
        print(f"{BAT}: {player_choice}! {VERBS[player_choice]}")
        print(f"{SNAKE}: {adversary_choice}! {VERBS[adversary_choice]}")
        print(msg)
        print()

        adversary.notify(player_choice, adversary_choice)
        counter += 1


VERBS = {
    ROCK: "THUD!",
    PAPER: "WHACK!",
    SCISSORS: "STAB!",
}


def input_choice[T](prompt: str, options: dict[str, T]) -> T:
    hint = f"({", ".join(options.keys())})"
    choice = input(f"{prompt} {hint}: ").upper()
    while choice not in options:
        choice = input(f"{hint}: ").upper()
    return options[choice]


class Adversary(object):
    def __init__(self):
        self._history: list[tuple[str, str]] = []

    def notify(self, player: str, adversary: str):
        self._history.append((player, adversary))

    def pick(self) -> Choice:
        return random.choice(CHOICES)


class Predictor(object):
    def __init__(self, context_size: int, include_player_moves: bool, include_adversary_moves: bool):
        self._context_size = context_size
        self._include_player_moves = include_player_moves
        self._include_adversary_moves = include_adversary_moves

    def predict(self, history: list[tuple[str, str]]) -> dict[str, float]:
        context_size = self._context_size

        context_size_in_items = context_size
        if self._include_player_moves and self._include_adversary_moves:
            context_size_in_items /= 2

        def _pluck_context(i: int) -> str:
            raw_context = history[i:i + context_size]
            context: str = ""
            for player, adversary in raw_context:
                if self._include_player_moves:
                    context += player
                if self._include_adversary_moves:
                    context += adversary
                
            return context[-self._context_size:]

        observations: dict[str, int] = {c: 1 for c in CHOICES}
        
        final_context = _pluck_context(len(history) - context_size)

        for i in range(len(history) - context_size):
            context = _pluck_context(i)
            observation, _ = history[i + context_size]
            if context == final_context:
                observations[observation] += 1

        sum_count = 0
        for r in CHOICES:
            sum_count += observations[r]

        return {
            r: observations[r] / sum_count
            for r in CHOICES
        }


class Adversary(object):
    def __init__(self):
        self._history: list[tuple[str, str]] = []

    def notify(self, player: str, adversary: str):
        self._history.append((player, adversary))

    def pick(self) -> Choice:
        player_probability = self.predict()
        expected_value: dict[Choice, float] = {c: 0.0 for c in CHOICES}
        for adversary in CHOICES:
            for player in CHOICES:
                expected_value[adversary] += (
                    1.0 if beats(adversary, player) else 
                    -1.0 if beats(player, adversary) else
                    0.0
                ) * player_probability[player]
        print(f"SNAKE PREDICTIONS: {player_probability}")
        print(f"SNAKE EV: {expected_value}")
        max_expected_value = max(expected_value.values())
        best_choices: list[Choice] = [c for c, v in expected_value.items() if v == max_expected_value]
        return random.choice(best_choices)

    def predict(self) -> dict[Choice, float]:
        def _merge(tables: list[dict[Choice, float]]):
            out = {}
            for c in CHOICES:
                sum_ = 0.0
                for t in tables:
                    sum_ += t[c]
                sum_ /= len(tables)
                out[c] = sum_
            return out


        tables = []
        for p in [
            Predictor(0, include_player_moves=False, include_adversary_moves=False),
            Predictor(1, include_player_moves=True, include_adversary_moves=False),
            Predictor(1, include_player_moves=False, include_adversary_moves=True),
            Predictor(2, include_player_moves=True, include_adversary_moves=True),
        ]:
            tables.append(p.predict(self._history))
        
        return _merge(tables)



if __name__ == "__main__":
    main()