import numpy as np
import os
import matplotlib.pylab as plt
import matplotlib
import cartopy.crs as ccrs
import cmocean.cm as cm
import cartopy.feature as cfeature

from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation
from IPython.display import Video

# font size for plotting
fs = 30
font = {'size'   : 18}
matplotlib.rc('font', **font)

# extend of the grid:
minlat = 20
maxlat = 50
minlon = 100
maxlon = 200
# define the grid
X,Y = np.meshgrid(np.arange(minlon,maxlon,3), np.arange(minlat, maxlat, 2))
# define time and Earth rotation speed
rotateperday = 0.5 # degrees per day
days = np.arange(0,180,0.15)
midlons = 100 + rotateperday*days

filedir = os.path.abspath(".")
vdir = os.path.abspath(os.path.join(filedir, "..", "images"))
vfile = 'RotatingEarth.mp4'
vfile_path = os.path.join(vdir, vfile)

# Functions that define the time varying flow field
def xcon(x, lat, r0):
    x =  (x - minlon) / (maxlon-minlon) * r0*np.pi
    return x

def ycon(y):
    y = (y-minlat) / (maxlat-minlat) * 8000 - 4000
    return y

def BickleyJetflow(x, y, t, midlat = 35):
    """
    Return the Bickley Jet flow field

    """
    ua = np.zeros(x.shape)
    va = np.zeros(x.shape)
    
    Ubj =  0.06266
    L = 1700.
    r0 = 6371.
    k1 = 2 * 1/ r0
    k2 = 2 * 2 / r0
    k3 = 2 * 3/ r0
    eps1 = 0.075
    eps2 = 0.4
    eps3 = 0.3
    c3 = 0.461 * Ubj
    c2 = 0.205 * Ubj
    c1 = 0.1446 * Ubj
    
    x = xcon(x, midlat, r0=r0)
    y = ycon(y)
    
    def v(x1, x2, t): #2D Bickley Jet velocity given single location
        f1 = eps1 * np.exp(-1j *k1 * c1 * t)
        f2 = eps2 * np.exp(-1j *k2 * c2 * t)
        f3 = eps3 * np.exp(-1j *k3 * c3 * t)
        F1 = f1 * np.exp(1j * k1 * x1)
        F2 = f2 * np.exp(1j * k2 * x1)
        F3 = f3 * np.exp(1j * k3 * x1)
        G = np.real(np.sum([F1,F2,F3]))
        G_x = np.real(np.sum([1j * k1 *F1, 1j * k2 * F2, 1j * k3 * F3]))
        u =  Ubj / (np.cosh(x2/L)**2)  +  2 * Ubj * np.sinh(x2/L) / (np.cosh(x2/L)**3) *  G
        vd = Ubj * L * (1./np.cosh(x2/L))**2 * G_x
        return [u,vd]
    
    for i in range(x.shape[0]):
        for j in range(y.shape[1]):
            ua[i,j], va[i,j] = v(x[i,j], y[i,j], t*36000)
    res = np.sqrt(ua**2+va**2)

    return res

# The initial plot

def plot(t): 
    data = BickleyJetflow(X, Y, days[t])
    
    fig = plt.figure(figsize=(6,5))
    projection = ccrs.NearsidePerspective(midlons[t], 20)

    ax = plt.subplot(111, projection=projection)
    ax.add_feature(cfeature.OCEAN, zorder=0)
    ax.add_feature(cfeature.LAND, zorder=0, edgecolor='black')
    ax.set_global()
    tit = ax.set_title('day %.1f'%(days[t]), fontsize=fs)

    im = ax.pcolor(X, Y, data, transform=ccrs.PlateCarree(), cmap=cm.speed,
                        vmin=0, vmax=0.1)

    #% add colorbar  
    fig.subplots_adjust(right=0.8)
    cbar_ax = fig.add_axes([0.8, 0.15, 0.1, 0.7])
    cbar_ax.set_visible(False)
    
    cbar = fig.colorbar(im, ax=cbar_ax, orientation = 'vertical', fraction = 0.8,
                    aspect=18, extend='max')
    cbar.ax.xaxis.set_label_position('bottom')
    cbar.ax.set_xlabel('m/s', fontsize=fs)
    return fig, tit, cbar_ax

# The animation

def animate(t):
    if(t%25==0):
        print(t/len(days),'%')
    data = BickleyJetflow(X, Y, days[t])
    projection = ccrs.NearsidePerspective(midlons[t], 20)

    ax = plt.subplot(111, projection=projection)
    ax.add_feature(cfeature.OCEAN, zorder=0)
    ax.add_feature(cfeature.LAND, zorder=0, edgecolor='black')
    ax.set_global()
    tit = ax.set_title('day %.1f'%(days[t]), fontsize=fs)

    ax.set_transform(projection)
    tit.set_text('day %.1f'%(days[t]))
    im = ax.pcolor(X, Y, data, transform=ccrs.PlateCarree(), cmap=cm.speed,
                        vmin=0, vmax=0.1)

if(__name__=='__main__'):
    if not os.path.exists(vfile_path):
        fig, tit, cbar_ax = plot(0)
        anim = FuncAnimation(fig, animate, frames = np.arange(len(days)), interval=10)
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=20, metadata=dict(artist='Me'), bitrate=1250)
        anim.save(vfile_path, writer=writer, dpi=92)
        plt.close()