from __future__ import print_function

print('Loading DYNAMICO modules ...')
from dynamico import unstructured as unst
from dynamico.precision import asnum
from dynamico.dev.meshes import MPAS_Format, Unstructured_Mesh
from dynamico.dev.numba import jit
from dynamico import time_step
print('...Done')

print('Loading modules ...')
import math as math
import matplotlib.pyplot as plt
import numpy as np
from numba import int32, float64
import numba
import time

print('...Done')

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

def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter
print('Reading MPAS mesh ...')
meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
mesh=Unstructured_Mesh(meshfile, llm, nqdyn, radius, f)
print('...Done')
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)

u0 = Omega*radius/12. # cf Williamson (1991), p.13
gh1 = radius*Omega*u0+.5*u0**2
print('u0=', u0)
ulon = u0*np.cos(mesh.lat_e)

dx = mesh.de.min()
dt = courant*dx/u0
nt = int(math.ceil(T/dt))
dt = T/nt
print(dx, dt, dt*u0/dx)
print(T, nt)

q0 = np.exp(-32.*(x+1.)**2)
q0, u0 = asnum([q0, mesh.ucov2D(ulon,0.*ulon)])

@jit
def upwind(mesh,u,q,flux):
    # compute upwind fluxes
    for edge in range(mesh.edge_num):
        left = mesh.left[edge]-1
        right = mesh.right[edge]-1
        ue = u[edge]*mesh.le_de[edge]
        if ue>0:
            flux[edge]=ue*q[left]
        else:
            flux[edge]=ue*q[right]

@jit
def advance(mesh, scheme, u, q, flux):
    scheme(mesh,u,q,flux)
    # advance by -divergence of flux
    for cell in range(mesh.primal_num):
        div=0.
        for iedge in range(mesh.primal_deg[cell]):
            edge = mesh.primal_edge[cell,iedge]-1
            div = div + flux[edge]*mesh.primal_ne[cell,iedge]
        q[cell] = q[cell] - dt*div/mesh.Ai[cell]
    
q = q0
flux = mesh.field_u()

mesh_data = mesh.get_data()

for i in range(N):    
    mesh.plot_i(q) ; plt.title('q'); plt.xlim([0.,360.])
    plt.savefig('fig_Godunov/q_%02d.png'%i)
    plt.close()
    print(i)
    start_time = time.time()
    for iter in range(100):
        advance(mesh_data, upwind, u0,q,flux)
    elapsed_time = time.time() - start_time
    print('elapsed time : %g ms/step'%(1000.*elapsed_time/iter))

print('Time spent in DYNAMICO (s) : ', unst.getvar('elapsed'))