import itertools

import numpy as np
import random
from collections import defaultdict
from collections.abc import Iterable
from pathlib import Path

from poker import Suit, Card
from shortdeck import ShortDeckHandEvaluator as HE
from shortdeck import ShortDeckRank


CHARS_SUIT = "shdc"
CHARS_RANK = "23456789TJQKA"

data_path = Path(".") / "ehs_data"
np_river = np.load(data_path / "river_ehs.npy")
np_turn = np.load(data_path / "turn_hist.npy")
np_flop = np.load(data_path / "flop_hist.npy")

cards = [Card(r, s) for r in ShortDeckRank for s in Suit]

Suit.__lt__ = lambda self, o: CHARS_SUIT.index(str(self)) < CHARS_SUIT.index(str(o))  # type: ignore
Card.__hash__ = lambda self: hash(str(self))  # type: ignore


class SuitMapping:
    def __init__(self):
        self.mapping = {}
        self.suits = list(reversed(Suit))

    def map_suit(self, s: Suit) -> Suit:
        if s not in self.mapping:
            self.mapping[s] = self.suits.pop()

        return self.mapping[s]


def to_iso(cards: list[Card], mapping: SuitMapping) -> list[Card]:
    def count_suit(card: Card) -> int:
        return sum(1 for other in cards if other.suit == card.suit)

    sorted_cards = sorted(cards, key=lambda c: (count_suit(c), c.rank, c.suit))

    res = []
    for card in sorted_cards:
        mapped_suit = mapping.map_suit(card.suit)
        res.append(Card(card.rank, mapped_suit))

    return sorted(res, key=lambda c: (c.rank, c.suit))


def card_to_u64(card: Card) -> int:
    r = list(CHARS_RANK).index(str(card.rank))
    s = list(CHARS_SUIT).index(str(card.suit))
    return (1 << r) << (16 * s)


def card_to_u8(card: Card) -> int:
    r = list(CHARS_RANK).index(str(card.rank))
    s = list(CHARS_SUIT).index(str(card.suit))
    return r | (s << 4)


def cards_to_u64(cards: list[Card]) -> int:
    from functools import reduce
    import operator

    return reduce(operator.or_, map(card_to_u64, cards))


def cards_to_u16(cards: list[Card]) -> int:
    return (card_to_u8(cards[1]) << 8) | card_to_u8(cards[0])


def calc_river_ehs(board: list[Card], player: list[Card]) -> float:
    player_hand = [*board, *player]
    player_ranking = HE.evaluate_hand(player_hand)

    acc = 0
    sum = 0
    for other in itertools.combinations(cards, 2):
        if set(other) & set(player_hand):
            continue
        if set(other) & set(board):
            continue
        other_ranking = HE.evaluate_hand([*board, *other])

        if player_ranking == other_ranking:
            acc += 1
        elif player_ranking > other_ranking:
            acc += 2

        sum += 2

    return acc / sum


CACHE = {}


def calc_river_ehs_cached(board, player) -> float:
    suit_map = SuitMapping()
    iso_board = to_iso(board, suit_map)
    iso_player = to_iso(player, suit_map)
    hand = "".join(map(str, [*iso_board, *iso_player]))
    if hand not in CACHE:
        CACHE[hand] = calc_river_ehs(board, player)

    return CACHE[hand]


def get_data(board: list[Card], player: list[Card]):
    def _get_data(data, board: list[Card], player: list[Card]):
        suit_map = SuitMapping()

        iso_board = to_iso(board, suit_map)
        iso_player = to_iso(player, suit_map)

        mask_board = data["board"] == cards_to_u64(iso_board)
        mask_player = data["player"] == cards_to_u16(iso_player)
        return data[mask_board & mask_player][0][2]

    match len(board):
        case 3:
            return _get_data(np_flop, board, player)
        case 4:
            return _get_data(np_turn, board, player)
        case 5:
            return _get_data(np_river, board, player)
        case _:
            raise NotImplementedError


def euclidean_dist(left, right):
    if isinstance(left, Iterable):
        v1 = np.sort(np.array(left, dtype=np.float32))
        v2 = np.sort(np.array(right, dtype=np.float32))
        return np.linalg.norm(v2 - v1)
    else:
        return np.abs(left - right) ** 2


def compare_data(sampled, board, player):
    d = euclidean_dist(get_data(board, player), sampled)

    if not np.isclose(d, 0.0):
        print(f"[{''.join(map(str, board))} {''.join(map(str, player))}]: {d}")


def analysis(flop, player, sampled):
    print(f"sampled flop: {''.join(map(str, flop))}")
    print(f"sampled player cards: {''.join(map(str, player))}")

    compare_data(
        [sampled[t][r] for t in sampled for r in sampled[t] if t > r],
        flop,
        player,
    )

    for turn in sampled:
        compare_data(list(sampled[turn].values()), [*flop, turn], player)

        for river in sampled[turn]:
            if turn > river:
                continue
            compare_data(sampled[turn][river], [*flop, turn, river], player)


def main():
    sample = random.sample(cards, 5)
    flop = sample[2:]
    player = sample[:2]
    sampled_ehs = defaultdict(dict)

    for turn in cards:
        if turn in flop or turn in player:
            continue

        for river in cards:
            if river in flop or river in player or river == turn:
                continue

            board = [*flop, turn, river]
            sampled_ehs[turn][river] = calc_river_ehs_cached(board, player)
            print(".", end="", flush=True)

    print("")
    analysis(flop, player, sampled_ehs)


main()