#
#
#
import numpy as np
import os, sys
import pickle
from mpi4py import MPI
import gc
#
import matplotlib as mpl
mpl.use('Agg')
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt
#
# Gert the information from the configuration file.
#
import configparser
config=configparser.ConfigParser({'SaveWeights':'true', "DiagLon":"0.0, 0.0", "DiagLat":"0.0, 0.0"})
config.read("run.def")
saveweights=config.get("OverAll", "SaveWeights")
nbasmax=int(config.getfloat("OverAll", "nbasmax"))
OutGraphFile=config.get("OverAll","GraphFile")
lonint=np.array(config.get("OverAll", "DiagLon").split(","),dtype=float)
latint=np.array(config.get("OverAll", "DiagLat").split(","),dtype=float)
#
import ModelGrid as MG
import Partition as PA
import RPPtools as RPP
import HydroGrid as HG
import diagplots as DP
import Interface as IF
#
# Logging in MPI : https://groups.google.com/forum/#!topic/mpi4py/SaNzc8bdj6U
#
import getargs
log_master, log_world = getargs.getLogger()
INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug
#
# Read full grid and partition the domain.
#
gg = MG.GlobalGrid()
#
comm = MPI.COMM_WORLD
szhalo=1
nbcore=comm.Get_size()
rank=comm.Get_rank()
#
# Verify directories
#
wdir="Weights"
if rank == 0 :
    if not os.path.exists(wdir) :
        os.mkdir(wdir)
comm.Barrier()
#
# Region of grid to be treated
#
part = PA.partition(gg.ni, gg.nj, gg.land, comm, nbcore, szhalo, rank)
#
# Weight file to be created
#
wfile=os.path.basename(config.get("OverAll", "ModelGridFile").replace(".","_"))+"_i"+str(part.istart).zfill(4)+"di"+str(part.ni).zfill(4)+"j"+str(part.jstart).zfill(4)+"dj"+str(part.nj).zfill(4)
#
#
modelgrid=MG.ModelGrid(part.ihstart+gg.igstart,part.nih,part.jhstart+gg.jgstart,part.njh)
INFO("Longitude interval on proc "+str(rank)+" = "+str(modelgrid.box_land[0][0])+" : "+str(modelgrid.box_land[0][1]))
INFO("Latitude interval on proc "+str(rank)+" = "+str(modelgrid.box_land[1][0])+" : "+str(modelgrid.box_land[1][1]))


hydrogrid=HG.HydroGrid(modelgrid.box_land)

icell = 0
sub_index = []
sub_area = []
sub_lon = []
sub_lat = []
#
# Only compute the weights if the file does not exist.
#
if not os.path.exists(wdir+"/"+wfile) :
    INFO("Computing weights")
    for cell, area in zip(modelgrid.polylist, modelgrid.area) :

        overlap=np.zeros((hydrogrid.jjm,hydrogrid.iim), dtype=float)
        
        for hydrocell,index in zip(hydrogrid.polylist, hydrogrid.index) :
            if cell.intersects_poly(hydrocell) :
                overlap[index[0],index[1]] = cell.overlap(hydrocell)
                
        INFO(str(icell)+" Sum of overlap "+str(np.sum(overlap))+
             " Nb of Hydro grid overlaping :"+str(np.count_nonzero(overlap)))
        sub_index.append(np.where(overlap > 0.0))
        sub_area.append(np.extract(overlap > 0.0, overlap*area))
        sub_lon.append(np.extract(overlap > 0.0, hydrogrid.lon))
        sub_lat.append(np.extract(overlap > 0.0, hydrogrid.lat))
        icell = icell + 1
    #
    # Save the weights to a dedicated file for future use
    #
    if ( saveweights.lower() == 'true' ) :
        with open(wdir+"/"+wfile, 'wb') as fp:
            pickle.dump(sub_index, fp)
            pickle.dump(sub_area, fp)
            pickle.dump(sub_lon, fp)
            pickle.dump(sub_lat, fp)
        fp.close()
else :
    #
    # Extract the weights from the file if it exists
    #
    INFO("Reading weights")
    with open (wdir+"/"+wfile, 'rb') as fp:
        sub_index = pickle.load(fp)
        sub_area = pickle.load(fp)
        sub_lon = pickle.load(fp)
        sub_lat = pickle.load(fp)
    fp.close()
#
#
#
nbpt = len(sub_index)
sub_pts = np.array(list(len(sub_index[i][0]) for i in range(len(sub_index))))
nbvmax = max(sub_pts)
print("nbpt : ", nbpt)
print("nbvmax : ", nbvmax)
print("nbasmax : ", nbasmax)
#
# Extract hydo data from file
#
INFO("hydrodata")
hydrodata = HG.HydroData(hydrogrid.ncfile, hydrogrid.box, sub_index)

INFO("nitiatmgrid")
IF.initatmgrid(rank, nbcore, nbpt, modelgrid)


INFO("hoverlap")
hoverlap = IF.HydroOverlap(nbpt, nbvmax, sub_pts, sub_index, sub_area, sub_lon, sub_lat, modelgrid, hydrodata)

#
# Do some memory management and synchronize procs.
#
comm.Barrier()
del sub_index
del sub_area
del sub_lon
del sub_lat
gc.collect()

if rank ==0 :
    if lonint[0] != lonint[1] and latint[0] != latint[1] :
        DP.MAPPLOT("MapHydroGrid", lonint, latint, hoverlap, hoverlap.hierarchy_bx, modelgrid.polyll, title="Distance to ocean")

hsuper = IF.HydroSuper(nbvmax, hydrodata, hoverlap)
#
# Plot the hydrological supermesh
#
if rank == 0 :
    if lonint[0] != lonint[1] and latint[0] != latint[1] :
        DP.SUPERMESHPLOT("MapSuperGrid_Beforelinkup", lonint, latint, hoverlap, hsuper, modelgrid.polyll)

print("=================== LINKUP ====================")
hsuper.linkup(hydrodata)
if rank ==0 :
    if lonint[0] != lonint[1] and latint[0] != latint[1] :
        DP.SUPERMESHPLOT("MapSuperGrid_Afterlinkup", lonint, latint, hoverlap, hsuper, modelgrid.polyll)

#
# Do some memory management and synchronize procs.
#
comm.Barrier()
del hoverlap
gc.collect()

print("=================== Compute fetch ====================")
hsuper.fetch(part)
hsuper.dumpnetcdf(OutGraphFile.replace(".nc","_HydroSuper.nc"), gg, modelgrid, part)

hgraph = IF.HydroGraph(nbasmax, hsuper, part)

hgraph.dumpnetcdf(OutGraphFile, gg, modelgrid, part)

IF.closeinterface(comm)