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

print ''
print(colored('********** Partition mesh with algorithm %d **********', "red")%(algo))
print ''
m2 = GModel()
m2.load(filename+'_tan' + '.msh')
pOpt = meshPartitionOptions()
pOpt.setNumOfPartitions(int(sys.argv[1]))

groups = dgGroupCollection(m2, dimension, order, True)
xyz = groups.getFunctionCoordinates();
lonLatDegrees = functionC(glib,"lonLatDegrees",3,[xyz])
lonLat = functionC(glib,"lonLat",3,[xyz])
bath = dgDofContainer(groups, 1);
if(not os.path.exists(bathname)):
  Msg.Fatal('No valid bathymetry file :' + bathname + ' found,  please launch "rundgpy DiffuseBathymetryTsunami.py"')

bath.importIdx(bathname)
# Define Conservation Law
claw = dgConservationLawShallowWater2d()
solution = dgDofContainer(groups, claw.getNbFields())

ld = functionConstant(0)
Cd = functionC(glib, "bottomDrag", 1, [solution.getFunction(), bath.getFunction()])
Di = functionConstant(diff)
Coriolis = functionC(glib, "coriolis", 1, [xyz])

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


rk = dgMultirateERK(groups, claw, RK_TYPE)
dt = rk.splitGroupsForMultirate(mL, solution, [solution, bath])
rk.printMultirateInfo()

rk.computeMultiConstrainedPartition(pOpt, algo)
PartitionMesh(m2, pOpt)
rk.updatePartitionInfo(int(sys.argv[1]))
rk.printPartitionInfo()

m2.save(filename + '_partitioned_tan%d.msh'%(int(sys.argv[1])))
print ''
print(colored('**********       Partition Done            **********', "red"))
print("%d %f %f %f %f %f"%(int(sys.argv[1]),rk.bestImbalance(),rk.globalImbalance(),rk.speedUp(),rk.speedUp()/rk.bestImbalance(),rk.speedUp()/rk.globalImbalance()))
print ''