#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" Plot mechanical energy evolution
    in the Marsaleix et al (2008) test case (IW1)
"""

from netCDF4 import Dataset
import matplotlib.pyplot as plt
import numpy as np

rau0 = 1026.
grav = 9.81
nexp = 5 

def comp_ke(filedatat, filedatau ):
    global rau0
    global grav

    # Open file at T-points:
    ncid = Dataset(filedatat)
    time = ncid.variables['time_counter'][:] # Time [s]
    lon = ncid.variables['nav_lon'][1, :] # Longitude [km]
    ssh = ncid.variables['ssh_inst'][:, 1, :] # Read sea level anomaly [m]
    e3t = ncid.variables['e3t_inst'][:, :, 1, :] # thicknesses [m]
    temp = ncid.variables['thetao_inst'][:, :, 1, :] # temperature [°C]
    e3t = np.where((temp == 0.), 0., e3t)
    ncid.close()

    # Open file at U-points:
    ncid = Dataset(filedatau)
#    ubar = ncid.variables['ubar_inst'][:, 1, :] # Read barotropic velocity [m/s]
    e3u = ncid.variables['e3u_inst'][:, :, 1, :] # U thicknesses [m]
    uu = ncid.variables['uo_inst'][:, :, 1, :] # velocity [m/s] 
    ubar = np.zeros(np.shape(ssh))
    hu = np.zeros(np.shape(ssh))
    e3u = np.where((uu == 0.), 0., e3u)
    hu = np.sum(e3u, axis=(1))
    hu = np.where((hu == 0.), 1., hu)
    ubar = np.sum(uu*e3u, axis=(1))/hu
    ncid.close()

    # Set velocity at T-points:
    (jpt, jpi) = np.shape(ssh)
    work = np.zeros(np.shape(ssh))
    work[:, 1:jpi - 1] = 0.5 * (ubar[:, 1:jpi - 1] + ubar[:, 0:jpi - 2])

    # mechanical energy:
    dx = (lon[1]-lon[0])*1.e3    # Horizontal resolution in m
    dt = (time[1]-time[0])
    dl = (jpi-2)*dx/1.e3
    h = np.zeros((jpt,1)) # Total depth
    h = np.sum(e3t, axis=(1))
    ke = np.zeros((jpt, 1)) # energy
    ke = np.sum(0.5*rau0*dx*dx* (h*work**2 + grav*ssh**2), axis=(1))
    ssh0 = lon / (0.5*dl)
#    ssh0 = np.sin(np.pi*lon/dl)
    ssh0[0] = 0.
    ssh0[jpi-1] = 0.
    ke0 = np.sum(0.5*rau0*dx*dx*grav*ssh0**2, axis=(0))
    ke = ke/ke0
    # Set time from time increment:
    time = np.arange(dt,(jpt+1)*dt,dt)

    return time, ke  

fig = plt.figure(figsize=(7, 5))
for exp in np.arange(1, nexp+1,1):
    filet='SEICHE_mar_rk3_dt_%i_grid_T.nc' %exp
    fileu='SEICHE_mar_rk3_dt_%i_grid_U.nc' %exp
    print(filet,fileu)
    (time, ke) = comp_ke(filet, fileu)
    plt.plot(time/86400., ke)

plt.ylim((0.6, 1.05))
plt.xlim((0., 12.))
plt.xlabel('Time [days]')
plt.ylabel(r'$\frac{E(t)}{E(0)}$',rotation=0)
plt.grid()
plt.title('DM19 test case - Barotropic energy vs baroclinic time step - RK3')
plt.legend(('960s','480s','240s','120s','60s'),loc='lower left')
plt.savefig('SEICHE_mar_rk3_dt.png')
plt.show()