from dgpy import *
from InputTsunami import *
from mpi4py import MPI
from math import *
import time, os, sys

model = GModel()
model.load(filename + '_%dm_%dp_%da.msh'%(int(sys.argv[1]), Msg.GetCommSize(), int(sys.argv[2])))

os.putenv("BATHFILE", "etopo2.bin")
groups = dgGroupCollection(model, dimension, order)
xyz = groups.getFunctionCoordinates();
bath = dgDofContainer(groups, 1);

bathname=outputDir+"/"+filename+"_bath_smooth"+"/"+filename+"_bath_smooth.idx"
if(os.path.exists(bathname)):
  bath.importIdx(bathname)
else: 
  Msg.Fatal("No corresponding bathymetry file,   launch DiffuseBathymetryTsunami.py")

Finit = functionC(tlib, "initial_condition_okada_honshu",3, [xyz])
FCoriolis= functionC(tlib, "coriolis", 1, [xyz])

claw = dgConservationLawShallowWater2d()
solution = dgDofContainer(groups, claw.getNbFields())
solution.interpolate(Finit)
#solution.importIdx(outputDir+"/"+filename+"_sol"+"/"+filename+"_sol.idx")
claw.setBathymetry(bath.getFunction())
claw.setBathymetryGradient(bath.getFunctionGradient())
claw.setCoriolisFactor(FCoriolis)
claw.setIsSpherical(6371220.)
wall=claw.newBoundaryWall()
claw.addBoundaryCondition("Coast", wall)
bottomDrag = functionC(tlib, "bottomDrag", 1, [function.getSolution(), bath.getFunction()])
claw.setQuadraticDissipation(bottomDrag)
#claw.setIsLinear(True)
#bottomDragLinear = functionC(tlib, "bottomDragLinear", 1, [function.getSolution(), bath.getFunction()])
#claw.setLinearDissipation(bottomDragLinear)

rk=dgMultirateERK(groups, claw, RK_TYPE)
dt=rk.splitGroupsForMultirate(int(sys.argv[1]), solution, [solution, bath])
#solution.exportGroupIdMsh()

if(Msg.GetCommRank() == 0):
  print 'SU=', rk.speedUp()
rk.printMultirateInfo(0)

t=Ti
nbSteps = int(ceil((Tf-Ti)*3600/dt)) 
export = 967

exporterSol = dgIdxExporter(solution, outputDir+"/"+filename+"_sol")
Msg.Barrier()
rk.initialize()
Msg.Barrier()

timer=dgTimer.root()
timerFull=dgTimer.root()
timer.reset()
timerFull.reset()
Msg.Barrier()
timer.start()
timerFull.start()
for i in range (1,nbSteps+1) :
  rk.iterate(solution, dt, t)
  t = t + dt
  if(i%1==0):
    timer.stop()
    timerFull.stop()
    Msg.Barrier()
    maxTime = MPI.COMM_WORLD.allreduce(timer.get(),   op=MPI.MAX)
    minTime = MPI.COMM_WORLD.allreduce(timer.get(),   op=MPI.MIN)
    maxTimeFull = MPI.COMM_WORLD.allreduce(timerFull.get(),   op=MPI.MAX)
    minTimeFull = MPI.COMM_WORLD.allreduce(timerFull.get(),   op=MPI.MIN)
    norm=solution.norm()
    if (Msg.GetCommRank()==0):
      print "%4d %3.8f %3.8f %.8f %.8f %.8f %.8f"%(i, t/3600, norm, minTime,  maxTime, minTimeFull, maxTimeFull)
    Msg.Barrier()
    timer.reset()
    Msg.Barrier()
    timer.start()
    timerFull.start()
  if(i%export==0):
    timer.stop()
    timerFull.stop()
    exporterSol.exportIdx()
    timer.start()
    timerFull.start()

Msg.Exit(0)