import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multinomial
from matplotlib.patches import RegularPolygon
import os
from tqdm import trange

M, N = 100000, 10
for N in trange(2, 200):
    p = np.array([0.2, 0.3, 0.5])

    samples = multinomial.rvs(N, p, size=M).T
    K = np.array([[-np.sqrt(1/2), np.sqrt(1/2), 0], [-np.sqrt(1/6), -np.sqrt(1/6), np.sqrt(4/6)]])
    result = np.dot(K, samples) / N

    triangle_vertices = np.array([K[:, 0], K[:, 1], K[:, 2]])

    def f(x, y):
        return -N/2 * np.sum((np.array([1/3, 1/3, 1/3]) + x * K[0,:] + y*K[1,:] - p)**2 / p, axis=0)

    x_values = np.linspace(-np.sqrt(1/2), np.sqrt(1/2), 50)
    y_values = np.linspace(-np.sqrt(1/6), np.sqrt(4/6), 50)
    X, Y = np.meshgrid(x_values, y_values)

    Z = np.zeros_like(X)
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Z[i, j] = f(X[i, j], Y[i, j])
            
    hexbin_x = result[0]
    hexbin_y = result[1]

    plt.figure(figsize=(10, 10 * np.sqrt(3)))
    plt.hexbin(hexbin_x, hexbin_y, gridsize=50, cmap='YlGnBu', extent=(min(result[0]), max(result[0]), min(result[1]), max(result[1])),
              bins='log', mincnt=1, alpha=0.7, edgecolors='gray', linewidths=0.1)

    # Overlay heatmap of function f within the equilateral triangle
    plt.imshow(Z, extent=(-np.sqrt(1/2), np.sqrt(1/2), -np.sqrt(1/6), np.sqrt(4/6)),
              origin='lower', cmap='coolwarm', alpha=0.5)

    # Plot equilateral triangle
    triangle = plt.Polygon(triangle_vertices, edgecolor='black', closed=True, fill=False)
    plt.gca().add_patch(triangle)

    plt.xlim(-np.sqrt(1/2), np.sqrt(1/2))
    plt.ylim(-np.sqrt(1/6), np.sqrt(4/6))

    plt.title(f"N={N}, p={p}")
    plt.gca().set_aspect('equal', adjustable='box')
    plt.axis('off')
    dir_path = f"./multinomial"
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    plt.savefig(f"{dir_path}/{N:03d}.png",bbox_inches='tight')

    plt.close()

import imageio.v3 as iio
import os
from natsort import natsorted
import moviepy.editor as mp

for dir_path in ["./multinomial"]:
    file_names = natsorted((fn for fn in os.listdir(dir_path) if fn.endswith('.png')))

    # Create a list of image files and set the frame rate
    images = []
    fps = 12

    # Iterate over the file names and append the images to the list
    for file_name in file_names:
        file_path = os.path.join(dir_path, file_name)
        images.append(iio.imread(file_path))

    filename = dir_path[2:]
    clip = mp.ImageSequenceClip(images, fps=fps)
    clip.write_videofile(f"{filename}.mp4")

!ffmpeg -i multinomial.mp4 -c:v libvpx-vp9 -b:v 0 -crf 10 -c:a libvorbis multinomial.webm