#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Feb 21 21:24:57 2021

@author: Xiaoting Yang
"""
## Welander1982_flip_flop
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# define constants
k0=5.0; eps=-0.01;
alpha=0.2; beta=1.0;
T_d=0.0; S_d=0.0;
Sa=1.0; Ta=1.0;
A=1.0; B=0.1;

# define functions to be used
def EOS(T,S):
    # equation of state
    res=-alpha*T+beta*S;
    return res

def K(rho):
    # convective mixing coefficient
    if rho>eps:
        res=k0;
    else:
        res=0.0;
    
    return res

def system_eqs(t,y):
    # RHS of time dependent ODEs for temperature and salinity in Welander model
    T=y[0]; S=y[1];
    rho=EOS(T,S);
    
    dTdt=A*(Ta-T)+K(rho)*(T_d-T)
    dSdt=B*(Sa-S)+K(rho)*(S_d-S)
    
    res=np.array([dTdt,dSdt])
    
    return res

# calcualte convective and nonconvective steady solutions:
T_nonconv=Ta; S_nonconv=Sa;
T_conv=(A*Ta+k0*T_d)/(A+k0)
S_conv=(B*Sa+k0*S_d)/(B+k0)

# initialize the model
time_max=600.0;
Tspan=(0,1*time_max);
X0=[T_conv+0.01,S_conv+0.01]

# solve for time evolution:
sol=solve_ivp(system_eqs,Tspan,X0,method='RK45',rtol=1.0e-7,atol=1.0e-8);

Time=sol.t;
Temp=sol.y[0];
Salt=sol.y[1];

# make figure
plt.figure(figsize=(9,9)); plt.clf();
ax1=plt.subplot(2,1,1); 
ax1.plot(Time,-alpha*Temp,color='r',linewidth=1.2,label=r'$-\alpha T$');
ax1.plot(Time,beta*Salt,color='g',linewidth=1.2,label=r'$\beta S$');
ax1.plot(Time,EOS(Temp,Salt),color='b',linewidth=1.2,label=r'$\rho$');
ax1.plot(Time,eps*np.ones(Time.shape),color='k',linestyle='--',linewidth=0.8,label=r'$\epsilon$')
plt.legend(loc='upper left')
ax1.set_xlabel('Time');
ax1.set_ylabel(r'$T,\,\, S,\,\, \rho$')
ax1.set_title('Welandar flip-flop oscillations');
plt.xlim(Time[0],np.minimum(20,Time[-1]))
plt.grid(lw=0.25)

ax2=plt.subplot(2,2,3);
ax2.plot(T_conv,S_conv,color='g',marker='*',markersize=10,label='convecting')
ax2.plot(T_nonconv,S_nonconv,color='b',marker='*',markersize=10,label='non-convecting')
b=round(0.5*len(Time));
T=Temp[b:]; S=Salt[b:]
for i in range(len(T)):
    rho=EOS(T[i],S[i])
    
    if rho>eps:
        ax2.plot(T[i],S[i],color='g',marker='.',markersize=1.5);
    else:
        ax2.plot(T[i],S[i],color='b',marker='.',markersize=1.5);

plt.legend(loc='upper left');
ax2.set_xlabel('T'); ax2.set_ylabel('S'); ax2.set_title('Phase space');

ax3=plt.subplot(2,2,4); 
for i in range(len(T)):
    rho=EOS(T[i],S[i]);
    
    if rho>eps:
        ax3.plot(T[i],S[i],color='g',marker='.',markersize=1.5);
    else:
        ax3.plot(T[i],S[i],color='b',marker='.',markersize=1.5);
    
dv=0.01; 
plt.axis([0.16,0.82,0.02,0.16]);
V=plt.axis();

x,y=np.meshgrid(np.arange(V[0],V[1],dv),np.arange(V[2],V[3],dv));
density_data=EOS(x,y);
VV=[-0.09,-0.06,-0.03,eps,0.0,0.03,0.06,0.09];
CS=ax3.contour(x,y,density_data,VV,colors='k',linewidths=1)
ax3.clabel(CS)
plt.ylim(0.02,0.15)
plt.show()

#plt.savefig('./Figures/Flip_flop_no_oscillation.png')