import numpy as np
import matplotlib.pyplot as plt

n_pts  = 5000
seed   = 2
alpha  = 0.1 
msize  = 10
nbins  = 40

rng = np.random.default_rng(seed)

# 1) sample data  (cls0 not Gaussian)
cls0 = np.column_stack([
    rng.uniform(-1, 0.3, n_pts),
    rng.uniform(-1.5,  1.5, n_pts)
])
cls1 = rng.normal(loc=[0.8, 0.0], scale=[0.6, 0.2], size=(n_pts, 2))

def plotter(cls0, cls1, theta):
    mean0, mean1 = cls0.mean(0), cls1.mean(0)
    grand_mean   = (cls0.sum(0) + cls1.sum(0)) / (2 * n_pts)
    cls0_c, cls1_c = cls0 - grand_mean, cls1 - grand_mean
    
    R = np.array([[np.cos(theta), -np.sin(theta)],
                  [np.sin(theta),  np.cos(theta)]])
    cls0_r, cls1_r = cls0_c @ R.T, cls1_c @ R.T
    
    # 4) projection to x‑axis
    x0, x1 = cls0_r[:, 0], cls1_r[:, 0]
    
    # variance decomposition on projection
    n0 = n1 = n_pts
    m0, m1  = x0.mean(), x1.mean()
    overall = (x0.sum() + x1.sum()) / (n0 + n1)
    within  = (n0 * x0.var(ddof=1) + n1 * x1.var(ddof=1)) / (n0 + n1)
    between = (n0 * (m0 - overall) ** 2 + n1 * (m1 - overall) ** 2) / (n0 + n1)
    ratio   = between / within
    
    # === plotting ===
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(7, 10),
                                   gridspec_kw={'height_ratios': [1, 1.3]},
                                   constrained_layout=True)
    
    # -- Plot 1 : histogram of projections --
    bins = np.histogram_bin_edges(np.concatenate([x0, x1]), bins=nbins)
    ax1.hist(x0, bins=bins, alpha=0.6, color='red',  label='Class 0')
    ax1.hist(x1, bins=bins, alpha=0.6, color='blue', label='Class 1')
    # ax1.legend(loc='upper left', frameon=False)
    ax1.set_xlim([-2, +2])
    ax1.set_ylim([0, +450])
    ax1.set_yticks([])
    
    ax1.text(0.75, 0.95,
             f"Within Var: {within:.2f}
Between Var: {between:.2f}
Ratio: {ratio:.2f}",
             transform=ax1.transAxes, ha='left', va='top', fontsize=10)
        
    # -- Plot 2 : scatter of rotated, centred data --
    ax2.scatter(cls0_r[:, 0], cls0_r[:, 1], c='red',  alpha=alpha,
                marker='x', s=msize, label='Class 0')
    ax2.scatter(cls1_r[:, 0], cls1_r[:, 1], c='blue', alpha=alpha,
                marker='x', s=msize, label='Class 1')
    
    # class means: filled squares with thin black edge
    m0_r = (mean0 - grand_mean) @ R.T
    m1_r = (mean1 - grand_mean) @ R.T
    ax2.set_aspect(1)
    ax2.set_xlim([-2, +2])
    ax2.set_ylim([-2, +2])
    
    ax2.scatter(m0_r[0], m0_r[1], marker='s', s=60,
                facecolor='red', edgecolor='k', linewidth=0.7)
    ax2.scatter(m1_r[0], m1_r[1], marker='s', s=60,
                facecolor='blue', edgecolor='k', linewidth=0.7)
    
    # grand mean: black square
    ax2.scatter(0, 0, marker='s', s=60, color='black')
    
    ax2.set_xticks([]); ax2.set_yticks([])
    for spine in ax2.spines.values():
        spine.set_visible(False)
    
    return fig

import os
from tqdm import tqdm
movie_name = 'lda'
dir_path = f"./{movie_name}"
tmax = 12
frames_per_t = 24
t_values = np.linspace(0, 2 * np.pi, frames_per_t * tmax)

for N, t in tqdm(enumerate(t_values)):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

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

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

for dir_path in [dir_path]:
    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 = 24

    # 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"output.mp4")

!ffmpeg -framerate 24 -i %03d.png output.webm
!ffmpeg -y -i output.mp4 -c:v libvpx-vp9 -b:v 0 -crf 10 -c:a libvorbis output.webm