#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jun  18 10:41:33 2018

@author: helge
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate
from timeit import default_timer as timer
from scipy import signal

# General setting
sheet_part_b = 0

t0 = 0
tend = 100

if sheet_part_b:
    mu = 1000
else:
    mu = 5

# Runge-Kutta 4th order own implementation
def rk4(f,t0,tend,dt,y0):
    t4 = np.empty(int(np.ceil((tend-t0)/dt))+1)
    y4 = np.empty([y0.size,int(np.ceil((tend-t0)/dt))+1])

    #Initialize
    t4[0] = t0
    cur_t = t0
    y4[:,0] = y0
    n=0
    while cur_t < tend:
        t4[n+1] = t4[n] + dt      
        k1 = f(t4[n], y4[:,n])
        k2 = f(t4[n] + dt/2, y4[:,n] + np.multiply(k1,dt/2))
        k3 = f(t4[n] + dt/2, y4[:,n] + np.multiply(k2,dt/2))
        k4 = f(t4[n] + dt, y4[:,n] + np.multiply(k3,dt) )
        phi = (1/6) * (k1 + 2*k2 + 2*k3 + k4)
        y4[:,n+1] = y4[:,n] + phi * dt;
        n += 1
        cur_t += dt
    return (t4,y4)

def vdp(t,y):
    dy = np.zeros(2)
    dy[0] = y[1]
    dy[1] = mu*(1-y[0]**2)*y[1] - y[0]
    return dy

#Start values for y
y0 = np.array(np.random.randn(2))

#Integration via own RK4
dt = 0.01
if sheet_part_b==0:
    RK4_out = rk4(vdp,t0,tend,dt,y0)

# PLOTTING of own RK4 integrator in time and phase-space
if sheet_part_b==0:
    plt.figure(1)
    plt.subplot(211)
    plt.xlabel("time")
    plt.ylabel("Y values")

    plt.plot(RK4_out[0].transpose(),RK4_out[1].transpose())
    plt.title("Van-der-Pol RK4 integration")

    plt.subplot(212)
    plt.xlabel("X values")
    plt.ylabel("Y values")
    RK4_ys = RK4_out[1]
    plt.plot(RK4_ys[0,:],RK4_ys[1,:])
    plt.title("Van-der-Pol in phase space")

#Integration via python RK45

start45 = timer()
RK45_out = integrate.RK45(vdp,t0,y0,tend)
t_45 = np.array([t0])
y_45 = np.empty([2,1])
y_45[:,0] = y0.transpose()
while RK45_out.status is 'running':
    RK45_out.step()
    t_45 = np.append(t_45,RK45_out.t)
    y_tmp = np.empty([2,1])
    y_tmp[:,0] = RK45_out.y
    y_45 = np.append(y_45,y_tmp,axis = 1)
end45 = timer() 
 
# Integration via python BDF
if sheet_part_b:
    startBDF = timer()
    BDF_out = integrate.BDF(vdp,t0,y0,tend)
    t_BDF = np.array([t0])
    y_BDF = np.empty([2,1])
    y_BDF[:,0] = y0.transpose()
    while BDF_out.status is 'running':
        BDF_out.step()
        t_BDF = np.append(t_BDF,BDF_out.t)
        y_tmp = np.empty([2,1])
        y_tmp[:,0] = BDF_out.y
        y_BDF = np.append(y_BDF,y_tmp,axis = 1)
    endBDF = timer() 

if sheet_part_b:
    print("RK4 integration time " + str(end45 - start45))
    print("BDF integration time " + str(endBDF - startBDF))

# PLOTTING of own and implement RK4, BDF integrators
plt.figure(2)
plt.subplot(211)

if sheet_part_b==1:
    plt.xlabel("time")
    plt.ylabel("Y values")
    plt.plot(t_BDF.transpose(),y_BDF.transpose())
    plt.title("Van-der-Pol implicit BDF integration")
else:
    RK45_steps = t_45[1:]-t_45[0:-1]
    plt.plot(t_45[1:],RK45_steps)
    plt.xlabel("time")
    plt.ylabel("step size")
    plt.title("Adaptive RK45 step sizes")
    
plt.subplot(212)
plt.plot(t_45.transpose(),y_45.transpose())
plt.title("Van-der-Pol RK4 integration")


# Additional challenge   
#ys = RK4_out[1]
#peaks = signal.find_peaks_cwt(ys[0,:],np.arange(1,1100))
#f, Van_peri = signal.periodogram(ys[0,np.arange(peaks[1],peaks[5])])
#plt.figure(3)
#plt.semilogy(f,Van_peri)