#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jun 18 22:25:10 2018

@author: mirjam
"""

import numpy as np
import matplotlib.pyplot as plt


def brownian(h=1, a = 1, M=1000):
    steps = np.int(100/h+1)
    times = np.zeros([1,steps])
    x = np.zeros([M,steps])
    mean_x = np.zeros([1,steps])
    var_x = np.zeros([1,steps])
    std_x = np.zeros([1,steps])
    
    for i in range(1,steps):
        x[:,i] = x[:,i-1] + a*np.random.randn(M)
        mean_x[:,i] = np.mean(x[:,i])
        var_x[:,i] = np.var(x[:,i])
        std_x[:,i] = np.std(x[:,i])
        times[:,i] = times[:,i-1] + h
    
    return x, times, mean_x, var_x, std_x


############# h = 1

bx, by, bmean, bvar, bstd = brownian(h=1,a=1, M=1000)

# plot of brownian trajectories
fig_traj = plt.figure()
for i in range(0,4):
    plt.plot(by.transpose(),bx[i,:].transpose(),'-')

# plot of mean, var and std
fig_mean = plt.figure()
plt.subplot(131)
plt.title("mean value")
plt.plot(by.transpose(),bmean.transpose(),'-')
plt.subplot(132)
plt.title("variance")
plt.plot(by.transpose(),bvar.transpose(),'-')
plt.subplot(133)
plt.title("std")
plt.plot(by.transpose(),bstd.transpose(),'-')


################ h = 0.01 or 0.1 => a=sqrt(h)

bh = 0.01 # 0.1
bx_h, by_h, bmean, bvar, bstd = brownian(h=bh, a=np.sqrt(bh), M=10)

# plot of brownian trajectories
fig_traj = plt.figure()
for i in range(0,4):
    plt.plot(by_h.transpose(),bx_h[i,:].transpose(),'-')

# plot of mean, var and std
fig_mean = plt.figure()
plt.subplot(131)
plt.title("mean value")
plt.plot(by_h.transpose(),bmean.transpose(),'-')
plt.subplot(132)
plt.title("variance")
plt.plot(by_h.transpose(),bvar.transpose(),'-')
plt.subplot(133)
plt.title("std")
plt.plot(by_h.transpose(),bstd.transpose(),'-')

################# now sampled plot

y_sample = range(0,101,2) # positions for sample step size

#get the subsets of the brownian motion processes
bx_sample = bx[:,y_sample]
bxh_sample = bx_h[:,[np.int(i/bh) for i in y_sample]] 

# plot of sampled trajectories
fig_samp = plt.figure()
plt.subplot(121)
plt.title("h=1")
for i in range(0,4):
    plt.plot(y_sample,bx_sample[i,:].transpose(),'-')
plt.subplot(122)
plt.title("h="+ str(bh))
for i in range(0,4):
    plt.plot(y_sample,bxh_sample[i,:].transpose(),'-')


#############################################################
############### part2 van der pol ###########################
#############################################################

# non stochastic van der Pol:
def vdp(y, mu):
    dy = np.zeros(2)
    dy[0] = y[1]
    dy[1] = mu*(1 - y[0]**2)*y[1] - y[0] 
    return dy

def euler(times, h, y0, mu = 1, a = 0):
    y = np.zeros([2,len(times)])
    y[:,0] = y0
    for i in range(0,len(times)-1):
        yr = y[:,i]
        t = times[i]
        while t < times[i+1]:
          yr = yr + h*vdp(yr,mu) + a*h**0.5*np.array([0,np.random.randn(1)])
          t = t + h
          
        y[:,i+1] = yr
    return y
    


def plot_vdP(h, a, mu = 1, maxtime = 50, y0 = np.array([1,1])):
    times = np.linspace(0,maxtime,num=np.int(maxtime/0.1 + 1))
    y_1 = euler(times, h, y0, mu = mu, a =a ) 
    print("h="+str(h)+", a="+str(a)+", mu="+str(mu)+":")
    plt.figure()
    plt.title("test")
    plt.subplot(121)
    plt.title("van der Pol (config space)")
    plt.plot(times, y_1[0,:], '-', label = "x1")
    plt.plot(times, y_1[1,:], '-', label = "x2")
    plt.xlabel("time")
    plt.legend(loc='best')
    plt.subplot(122)
    plt.title("van der Pol (phase space)")
    plt.plot(y_1[0,:], y_1[1,:], '-')
    plt.show()

# plot some h, a, mu combinations:
plot_vdP(h = 0.01, a = 0)
plot_vdP(h = 0.01, a = 0.5)
plot_vdP(h = 0.01, a = 1)
plot_vdP(h = 0.01, a = 2)
plot_vdP(h = 0.1, a = 0)
plot_vdP(h = 0.001, a = 0)
#plot_vdP(h = 0.1, a = 0, mu = 5) # crashes
plot_vdP(h = 0.01, a = 0, mu = 5)
plot_vdP(h = 0.01, a = 1, mu = 5)