#!/usr/bin/env python

from itertools import combinations, product
from random import sample

import matplotlib.pyplot as plt
import numpy as np

SAMPLE_SIZE = 10000

RANKS = ["2", "3", "4", "5", "6", "7", "8", "9", "10", "J", "Q", "K", "A"]
LOOPED_RANKS = [RANKS[-1]] + RANKS
SUITS = ["♥", "♦", "♣", "♠"]
DECK = list(product(RANKS, SUITS))

def count_rank(table, rank):
    return len(list(filter(lambda card: card[0] == rank, table)))
def count_suit(table, suit):
    return len(list(filter(lambda card: card[1] == suit, table)))

def find_high_card(table):
    output = {}
    for rank in RANKS:
        output[rank] = count_rank(table, rank) >= 1
    return output
def find_one_pair(table):
    output = {}
    for rank in RANKS:
        output[rank] = count_rank(table, rank) >= 2
    return output
def find_two_pair(table):
    output = {}
    for (rank1, rank2) in combinations(RANKS, 2):
        output[rank1 + '+' + rank2] = count_rank(table, rank1) >= 2 and count_rank(table, rank2) >= 2
    return output
def find_three_of_a_kind(table):
    output = {}
    for rank in RANKS:
        output[rank] = count_rank(table, rank) >= 3
    return output
def find_straight(table):
    output = {}
    for i in range(4, len(LOOPED_RANKS)):
        output[LOOPED_RANKS[i]] = (count_rank(table, LOOPED_RANKS[i-4]) >= 1 and
                     count_rank(table, LOOPED_RANKS[i-3]) >= 1 and
                     count_rank(table, LOOPED_RANKS[i-2]) >= 1 and
                     count_rank(table, LOOPED_RANKS[i-1]) >= 1 and
                     count_rank(table, LOOPED_RANKS[i]) >= 1)
    return output
def find_flush(table):
    output = {}
    for suit in SUITS:
        output[suit] = count_suit(table, suit) >= 5
    return output
def find_full_house(table):
    output = {}
    for (rank1, rank2) in product(RANKS, RANKS):
        output[rank1 + ' over ' + rank2] = count_rank(table, rank1) >= 3 and count_rank(table, rank2) >= 2
    return output
def find_four_of_a_kind(table):
    output = {}
    for rank in RANKS:
        output[rank] = count_rank(table, rank) >= 4
    return output
def find_straight_flush(table):
    output = {}
    for i in range(4, len(LOOPED_RANKS)):
        found = False
        for s in SUITS:
            if ((LOOPED_RANKS[i-4], s) in table and
                (LOOPED_RANKS[i-3], s) in table and
                (LOOPED_RANKS[i-2], s) in table and
                (LOOPED_RANKS[i-1], s) in table and
                (LOOPED_RANKS[i], s) in table):
                found = True

        output[LOOPED_RANKS[i]] = found
    return output

def find_hands(table):
    return {
        "high card": find_high_card(table),
        "one pair": find_one_pair(table),
        "two pair": find_two_pair(table),
        "three of a kind": find_three_of_a_kind(table),
        "straight": find_straight(table),
        "flush": find_flush(table),
        "full house": find_full_house(table),
        "four of a kind": find_four_of_a_kind(table),
        "straight flush": find_straight_flush(table),
    }

def sanity_check(table, result):
    for hand_type in result["one pair"]:
        if result["one pair"][hand_type] and not result["high card"][hand_type]:
            print(f"Invalid result at one pair of {hand_type} but not high card")
            print(table)
            print(result)

    for hand_type in result["two pair"]:
        [rank1, rank2] = hand_type.split("+")
        if result["two pair"][hand_type] and not (result["one pair"][rank1] and result["one pair"][rank2]):
            print(f"Invalid result at two pair of {hand_type} but not respective one pairs")
            print(table)
            print(result)

    for hand_type in result["three of a kind"]:
        if result["three of a kind"][hand_type] and not result["one pair"][hand_type]:
            print(f"Invalid result at three of a kind of {hand_type} but not one pair")
            print(table)
            print(result)

    for hand_type in result["straight"]:
        i = LOOPED_RANKS.index(hand_type)
        if result["straight"][hand_type] and not (result["high card"][LOOPED_RANKS[i-4]] and 
                                                  result["high card"][LOOPED_RANKS[i-3]] and
                                                  result["high card"][LOOPED_RANKS[i-2]] and
                                                  result["high card"][LOOPED_RANKS[i-1]] and
                                                  result["high card"][LOOPED_RANKS[i]]):
            print(f"Invalid result at straight of {hand_type} but not high cards")
            print(table)
            print(result)

    for hand_type in result["full house"]:
        [rank1, rank2] = hand_type.split(" over ")
        if result["full house"][hand_type] and not (result["three of a kind"][rank1] and result["one pair"][rank2]):
            print(f"Invalid result at full house of {hand_type} but not respective three of a kind and one pair")
            print(table)
            print(result)

    for hand_type in result["four of a kind"]:
        if result["four of a kind"][hand_type] and not result["three of a kind"][hand_type]:
            print(f"Invalid result at four of a kind of {hand_type} but not three of a kind")
            print(table)
            print(result)

    for hand_type in result["straight flush"]:
        if result["straight flush"][hand_type] and not result["straight"][hand_type]:
            print(f"Invalid result at straight flush of {hand_type} but not straight")
            print(table)
            print(result)

# def find_hand_averages(result):
#     output = {}
#     for hand in result:
#         output[hand] = 0
#         for hand_type in result[hand]:
#             output[hand] += result[hand][hand_type]
#         output[hand] /= len(result[hand])
#     return output
def find_probabilities(size):
    results = []
    for _ in range(SAMPLE_SIZE):
        table = sample(DECK, size)
        result = find_hands(table)
        sanity_check(table, result)
        results += [result]

    # The idea was to find the probability of each hand this way:
    #   for each hand
    #       for each hand_type of this hand
    #           find the probability of this hand_type with occurrences / SAMPLE_SIZE
    #       find the probability of the hand by averaging the probability of each hand_type, that is
    #       sum everything and divide by the number of types (= len(results[0][hand]))
    # We don't use this approach becuase finding the intermediate probabilities for each hand_type introduces
    # division, which comes with rounding errors.
    # To avoid this we postpone division as much as possible, working with integers as much as possible
    probabilities = {}
    for hand in results[0]:
        occurrences = 0
        for hand_type in results[0][hand]:
            occurrences_per_type = 0
            for i in range(SAMPLE_SIZE):
                if results[i][hand][hand_type]:
                    occurrences_per_type += 1
            # probability_per_type should be occurrences_per_type / len(results[0][hand]) but
            # it's better to postpone divisions as much as possible to prevent rounding errors
            occurrences += occurrences_per_type
        probabilities[hand] = occurrences / (SAMPLE_SIZE * len(results[0][hand]))

    return probabilities

    # output = {}
    # for hand in results[0]:
    #     output[hand] = {}
    #     for hand_type in results[0][hand]:
    #         output[hand][hand_type] = 0
    #         for i in range(SAMPLE_SIZE):
    #             if results[i][hand][hand_type]:
    #                 output[hand][hand_type] += 1
    #         output[hand][hand_type] /= SAMPLE_SIZE
    # return find_hand_averages(output)

groups = tuple(range(5, 52, 5))
global_results = {
    "high card": (),
    "flush": (),
    "one pair": (),
    "straight": (),
    "two pair": (),
    "three of a kind": (),
    "full house": (),
    "straight flush": (),
    "four of a kind": (),
}
for size in groups:
    results = find_probabilities(size)
    for hand in global_results:
        global_results[hand] += (results[hand],)

    results = sorted(results.items(), key=lambda hand: -hand[1])
    print("===========================")
    print(" cards on table:", str(size).rjust(9))
    print("---------------------------")
    for (hand, p) in results:
        str_p = f"{p:.4%}"
        # str_p = f"{p}"
        print(" " + hand.ljust(15) + str_p.rjust(10))
    print("===========================")
    print("")
    print("")

x = np.arange(len(groups))  # the label locations
width = 0.08  # the width of the bars
multiplier = 0

fig, ax = plt.subplots(layout='constrained')

for attribute, measurement in global_results.items():
    offset = width * multiplier
    rects = ax.bar(x + offset, measurement, width, label=attribute)
    # ax.bar_label(rects, padding=3)
    multiplier += 1

# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_ylabel('Probabilità')
ax.set_title('Probabilità mani poker polacco in funzione del numero di carte in gioco')
ax.set_xticks(x + width*4, groups)
ax.legend(loc='upper left', ncols=3)
ax.set_ylim(0, 1.05)

plt.show()