"""
    Simulations on the GBR with multirate explicit Runge-Kutta time-stepping.
    All the parameters can be found and modified in the file "InpuGBR.py"
"""

"""
    *** How to launch the testcase ***

    Simulation on 1 processor  : rundgpy MultirateGBR.py
    Simulation on n processors : mpirun -np n rundgpy MultirateGBR.py   ( !!! You need MPI !!! )

    **********************************
"""

# Import for python
from dgpy import *
from math import *
from ProjectMesh import *
from InputGBR import *
from termcolor import colored
import time, os, sys

model = GModel()
model.load(filename + '_partitioned_tan%d.msh'%(Msg.GetCommSize()))
#loadAndPartitionMesh(model,  filename,  mL, algo)

# Coordinate system
XYZ = function.getCoordinates();
lonLatDegrees = functionC(glib,"lonLatDegrees",3,[XYZ])
lonLat = functionC(glib,"lonLat",3,[XYZ])

latLonDegrees = functionC(glib,"latLonDegrees",3,[XYZ])
latLon = functionC(glib,"latLon",3,[XYZ])

groups = dgGroupCollection(model, dimension, order)

if(not os.path.exists(bathname)):
  Msg.Fatal('No valid bathymetry file :' + bathname + ' found,  please launch "rundgpy DiffuseBathymetryTsunami.py"')

bath = dgDofContainer(groups, 1);
bath.importIdx(bathname)

# ************************************************

t = Ti
timeFunction = functionConstant(t)

# ****************** Forcings ******************

# .................. Tides    ..................

# (a) Sinus tide
def tide (f):
  for i in range(0,f.size1()):
    f.set(i, 0, sin((t-Ti)*2*pi/3600/12))
    f.set(i, 1, 0)
    f.set(i, 2, 0)

# (b) Tpxo tide 
tideEta = slimFunctionTpxo("h_tpxo7.2.nc","ha","hp","lon_z","lat_z")
tideEta.setCoordAndTimeFunctions(lonLatDegrees,timeFunction)
tideU = slimFunctionTpxo("u_tpxo7.2.nc","Ua","up","lon_u","lat_u")
tideU.setCoordAndTimeFunctions(lonLatDegrees,timeFunction)
tideV = slimFunctionTpxo("u_tpxo7.2.nc","Va","vp","lon_v","lat_v")
tideV.setCoordAndTimeFunctions(lonLatDegrees,timeFunction)
tideUV = functionC(glib,"latLonVector",3,[latLon,tideU,tideV])

# ..............................................

# .................. Wind     ..................

windInterpolator = slimStructDataInterpolatorMultilinear()
windUcontainer = slimStructDataContainerNetcdf("uwind-2000-2011.nc", "uwnd", "lat", "lon")
windVcontainer = slimStructDataContainerNetcdf("vwind-2000-2011.nc", "vwnd", "lat", "lon")
windU = slimFunctionStructData(windUcontainer, windInterpolator, latLonDegrees)
windV = slimFunctionStructData(windVcontainer, windInterpolator, latLonDegrees)
windU.setTimeFunction(timeFunction)
windV.setTimeFunction(timeFunction)
windUV = functionC(glib,"latLonVector",3,[latLon,windU,windV])

# ..............................................

# **********************************************


# ****************** Cons Law ******************

claw = dgConservationLawShallowWater2d()
solution = dgDofContainer(groups, claw.getNbFields())

ld = functionConstant(0)
Cd = functionC(glib, "bottomDrag", 1, [function.getSolution(), bath.getFunction()])
if(gbrDiff == 0):
  Di = functionConstant(diff)
elif(gbrDiff == 1):
  Di = functionC(glib, "smagorinsky", 1, [function.getSolutionGradient()])
Coriolis = functionC(glib, "coriolis", 1, [XYZ])
if(gbrWind == 0):
  source = functionC(glib, "wind", 2, [XYZ])
elif(gbrWind == 1):
  source = functionC(glib, "windStress", 2, [windUV, function.getSolution(), bath.getFunction()])


claw.setCoriolisFactor(Coriolis)
claw.setQuadraticDissipation(Cd)
claw.setLinearDissipation(ld)
claw.setDiffusivity(Di)
claw.setSource(source)
claw.setBathymetry(bath.getFunction())
claw.setBathymetryGradient(bath.getFunctionGradient())

Os = functionPython(3, tide,[])
tideuv = functionC(glib,"transport2velocity",3,[tideUV,bath.getFunction(),function.getSolution()]);
tpxoTide = functionC(glib,"merge",3,[tideEta,tideuv]);

if(gbrTide == 0):
  claw.addBoundaryCondition('Open Sea North',claw.newOutsideValueBoundary("Surface", Os))
  claw.addBoundaryCondition('Open Sea Central',claw.newOutsideValueBoundary("Surface", Os))
  claw.addBoundaryCondition('Open Sea South',claw.newOutsideValueBoundary("Surface", Os))
elif(gbrTide == 1):
  claw.addBoundaryCondition('Open Sea North',claw.newOutsideValueBoundary("Surface", tpxoTide))
  claw.addBoundaryCondition('Open Sea Central',claw.newOutsideValueBoundary("Surface", tpxoTide))
  claw.addBoundaryCondition('Open Sea South',claw.newOutsideValueBoundary("Surface", tpxoTide))

claw.addBoundaryCondition('Coast',claw.newBoundaryWall())
claw.addBoundaryCondition('Islands',claw.newBoundaryWall())

# **********************************************


# ****************** Time Int ******************
rk = dgMultirateERK(groups, claw, RK_TYPE)
dt = rk.splitGroupsForMultirate(mL, solution, [solution, bath])
rk.printMultirateInfo(0)

# **********************************************

currentFunction = functionC(glib,"current",3,[function.getSolution()])
#tideCurrentFunction = functionC(glib,"current",3,[tpxoTide])

nbSteps = int(ceil((Tf-Ti)/dt))
nbExport = 0

if(Msg.GetCommRank() == 0):
  if(os.path.exists(outputDir+'/mr-gbr')):
    try : os.system("rm -r "+outputDir+'/mr-gbr');
    except: 0;
exporterSol=dgIdxExporter(solution, outputDir+'/mr-gbr')

norm=solution.norm()

mass=functionPython(1,  massIntegrator,  [functionSolution.get(), bath.getFunction()])
integrator = dgFunctionIntegrator(groups,  mass,  solution)
totalMass=fullMatrixDouble(1, 1)
totalMassRef=fullMatrixDouble(1, 1)
integrator.compute(totalMassRef, "")

if (Msg.GetCommRank() == 0):
  print ''
  print(colored('******************** Initializing ********************', "blue"))
  print ''
  print '     - Number of processors:', Msg.GetCommSize()
  print '     - Multirate Runge Kutta Scheme:', RK_TYPE
  print '     - Initial Time (Ti):',  printDate(Ti)
  print '     - Final Time (Tf):',  printDate(Tf)
  print '     - Reference Time Step:', printTime(dt)
  print '     - Number of Iterartions:', nbSteps
  print '     - Theoretical Speedup=', rk.speedUp()
  print '     - Norm of Solution at Ti:',  norm
  print '     - Mass of Solution at Ti:',  totalMassRef(0, 0)
  print ''
  print(colored('******************************************************', "blue"))
  print ''


# ****************** Iterate  ******************


startcpu = time.clock()
for i in range(0, nbSteps):
  if(i == nbSteps - 1):
    dt=Tf-t
  norm = rk.iterate (solution, dt, t)
  t = t +dt
  timeFunction.set(t)
  if ( i%iter == 0 ):
     norm=solution.norm()
     integrator.compute(totalMass, "")
     relMass=totalMass(0, 0)
     if(Msg.GetCommRank() == 0):
       print(colored('|ITER| %d',  "red")%(i))
       print(colored('------------------------------------------------------', "red"))
       print'|TIME|', printDate(t)
       print'|NORM|',norm
       print'|MASS|',relMass 
       print'|CPUT|',printTime(time.clock() - startcpu)
       print(colored('------------------------------------------------------', "red"))
       print''
  if ( i%export  == 0 ):
    #exporterSol.exportIdx(i, t)
    solution.exportMsh(outputScratch+"mr_gbr-%06d"% (i),t, nbExport)
    solution.exportFunctionMsh(currentFunction, outputScratch+'mr_gbr_current-%06d'%(i), t, nbExport, '')
    nbExport = nbExport + 1

endcpu=time.clock()
norm=solution.norm()
integrator.compute(totalMass, "")
relMass=totalMass(0, 0)
Msg.Barrier()
if (Msg.GetCommRank() == 0):
  print ''
  print(colored('********************     End      ********************', "blue"))
  print ''
  print '     - Final Time (Tf):',  printDate(t)
  print '     - Norm of Solution at Tf:',  norm
  print '     - Mass of Solution at Tf:', relMass
  print '     - Total CPU Time:', endcpu-startcpu
  print ''
  print(colored('******************************************************', "blue"))
  print ''
Msg.Exit(0)