#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Feb 15 11:38:46 2021
author: Xiaoting Yang
"""
import numpy as np
import matplotlib.pyplot as plt

plotdir='./Figures/';

# Problem 3(b)
def calc_rhs_VDP(x,y,t,gamma,tau,alpha,beta):
    P1=1; # years, time constant for forcing
    
    # forcing
    F=gamma*np.sin(2*np.pi*t/P1)
    
    dx=(1/tau)*(F+beta-y);
    dy=(alpha/tau)*(x+y-y**3/3)
    
    return dx,dy

P1=1.0;
dt=0.01; # time step in years
Tend=50.0; # end time of simulation in years

Nt=int(np.round(Tend/dt)) # the total number of time steps
T=np.linspace(0,Tend,Nt+1);

# initialize vectors
X=0*T; Y=0*T; DX=0*X; DY=0*Y; 
X[0]=-0.5; Y[0]=0.5;

# define parameters
gamma=3.0; tau=3.6; alpha=30.0; beta=0.75; 
tail='strong_forcing';

for i in range(Nt):
    dx,dy=calc_rhs_VDP(X[i],Y[i],T[i],gamma,tau,alpha,beta);
    
    DX[i]=dx; DY[i]=dy;
    
    X[i+1]=X[i]+dx*dt
    Y[i+1]=Y[i]+dy*dt

# make figure
plt.figure(figsize=(16,8)); plt.clf();
plt.subplot(1,2,1);
plt.plot(X,Y,color='b',linewidth=0.8);
plt.xlabel('X(t)'); plt.ylabel('Y(t)');

plt.subplot(1,2,2);
plt.plot(X[::(int(1/dt))],Y[::(int(1/dt))],'ko',markersize=6)
plt.xlabel('X(yearly)'); plt.ylabel('Y(yearly)');

#plt.savefig(plotdir+'VDP_'+tail+'1.png')

# make figure 2
plt.figure(figsize=(16,8)); plt.clf();
plt.subplot(1,2,1);
plt.plot(X,np.sin(2*np.pi*T/P1),color='b',linewidth=0.8);
plt.xlabel('X(t)'); plt.ylabel('Normalized forcing');

plt.subplot(1,2,2);
plt.plot(T,X,color='r',linewidth=0.8)
plt.xlabel('Time (year)'); plt.ylabel('X(t)');

#plt.savefig(plotdir+'VDP_'+tail+'2.png')