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


# 设置matplotlib支持中文的字体，这里使用SimHei字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False    # 解决保存图像是负号'-'显示为方块的问题

# Assuming an 8x8 grid representing the graphical password pad
grid_size = 8
np.random.seed(0)

# Generate a frequency matrix where the probabilities are influenced by the elements in the image
# Let's say trees are more likely to be chosen, and houses are less likely
# For simplicity, we'll assign higher probabilities to the "tree" locations and lower to "house" locations
# Randomly generated for demonstration; in a real scenario, this would be based on user data

# Example positions based on the image pattern
tree_positions = [(1, 0), (1, 3), (1, 5), (2, 2), (3, 6), (4, 1), (4, 4), (5, 0), (5, 3), (6, 5)]
house_positions = [(0, 1), (0, 4), (2, 1), (2, 4), (2, 6), (3, 0), (3, 3), (4, 7), (6, 2), (7, 3)]

# Initialize a matrix with small base probabilities
base_probability = 0.01
probability_matrix = np.full((grid_size, grid_size), base_probability)

# Assign higher probabilities to trees and lower to houses
tree_probability = 0.07
house_probability = 0.03
for position in tree_positions:
    probability_matrix[position] = tree_probability
for position in house_positions:
    probability_matrix[position] = house_probability

# Fixed unselectable points have a probability of 0
fixed_unselectable_points = [(0, 0), (7, 0), (0, 7), (7, 7), (3, 3), (3, 4), (4, 3), (4, 4)]
for point in fixed_unselectable_points:
    probability_matrix[point] = 0

# Normalize the probabilities to ensure they sum to 1
probability_matrix /= probability_matrix.sum()

# Create the heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(probability_matrix, annot=True, fmt=".2f", cmap="YlGnBu", square=True)
# plt.title('Heatmap of Grid Selection Probabilities Influenced by Image Elements')
# plt.xlabel('Grid Points (X Coordinate)')
# plt.ylabel('Grid Points (Y Coordinate)')

plt.title('RUGP：4-Gram模型概率的热图表示')
plt.xlabel('Grid Points (X 轴)')
plt.ylabel('Grid Points (Y 轴)')
plt.show()