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

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

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

amu_list = [(10, 0), (0, 0), (-10, 0), (0, -10)]
sigmaX = 1. # standard deviation of X
tau = 1. # relaxation time scale
theta = 1. / tau # mean reversion rate
D = sigmaX**2 * theta # diffusion constant
sigma = sqrt(2 * D) # diffusion
t = np.linspace(0, 12, 2001)
dt = t[1:] - t[:-1]

fig = plt.figure(figsize=(520 / 90.0, 340 / 90.0), dpi=72)

for a, mu in amu_list:
    randnorm = np.random.normal(0, 1, len(t))
    X = np.empty_like(t)
    X[0] = a

    for i in range(1, len(t)):
        X[i] = X[i-1] + (mu - X[i-1]) * dt[i-1] * theta + sigma * sqrt(dt[i-1]) * randnorm[i]
    X_expect = mu + (a - mu) * np.exp(-t * theta)
    
    plt.plot(t, X, lw=1, label=rf'a ={a:3.0f}  $\mu$ ={mu:3.0f}'.replace(' ', '  ').replace('-', u'\u2212'))
    plt.plot(t, X_expect, lw=2.5, color='#cccccc', zorder=-1,
        label={False:None, True:'E(X)'}[mu==-10])

plt.grid(True)
plt.xlim(t[0], t[-1])
plt.ylim(-13, 13)
plt.xlabel(r't / $\theta^{-1}$')
plt.ylabel(r'X / $\sigma_X$')
plt.legend(loc='upper right', framealpha=1)
plt.tight_layout()
plt.savefig('Ornstein-Uhlenbeck-traces-a-mu.svg')