#!/usr/bin/env python3

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np


def calculate_probabilities(card1, card2, num_decks=1):
	deck = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] * int(4 * num_decks)

	deck.remove(card1)
	deck.remove(card2)

	is_pair = card1 == card2

	if is_pair:
		higher_cards = [card for card in deck if card > card1]
		lower_cards = [card for card in deck if card < card1]
		collision_cards = [card for card in deck if card == card1]

		higher_prob = len(higher_cards) / len(deck)
		lower_prob = len(lower_cards) / len(deck)
		collision_prob = len(collision_cards) / len(deck)
		print('higher_prob: ', higher_prob)
		print('lower_prob: ', lower_prob)
		print('collision_prob: ', collision_prob)

		higher_expected_value = higher_prob * 1 - collision_prob * 3 - lower_prob * 1
		lower_expected_value = lower_prob * 1 - collision_prob * 3 - higher_prob * 1

		return {'higher': (higher_prob, higher_expected_value), 'lower': (lower_prob, lower_expected_value)}

	else:
		lower = min(card1, card2)
		upper = max(card1, card2)
		winning_cards = [card for card in deck if lower < card < upper]
		collision_cards = [card for card in deck if card == lower or card == upper]
		
		win_probability = len(winning_cards) / len(deck)
		collision_probability = len(collision_cards) / len(deck)
		print('win_probability: ', win_probability)
		print('collision_probability: ', collision_probability)

		expected_value = win_probability * 1 - collision_probability * 2 - (1 - win_probability - collision_probability) * 1
		return {'between': (win_probability, expected_value)}

def choose_best_strategy_probabilities(card1, card2, num_decks=1):
	probabilities = calculate_probabilities(card1, card2, num_decks)
	print('probabilities: ', probabilities)
	if 'between' in probabilities:
		return probabilities['between']
	elif probabilities['higher'][1] > probabilities['lower'][1]:
		return probabilities['higher']
	else:
		return probabilities['lower']

def plot_probabilities_with_hotmap(num_decks=1):
	expectations = np.zeros((13, 13))
	for i in range(13):
		for j in range(13):
			probabilities = choose_best_strategy_probabilities(i + 1, j + 1, num_decks)
			expectations[i, j] = probabilities[1]

	plt.figure(figsize=(10, 10))

	# flip the expectations matrix by y-axis to align with the card values
	expectations = np.flipud(expectations)
	ax = sns.heatmap(expectations, annot=True, fmt='.2f', cmap='coolwarm', center=0,
		xticklabels=['A', '2', '3', '4', '5', '6', '7', '8', '9', '10', 'J', 'Q', 'K'],
		yticklabels=['K', 'Q', 'J', '10', '9', '8', '7', '6', '5', '4', '3', '2', 'A'],
	)
	ax.set(xlabel='Card 1', ylabel='Card 2', title=f'Expectation Values Distribution for Dragon Gate Game with {num_decks} Deck{"s" if num_decks > 1 else ""}')
	plt.show()