from dynamico import meshes
from dynamico import unstructured as unst
from dynamico import time_step
from dynamico import precision as prec
from dynamico import xios

import math as math
import numpy as np
import argparse

#------------------ Command-line arguments ------------------

parser = argparse.ArgumentParser()

parser.add_argument("--mpi_ni", type=int, default=1,
                    help="number of x processors")
parser.add_argument("--mpi_nj", type=int, default=1,
                    help="number of y processors")
parser.add_argument("-T", type=float, default=10.,
                    help="Length of time slice")
args = parser.parse_args()

#------------------------- Functions ----------------------

def globmin(x):
   xloc = x[mesh.primal_own_loc]
   locmin= xloc.min()
   return client.min(locmin)

def globmax(x):
   xloc = x[mesh.primal_own_loc]
   locmin= xloc.max()
   return client.max(locmin)

def minmax(name, x):
   gmin,gmax  = globmin(x),globmax(x)
   if(mpi_rank==0) : print('Min/max %s :'%name, gmin, gmax ) 

def laplace(mesh, p):
   left, right, ne = mesh.left, mesh.right, mesh.primal_ne
   gradp = mesh.field_u().flatten()
   for ij in range(mesh.left.size):
      p_left = p[left[ij]-1] # left,right start at 1 (CHECK)
      p_right = p[right[ij]-1] # left,right start at 1 (CHECK)
      gradp[ij]= mesh.le_de[ij]*(p_right-p_left) # convert to contravariant
   lap = mesh.field_mass().flatten()
   for ij in range(lap.size):
      lap_ij=0.
      for k in range(mesh.primal_deg[ij]):
         edge=mesh.primal_edge[ij,k]-1 # primal_edge starts at 1 (CHECK)
         lap_ij = lap_ij + mesh.primal_ne[ij,k]*gradp[edge]
      lap[ij]=lap_ij
   return lap/mesh.Ai

def curlgrad(mesh,p):
   left, right, ne = mesh.left, mesh.right, mesh.primal_ne
   gradp = mesh.field_u().flatten()
   for ij in range(mesh.left.size):
      p_left = p[left[ij]-1] # left,right start at 1 (CHECK)
      p_right = p[right[ij]-1] # left,right start at 1 (CHECK)
      gradp[ij]= p_right-p_left # covariant
   curl = mesh.field_z().flatten()
   for ij in range(curl.size):
      curl_ij=0.
      for k in range(mesh.dual_deg[ij]):
         edge=mesh.dual_edge[ij,k]-1 # dual_edge starts at 1 (CHECK)
         curl_ij = curl_ij + mesh.dual_ne[ij,k]*gradp[edge]
      curl[ij]=curl_ij
   return curl/mesh.Av

def curl(mesh, ucov):
   curl = mesh.field_z().flatten()
   for ij in range(curl.size):
      curl_ij=0.
      for k in range(mesh.dual_deg[ij]):
         edge=mesh.dual_edge[ij,k]-1 # dual_edge starts at 1 (CHECK)
         curl_ij = curl_ij + mesh.dual_ne[ij,k]*ucov[edge]
      curl[ij]=curl_ij
   return curl/mesh.Av

def test_laplace_seq():
   p=p.flatten()
   print p.shape
   curl = curlgrad(mesh,p)
   print 'curl(grad(p)) min max :', curl.min(), curl.max()       
   lap = laplace(mesh,p)
   p  = p+lap/vp
   print 'p+lap(p)/vp min max :', p.min(), p.max()

def test_laplace_parallel():
   transfer_primal, transfer_edge = mesh.com_primal.index, mesh.com_edges.index
   with xios.Context_Curvilinear(mesh,1, 24*3600) as context:
      # now XIOS knows about the mesh and we can write to disk
      context.update_calendar(0)
      for it in range(4):
         context.update_calendar(it+1)
         unst.ker.dynamico_update_halo(mesh.com_primal.index, 1, mesh.primal_num, p)
         lap = laplace(mesh,p)
          # compute eigenvalue and save residual (should be 0.)
         context.send_field_primal('p', p)
         p  = p+lap/vp
         context.send_field_primal('ps', p)
          # next iteration will start from lap(p)
         p=-lap/vp

def periodize(z,nz): return (z+2*nz)%nz
def index(x,y):      return periodize(x,nx)+nx*periodize(y,ny)
def indexu(x,y):     return 2*index(x,y)
def indexv(x,y):     return 2*index(x,y)+1

def test_RSW():
   x1,x2= xx-1., xx+1.
   if False:
      u0 = mesh.field_u()  # flow initally at rest
      h0 = 1+0.1*(np.exp(-2.*(x1*x1+yy*yy))+
               np.exp(-2.*(x2*x2+yy*yy)))
   else: # uniform flow
      ulon = 1.+0.*mesh.lat_e
      ulat = 0.*ulon
      u0 = mesh.ucov2D(ulon,ulat)
      h0 = mesh.field_mass()
      h0[:]=1.

   dt=0.1
   scheme = time_step.RK1(None, dt)
   step = unst.caldyn_step_TRSW(mesh,scheme,1)

   # check halo exchange
   step.u[:] = mesh.edges_E2 # global index of edges
   unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, step.u)
   for i in range(step.u.size):
      if step.u[i] != mesh.edges_E2[i]:
         print '[%d] minmax halo exchanges mismatch : ', i, step.u[i], mesh.edges_E2[i]

   step.mass[:]=h0
   step.theta_rhodz[:]=h0
   step.u[:]=u0

   def minmax(name, data): print '[%d] minmax %s :'%(mpi_rank,name), data.min(), data.max()

   print 'minmax dual_deg : ', mesh.dual_deg.min(), mesh.dual_deg.max()
   with xios.Context_Curvilinear(mesh,1, 24*3600) as context:
      for it in range(3):
         step.next()
         context.update_calendar(it+1)
         u, dh, fast, slow = step.u[:], step.drhodz[:], step.du_fast[:], step.du_slow[:]
         qv = step.qv[:]
         # better exchange halos before evaluating min/max
         unst.ker.dynamico_update_halo(mesh.com_primal.index, 1, mesh.primal_num, dh)
         unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, fast)
         unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, slow)
#         unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, qu)
#         unst.ker.dynamico_update_halo(mesh.com_dual.index, 1, mesh.dual_num, qv)
         minmax('dh',dh)
         minmax('du_fast',fast)
         minmax('du_slow',slow)
#         minmax('qu',qu)
         minmax('qv',qv)
         context.send_field_primal('p', dh )

         unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, mesh.edge_num, u)
         zeta = curl(mesh, u)
         minmax('zeta', zeta)
         for i in range(zeta.size):
            if zeta[i]<0. :
               if mpi_rank==0:
                  ij = mesh.vertices_V1[i] # global index of dual cell (starting at 0)
                  x,y = ij%nx, ij/nx
                  edges1 = [ indexv(x+1,y), indexu(x,y+1), indexv(x,y), indexu(x,y) ]
                  edges2 = [ mesh.edges_E2[k-1] for k in mesh.dual_edge[i,:] ]
                  if edges1!=edges2: print '[%d] minmax zeta '%mpi_rank, ij, x, y
                  print '[%d] minmax zeta '%mpi_rank, [ u[edge-1] for edge in mesh.dual_edge[i,:] ]
                  print '[%d] minmax zeta '%mpi_rank, mesh.dual_ne[i,:]

#--------------------------------- Main program -----------------------------

with xios.Client() as client: # setup XIOS which creates the DYNAMICO communicator
    comm = client.comm
    mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
    print '%d/%d starting'%(mpi_rank,mpi_size)

    #    nx,ny,Lx,Ly=200,30,4e7,6e6
    nx,ny,Lx,Ly      = 32,32,8.,8.
    llm, nqdyn, g, f = 1,1,1.,0.
    dx,dy=Lx/nx,Ly/ny

    cfl = .8
    dt = cfl/math.sqrt((nx/Lx)**2+(ny/Ly)**2)
    T=args.T
    N=int(T/dt)+1
    dt=T/N
    print N,dt,Lx/nx

    filename, radius = 'cart_%03d_%03d.nc'%(nx,ny), None
    unst.setvar('g',g)

    def coriolis(lon,lat):
       return f+0.*lon

    print 'Reading Cartesian mesh ...'
    meshfile = meshes.DYNAMICO_Format(filename)

    parallel=True
    laplace=False

    if parallel:
       pmesh = meshes.Unstructured_PMesh(comm,meshfile)
       pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
    #    pmesh.partition_metis()
       mesh  = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis)
    else:
       mesh = meshes.Cartesian_mesh(nx,ny,llm,nqdyn,Lx,Ly,f)

    xx,yy = mesh.lon_i, mesh.lat_i

    if laplace:
       p = np.cos(2*np.pi*xx/Lx)*np.cos(2.*np.pi*yy/Ly)

    # analytic eigenvalue : vp = -((2*np/pi/Lx)**2+(2*np.pi/Ly)**2)
       kx = 2.*np.sin(dx*np.pi/Lx)/dx
       ky = 2.*np.sin(dy*np.pi/Ly)/dy
       vp = kx**2+ky**2

       if parallel:
          test_laplace_parallel()
       else:
          test_laplace_seq()

    else:
       test_RSW()