import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import os
from math import gcd

# Ensure output directory exists
os.makedirs("/mnt/data", exist_ok=True)

# Define the list S
S = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
n = len(S)

# Create the GCD matrix
gcd_matrix = np.zeros((n, n), dtype=int)
for i in range(n):
    for j in range(n):
        gcd_matrix[i, j] = gcd(S[i], S[j])

# Plot the heatmap
plt.style.use('seaborn-v0_8')
fig, ax = plt.subplots(figsize=(8, 6))
cax = ax.matshow(gcd_matrix, cmap='viridis')

# Add colorbar
fig.colorbar(cax)

# Set axis labels
ax.set_xticks(np.arange(n))
ax.set_yticks(np.arange(n))
ax.set_xticklabels(S)
ax.set_yticklabels(S)
plt.xlabel("S values")
plt.ylabel("S values")
plt.title("Heatmap of GCD Matrix")

# Rotate x-axis labels for better readability
plt.xticks(rotation=45)

# Save the figure
output_path = "/mnt/data/gcd_matrix_heatmap.png"
plt.tight_layout()
plt.savefig(output_path)
plt.close()