#! /usr/bin/env python3
# -*- coding:utf8 -*-

import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
from math import *

plt.rcParams['font.sans-serif'] = 'DejaVu Sans'
np.random.seed(1)

fname = "Pink-noise-density-animation"
nframes = 23
nsamples = 2**(nframes+1)
t = (2 / nsamples) * np.arange(nsamples)
white_noise = np.random.normal(0, 1, nsamples)
fourier_amplitudes = np.fft.rfft(white_noise)

def white_resample(X0, n):
    N = len(X0)
    resampled = sqrt(n) * np.roll(X0, n // 2).reshape((N // n, n)).mean(axis=1)
    return resampled

def pink_filter(X0):
    frequencies = np.fft.rfftfreq(len(X0), d=2 / len(X0))
    fourier_amplitudes = np.fft.rfft(X0)
    fourier_amplitudes[1:] /= np.sqrt(frequencies)[1:] # 1/sqrt(f) amplitude spectrum
    X = np.fft.irfft(fourier_amplitudes, n=len(X0), norm='ortho')
    return X

# determine standard deviation at 512 samples resolution
X = pink_filter(white_resample(white_noise, 2**(nframes - 9)))
pre = 1 / np.std(X[:len(X) // 2 + 1])


def animate(nframe):
    exponent = (nframe + 9) % nframes # start with frame 9
    plt.cla()
    dsample = 2**(nframes - exponent)
    white_reduced = white_resample(white_noise, dsample)
    X = pre * pink_filter(white_reduced)
    plt.plot(t[::dsample], X, '.-', markersize=4,
        label=f"$\Delta$t = 1/{nsamples//(2*dsample):.0f}")
    plt.xlim(0, 1)
    plt.ylim(-3, 3)
    plt.grid(True)
    plt.xlabel('t')
    plt.ylabel('X')
    plt.legend(loc="upper center", framealpha=1)

fig = plt.figure(figsize=(520 / 90.0, 340 / 90.0), dpi=72)
plt.xlim(0, 1)
plt.ylim(-3, 3)
plt.grid(True)
plt.xlabel('t')
plt.ylabel('X')
plt.tight_layout()
anim = animation.FuncAnimation(fig, animate, frames=nframes)
anim.save(fname + '.gif', writer='imagemagick', fps=2)