#
# Script to generate partitioned mesh for use with MultirateGBR_Stereo.py
#   - run using the command: rundgpy PartitionGBR_StandaloneStereo.py <NB_PROCESSORS>
#

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

#-------------
#filename = "gbr300K_sgbr"
#filename = "gbr10K"
filename = sys.argv[1]
#-------------
# Multirate and Partition Parameters
RK_TYPE = ERK_22_C
mL = 1000
algo = 5
# Mesh Parameters
order = 1
dimension = 2
"""
-gbrDiff= 0: Constant Diffusion diff || 1: Parametrization of Smagorinsky 
"""
gbrDiff=1
#-------------

try : os.mkdir('./Meshes');
except: 0;
if(not os.path.exists("./Meshes/"+filename+".msh")):
  try : os.system("make "+filename+".msh");
  except: 0;

print"Generating partitioned mesh from CARTESIAN mesh:", filename

nbProcessors = int(sys.argv[2])

print''
print'Number of processors partitioning for:', nbProcessors
print''

#Compile libraries
if (Msg.GetCommRank() == 0):
  functionC.buildLibraryFromFile ("./Libraries/GBR.cc", "./Libraries/lib_gbr.so")
glib = "./Libraries/lib_gbr.so"

#First project mesh in planar space
if(not os.path.exists("./Meshes/"+filename+"_tan.msh")):
  print'Projecting mesh in planar space ...'
  m = GModel()
  m.load("./Meshes/"+filename+".msh")
  projectMeshPlane(m)
  m.save("./Meshes/"+filename+"_tan.msh")

#Set up model
print(colored('********** Partition mesh with algorithm %d **********', "red")%(algo))
print ''
m2 = GModel()
m2.load('./Meshes/'+filename+'_tan.msh')
pOpt = meshPartitionOptions()
pOpt.setNumOfPartitions(nbProcessors)

groups = dgGroupCollection(m2, dimension, order, True)
xyz = groups.getFunctionCoordinates()

#Read bathymetry
bathDC = dgDofContainer(groups, 1);
if(not os.path.exists("./Bath/"+filename+"_bath_smooth/"+filename+"_bath_smooth.idx")):
  Msg.Fatal('No valid bathymetry file found,  please launch "rundgpy DiffuseBathGBR.py"')
Msg.Info("Smoothed bathymetry already exists.")
bathDC.importIdx("./Bath/"+filename+"_bath_smooth/"+filename+"_bath_smooth.idx")
Msg.Info("Smoothed bathymetry read.")

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

ld = functionConstant(0)
Cd = functionC(glib, "bottomDrag_old", 1, [solution.getFunction(), bathDC.getFunction()])
#Cd = functionC(glib, "bottomDrag", 1, [ solution.getFunction(), bathDC.getFunction(), reefsDC.getFunction() ]) #***
if(gbrDiff == 0):
  Di = functionConstant(diff)
elif(gbrDiff == 1):
  Di = functionC(glib, "smagorinsky", 1, [solution.getFunctionGradient()])
Coriolis = functionC(glib, "coriolisConst", 1, [xyz])

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

rk = dgMultirateERK(groups, claw, RK_TYPE)
Msg.Info("Splitting groups for multirate ...")
dt = rk.splitGroupsForMultirate(mL, solution, [solution, bathDC])
rk.printMultirateInfo()
solution.exportGroupIdMsh()

Msg.Info("Computing optimal mesh partitions ...")
rk.computeMultiConstrainedPartition(pOpt, algo)
PartitionMesh(m2, pOpt)
rk.printPartitionInfo()

m2.save('./Meshes/'+filename + '_tan_partitioned_'+str(nbProcessors)+'.msh')
print ''
print(colored('**********       Partition Done            **********', "red"))
print ''