from dgpy import *
from SRInputTsunami import *
from termcolor import colored
import time, math, os, sys
print(colored("*** Partition mesh with algorithm %d ***", "red")%(algo))
model = GModel()
model.load(filename + '.msh')
pOpt = meshPartitionOptions()
pOpt.setNumOfPartitions(int(sys.argv[1]))

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

if(os.path.exists(outputDir+"/"+filename+"_bath_smooth"+"/"+filename+"_bath_smooth.idx")):
  Msg.Info('.... Import Default Smoothed Bathymetry')
  bath.importIdx(outputDir+"/"+filename+"_bath_smooth"+"/"+filename+"_bath_smooth.idx")
else: 
  Msg.Fatal("No corresponding bathymetry file")
claw = dgConservationLawShallowWater2d()
solution = dgDofContainer(gr, claw.getNbFields())

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

claw.setBathymetry(bath.getFunction())
claw.setBathymetryGradient(bath.getFunctionGradient())
claw.setCoriolisFactor(FCoriolis)
claw.setIsSpherical(6371220.)
claw.setIsLinear(True)
#bottomDrag = functionC(tlib, "bottomDrag", 1, [function.getSolution(), bath.getFunction()])
#claw.setQuadraticDissipation(bottomDrag)
bottomDragLinear = functionC(tlib, "bottomDragLinear", 1, [function.getSolution(), bath.getFunction()])
claw.setLinearDissipation(bottomDragLinear)

rk=dgMultirateERK(gr,  claw,  RK_TYPE)
dt=rk.splitGroupsForMultirate(mL,  solution,  [solution,  bath])
solution.exportGroupIdMsh()
print 'SU=', rk.speedUp()
rk.printMultirateInfo()

rk.computeMultiConstrainedPartition(pOpt,  algo)
PartitionMesh(model, pOpt)
rk.printPartitionInfo(int(sys.argv[1]))
model.save(filename + '_sr_partitioned%d.msh'%(int(sys.argv[1])))
Msg.Exit(0)