# demo of langevin sampling
# https://yang-song.net/blog/2021/score/

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

np.random.seed(42)

TIMESTEPS = 100
NUM_SAMPLES = 100

std = 0.1

mean1 = 0
mean2 = 1

mixing_coefficient = 0.5

# plotting setup
fig, ax = plt.subplots()
ax.set_xlim(-1, 2)
ax.set_ylim(-1, 2)

# initialize chain from arbitrary distribution, let's say uniform
initial = np.random.uniform(-1, 2, (NUM_SAMPLES, 2))

step_size = 1e-3


def gradient_log_likelihood(x):
    def score_component(x, mu):
        return -(x - mu) / (std**2)

    def density_component(x, mu):
        return np.exp(-np.sum((x - mu) ** 2, axis=1) / (2 * std**2))

    p1 = density_component(x, mu=np.array([mean1, mean1]))
    p2 = density_component(x, mu=np.array([mean2, mean2]))

    score1 = score_component(x, mu=np.array([mean1, mean1]))
    score2 = score_component(x, mu=np.array([mean2, mean2]))

    numerator = (
        mixing_coefficient * p1[:, None] * score1
        + (1 - mixing_coefficient) * p2[:, None] * score2
    )
    denominator = (mixing_coefficient * p1 + (1 - mixing_coefficient) * p2)[:, None]

    return numerator / denominator


def update_particles(previous):
    gll = gradient_log_likelihood(previous)
    noise = np.sqrt(2 * step_size) * np.random.normal(0, 1, (NUM_SAMPLES, 2))
    return previous + step_size * gll + noise


current = initial.copy()
scat = ax.scatter(initial[:, 0], initial[:, 1], alpha=0.3, label="Particles")


def update(frame):
    global current
    current = update_particles(current)
    scat.set_offsets(current)
    return (scat,)


ax.scatter(np.array([0, 1]), np.array([0, 1]), label="Ground-Truth Means")

ani = FuncAnimation(fig, update, frames=TIMESTEPS, interval=50, blit=True, repeat=False)
plt.legend()
plt.title("Langevin Sampling on Mixture of Two Gaussians")
ani.save("langevin.gif", writer="ffmpeg", fps=20)