from dgpy import *
from InputTsunami import *
from termcolor import colored
import time, math 
import os, sys

model = GModel()
model.load(filename + '.msh')

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)
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()

pOpt = meshPartitionOptions()
pOpt.setNumOfPartitions(int(sys.argv[2]))
if(int(sys.argv[3])==1):
  rk.computeClassicalContrainedPartition(pOpt,  algo)
if(int(sys.argv[3])==2):
  rk.computeSingleContrainedPartition(pOpt,  algo)
if(int(sys.argv[3])==3):
  rk.computeMultiConstrainedPartition(pOpt,  algo)
PartitionMesh(model, pOpt)
rk.printPartitionInfo(int(sys.argv[2]))
model.save(filename + '_%dm_%dp_%da.msh'%(int(sys.argv[1]), int(sys.argv[2]), int(sys.argv[3])))
print("%f %f %f %f %f"%(fact, rk.dtMin(), rk.dtMax(), rk.dtRef(), rk.speedUp()))
Msg.Exit(0)