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

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

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

fname = "Wiener-process-zoom-animation"
nframes = 80
fx = 0.25
nsamples = int(1000 / fx)

# create pseudo-random sequence for continous adding of time points
golden = (sqrt(5) - 1) / 2
t_unsorted = (golden * np.arange(nsamples)) % 1
sort_idx = np.argsort(t_unsorted)
t = t_unsorted[sort_idx]
randnorm = np.random.normal(0, 1, nsamples)
X = np.zeros(nsamples)
X[-1] = sqrt(t[-1]) * randnorm[-1]
X[int(fx * nsamples)] = sqrt(fx) * X[-1]

# use Brownian bridges
for i in range(int(fx * nsamples) + 1, nsamples - 1):
    mean = X[i-1] + (t[i] - t[i-1]) / (t[-1] - t[i-1]) * (X[-1] - X[i-1])
    var = (t[-1] - t[i]) * (t[i] - t[i-1]) / (t[-1] - t[i-1])
    X[i] = mean + sqrt(var) * randnorm[i]

# use self-similar signal for small t-values
mask0 = sort_idx < nsamples * fx**(1 - 1 / nframes)
t0, X0 = t[mask0], X[mask0]
for i in range(int(fx * nsamples) - 1, 0, -1):
    tfx = t[i] / fx
    i2 = np.searchsorted(t0, tfx)
    i1 = i2 - 1
    if tfx - t0[i1] < t0[i2] - tfx:
        X2 = X0[i1]
    else:
        X2 = X0[i2]
    X[i] = sqrt(fx) * X2
    X0 = X[mask0] # update masked X0 from X (a bit inefficient)

ylim = -0.6, 1.4
def animate(nframe):
    plt.cla()
    xscale = fx**(nframe / nframes)
    yscale = sqrt(xscale)
    mask = sort_idx < nsamples * fx / xscale
    plt.plot(t[mask], X[mask], '-')
    plt.xlim(0, xscale)
    plt.ylim(ylim[0] * yscale, ylim[1] * yscale)
    plt.grid(True)
    plt.xlabel('t')
    plt.ylabel('X')
    if xscale > sqrt(fx):
        muly = 0.125
    else:
        muly = 0.0625
    plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(muly))
    plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(0.3))
    plt.gca().axes.xaxis.set_ticklabels([])
    plt.gca().axes.yaxis.set_ticklabels([])


fig = plt.figure(figsize=(520 / 90.0, 340 / 90.0), dpi=72)
plt.xlim(0, 1)
plt.ylim(ylim[0], ylim[1])
plt.grid(True)
plt.xlabel('t')
plt.ylabel('X')
plt.gca().axes.xaxis.set_ticklabels([])
plt.gca().axes.yaxis.set_ticklabels([])
plt.tight_layout()
anim = animation.FuncAnimation(fig, animate, frames=nframes)
anim.save(fname + '.gif', writer='imagemagick', fps=20)