import sys
print 'Loading modules ...'
sys.stdout.flush()

import math as math
# select non-interactive backend, cf http://stackoverflow.com/questions/4931376/generating-matplotlib-graphs-without-a-running-x-server
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt
import numpy as np

print 'Loading DYNAMICO modules ...'
sys.stdout.flush()
from dynamico import unstructured as unst
from dynamico import time_step
print '...Done'
sys.stdout.flush()

grid, llm, nqdyn = 10242, 1,1 # 2562, 10242, 40962
Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
N, T, courant = 40, 10800., 1.2 # simulation length = N*T

print 'Omega, planetary PV', Omega, 2*Omega/gh0
sys.stdout.flush()

def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter
print 'Reading MPAS mesh ...'
sys.stdout.flush()
mesh = unst.MPAS_Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f)
print '...Done'
sys.stdout.flush()
lon, lat = mesh.lon_i, mesh.lat_i
x,y,z = np.cos(lat)*np.cos(lon), np.cos(lat)*np.sin(lon), np.sin(lat)

unst.setvar('g',g)
caldyn = unst.Caldyn_RSW(mesh)

c0 = math.sqrt(gh0) # phase speed of barotropic mode
dx = mesh.de.min()
dt = courant*dx/c0
nt = int(math.ceil(T/dt))
dt = T/nt
#scheme = time_step.RKn_simple(1,caldyn.bwd_fast_slow, dt)
scheme = time_step.RK4(caldyn.bwd_fast_slow, dt)
    
print dx, dt, dt*c0/dx
print T, nt

#mesh.plot_e(mesh.le_de) ; plt.title('le/de') ; plt.show()
#mesh.plot_i(mesh.Ai) ; plt.title('Ai') ; plt.show()

#Phi0 = gh0*(1+0.1*np.exp(-2000.*(y+1.)**2))
#flow = (Phi0,mesh.field_u())
#for i in range(N):
#    h,u=flow
#    plot_i(h) ; plt.title('h'); plt.show()
#    junk, fast, slow = caldyn.bwd_fast_slow(flow,0.)
#    print i
#    flow = scheme.advance(flow, nt)

u0 = Omega*radius/12. # cf Williamson (1991), p.13
gh1 = radius*Omega*u0+.5*u0**2
print 'Williamson (1991) test 2, u0=', u0
ulon = u0*np.cos(mesh.lat_e)
Phi0 = gh0 - gh1*(np.sin(mesh.lat_i)**2)
zeta0 = (2*u0/radius+2*Omega)*np.sin(mesh.lat_v)
Phi0v = gh0 - (radius*Omega*u0+.5*u0**2)*(np.sin(mesh.lat_v)**2)
q0 = zeta0/Phi0v
fu_perp = mesh.ucov2D(0.*ulon,2*Omega*np.sin(mesh.lat_e)*ulon)

flow = (Phi0,mesh.ucov2D(ulon,0.*ulon))
for i in range(N):
    h,u=flow
    mesh.plot_i(h-Phi0) ; plt.title('err(gh)'); 
    plt.savefig('fig_RSW_MPAS_W02/err_gh_%02d.png'%i)
    plt.close()
#    junk, fast, slow = caldyn.bwd_fast_slow(flow,0.)
#    plot_i(slow[0]) ; plt.title('dh/dt'); plt.show()
#    plot_e(slow[1]+fast[1]) ; plt.title('du/dt'); plt.show()
#    plt.figure(); plt.plot(fu_perp,slow[1],'.');
#    plt.xlabel('fu_perp'); plt.ylabel('u_slow'); plt.show()
#    plt.figure(); plt.plot(fu_perp,fast[1],'.');
#    plt.xlabel('fu_perp'); plt.ylabel('u_fast'); plt.show()
#    plot_e(fast[1]) ; plt.title('u_fast'); plt.show()
#    plot_e(slow[1]) ; plt.title('u_slow'); plt.show()
#    plt.figure(); plt.plot(q0,caldyn.qv,'.');
#    plt.xlabel('q0'); plt.ylabel('qv'); plt.show()
    print i, h.min(), h.max()
    flow = scheme.advance(flow, nt)