import numpy as np
import sys
#
import getargs
log_master, log_world = getargs.getLogger(__name__)
INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug
#
def halfpartition(partin) :
    partout=[]
    if partin[0]["nbi"] > partin[0]["nbj"] :
        for dom in partin :
            h=int(np.trunc(dom["nbi"]/2))
            partout.append({"nbi":h,           "nbj":dom["nbj"],"istart":dom["istart"],  "jstart":dom["jstart"]})
            partout.append({"nbi":dom["nbi"]-h,"nbj":dom["nbj"],"istart":dom["istart"]+h,"jstart":dom["jstart"]})
    else :
        for dom in partin :
            h=int(np.trunc(dom["nbj"]/2))
            partout.append({"nbi":dom["nbi"],"nbj":h,           "istart":dom["istart"],"jstart":dom["jstart"]})
            partout.append({"nbi":dom["nbi"],"nbj":dom["nbj"]-h,"istart":dom["istart"],"jstart":dom["jstart"]+h})
    return partout
#
# Add number of land points to partitions
#
def addnbland(partin, land) :
    for dom in partin :
        dom['nbland']=int(np.sum(land[dom["jstart"]:dom["jstart"]+dom["nbj"],dom["istart"]:dom["istart"]+dom["nbi"]]))
#
# Divide in two the domain with the most number of land points
#
def halflargest(part) :
    nbland=np.array([dom['nbland'] for dom in part])
    imax = np.argmax(nbland)
    if part[imax]["nbi"] > part[imax]["nbj"] :
        h=int(np.trunc(part[imax]["nbi"]/2))
        part.append({"nbi":part[imax]["nbi"]-h,"nbj":part[imax]["nbj"],"istart":part[imax]["istart"]+h,"jstart":part[imax]["jstart"]})
        part[imax]["nbi"]=h
    else :
        h=int(np.trunc(part[imax]["nbj"]/2))
        part.append({"nbi":part[imax]["nbi"],"nbj":part[imax]["nbj"]-h,"istart":part[imax]["istart"],"jstart":part[imax]["jstart"]+h})
        part[imax]["nbj"]=h

    return
#
#
#
def adjustpart(partin, land, nbcore) :
    nbland=np.array([dom['nbland'] for dom in partin])
    nbptproc=np.array([dom['nbi']*dom['nbj'] for dom in partin])
    zeros=np.nonzero(nbland < 1)[0]
    #
    # Delete partitions without land points.
    #
    for i in range(len(zeros)) :
        partin.pop(zeros[i]-i)
    #
    # Divide domains with the largest number of land points
    #
    freecores = nbcore-len(partin)
    while nbcore-len(partin) > 0 :
        halflargest(partin)
        addnbland(partin, land)
        
    nbptproc=np.array([dom['nbi']*dom['nbj'] for dom in partin])
    INFO("Nb of land points per proc : Mean = "+str(np.mean(nbptproc))+" Min = "+str(np.min(nbptproc))+\
         " Max = "+str(np.max(nbptproc))+" Var :"+str((np.max(nbptproc)-np.min(nbptproc))/np.mean(nbptproc)*100)+"%")
    if np.min(nbptproc) < 3 :
        ERROR("The smallest domain in the partition is of a size less than 5. This will probably not work.")
        sys.exit()        
    return
#
# Partition domain in two steps : 1) halving, 2) adjusting by nb land points
#
def partitiondom(nig, njg, land, nbcore) :
    partout=[{"nbi":nig,"nbj":njg,"istart":0,"jstart":0, "nbland":int(np.sum(land))}]
    while len(partout) <= nbcore/2 :
        partout = halfpartition(partout)
    addnbland(partout, land)
    #
    adjustpart(partout, land, nbcore)
    #
    procmap=np.zeros((njg,nig), dtype=np.int8)
    procmap[:,:] = -1
    for i in range(len(partout)) :
        for ij in range(partout[i]["nbj"]) :
            for ii in range(partout[i]["nbi"]) :
                procmap[ij+partout[i]["jstart"],ii+partout[i]["istart"]] = i

    return partout, procmap
#
# Add halo to the partition
#
def addhalo(nig, njg, part, procmap, nbh) :
    nbcore=len(part)
    halosource=np.zeros((njg,nig,nbcore), dtype=np.int8)
    halosource[:,:,:] = -1
    coresend=np.zeros((njg,nig,nbcore), dtype=np.int8)
    coresend[:,:,:] = -1
    #
    for proc, dom in enumerate(part) :
        dom['ihstart']=max(dom["istart"]-nbh,0)
        dom['nbih']=min(dom["istart"]+dom['nbi']-1+nbh, nig-1)-dom['ihstart']+1
        dom['jhstart']=max(dom["jstart"]-nbh,0)
        dom['nbjh']=min(dom["jstart"]+dom['nbj']-1+nbh, njg-1)-dom['jhstart']+1
        for i in range(dom['nbih']) :
            #
            ic=dom['ihstart']+i
            #
            if dom['jhstart'] != dom["jstart"] :
                sproc = procmap[dom['jhstart'],ic]
                halosource[dom['jhstart'],ic,proc] = sproc
                coresend[dom['jhstart'],ic,proc] = proc
            #
            if dom['jhstart']+dom['nbjh']-1 != dom['jstart']+dom['nbj']-1 :
                sproc = procmap[dom['jhstart']+dom['nbjh']-1,ic]
                halosource[dom['jhstart']+dom['nbjh']-1,ic,proc] = sproc
                coresend[dom['jhstart']+dom['nbjh']-1,ic,proc] = proc
        for j in range(dom['nbjh']) :
            #
            jc=dom['jhstart']+j
            #
            if dom['ihstart'] != dom['istart'] :
                sproc = procmap[jc,dom['ihstart']]
                halosource[jc,dom['ihstart'],proc] = sproc 
                coresend[jc,dom['ihstart'],proc] = proc
            if dom['ihstart']+dom['nbih']-1 != dom['istart']+dom['nbi']-1 :
                sproc = procmap[jc,dom['ihstart']+dom['nbih']-1]
                halosource[jc,dom['ihstart']+dom['nbih']-1,proc] = sproc
                coresend[jc,dom['ihstart']+dom['nbih']-1,proc] = proc
    return halosource, coresend
#
# Get the list of the core points which need to be sent to another core.
#
def coresendlist(innersend_map, istart, ihstart, ni, jstart, jhstart, nj) :
    dcore=innersend_map[jstart:jstart+nj,istart:istart+ni,:]
    sendto = np.unique(dcore)[np.where(np.unique(dcore) >= 0)]
    innersend = []
    for ic in sendto :
        innersend.append(np.where(dcore[:,:,ic]==ic))
    #
    # Offset by halo
    #
    ihoff=istart-ihstart
    jhoff=jstart-jhstart
    for ic in range(len(sendto)) :
        innersend[ic][0][:] = innersend[ic][0][:]+jhoff
        innersend[ic][1][:] = innersend[ic][1][:]+ihoff
    return sendto, innersend
#
# All points belonging to the core of the domain
#
def listcoreland(istart, ihstart, ni, jstart, jhstart, nj, landind) :
    corelandlist=[]
    ihoff=istart-ihstart
    jhoff=jstart-jhstart
    for i in range(ni) :
        for j in range(nj) :
            if landind[jhoff+j,ihoff+i] >= 0 :
                corelandlist.append(landind[jhoff+j,ihoff+i])
    return corelandlist
#
# Get list of halo points which need to receive data
#
def haloreceivelist(halosource_map, rank) :
    receivefrom = np.unique(halosource_map[:,:,rank])[np.where(np.unique(halosource_map[:,:,rank]) >= 0)]
    halosource_g = []
    for ic in receivefrom :
        halosource_g.append(np.where(halosource_map[:,:,rank]==ic))
    return receivefrom, halosource_g
#
# Get 1D indices for the land points 
#
def landindexmap(istart, ni, jstart, nj, land) :
    gnj,gni=land.shape
    gindland=np.zeros((gnj,gni), dtype=np.int32)
    gindland[:,:]=-1
    n=0
    for i in range(gni) :
        for j in range(gnj) :
            if (land[j,i] > 0 ) :
                gindland[j,i] = n
                n += 1
    #
    indland=np.zeros((nj,ni), dtype=np.int32)
    indland[:,:]=-1
    local2global=[]
    n=0
    for i in range(ni) :
        for j in range(nj) :
            if (land[jstart+j,istart+i] > 0 ) :
                indland[j,i] = n
                local2global.append(gindland[jstart+j,istart+i])
                n += 1
    return n, indland, np.array(local2global, dtype=np.int32)
#
# Convert indices to local
#
def tolocal_index(x, istart, jstart) :
    xout = []
    for c in x :
        xout.append([c[0]-jstart,c[1]-istart])
    return xout
#
# Convert indices to localglobal
#
def toglobal_index(x, istart, jstart) :
    xout = []
    for c in x :
        xout.append([c[0]+jstart,c[1]+istart])
    return xout
#
#
#
def toland_index(x,landmap) :
    xout=[]
    for c in x :
        vl = landmap[c[0],c[1]]
        if np.count_nonzero(vl >= 0) > 0 :
            xout.append(vl[np.where(vl >= 0)])
        else :
            xout.append([])
    return xout
#
# 
#
class partition :
    def __init__ (self, nig, njg, land, mpicomm, nbcore, halosz, rank, wunit="None") :
        #
        part, procmap = partitiondom(nig, njg, land, nbcore)
        halosource_map, innersend_map = addhalo(nig, njg, part, procmap, halosz)
        #
        # Self varibales
        #
        # MPI info
        self.comm = mpicomm
        self.size = nbcore
        self.rank = rank
        # Global info
        self.nig = nig
        self.njg = njg
        self.allistart = []
        self.alljstart = []
        for i in range(self.size) :
            self.allistart.append(part[i]["istart"])
            self.alljstart.append(part[i]["jstart"])
        # Local info
        self.ni = part[rank]["nbi"]
        self.nj = part[rank]["nbj"]
        self.istart = part[rank]["istart"]
        self.jstart = part[rank]["jstart"]
        self.nih = part[rank]["nbih"]
        self.njh = part[rank]["nbjh"]
        self.ihstart = part[rank]["ihstart"]
        self.jhstart = part[rank]["jhstart"]
        #
        self.nbland, landindmap, self.l2glandind = landindexmap(self.ihstart, self.nih, self.jhstart, self.njh, land)        
        #
        if wunit != "None" :
            wunit.write("Offsets with halo (j-i) : "+str(self.jhstart)+"-"+str(self.ihstart)+'\n')
            wunit.write("Offsets without halo (j-i) : "+str(self.jstart)+"-"+str(self.istart)+'\n')
            wunit.write("Sizes without halo (j-i) : "+str(self.nj)+"-"+str(self.ni)+'\n')
            wunit.write("Shape of halosource_map"+str(halosource_map.shape)+"\n")
            for j in range(self.njh) :
                wunit.write(str(halosource_map[self.jhstart+j,self.ihstart:self.ihstart+self.nih,rank])+"\n")
            wunit.write("=============\n")
            wunit.write("Shape of innersend_map"+str(innersend_map.shape)+" from proc "+str(rank)+"\n")
            for r in range(self.size) :
                smtmp = innersend_map[self.jstart:self.jstart+self.nj,self.istart:self.istart+self.ni,r]
                if np.sum(smtmp) > -1.0*self.nj*self.ni :
                    wunit.write(str(r)+" ---------------- "+str(r)+"\n")
                    for j in range(self.nj) :
                        wunit.write(str(smtmp[j,:])+"\n")
            wunit.write("=============\n")
            wunit.write("Shape of landinmap :"+str(landindmap.shape)+"\n")
            wunit.write(str(landindmap)+"\n")
            wunit.write("=============\n")
        #
        self.receivefrom, self.halosource_g = haloreceivelist(halosource_map, rank)
        self.nbreceive = len(self.receivefrom)
        # To local indexing
        self.halosource = tolocal_index(self.halosource_g, self.ihstart, self.jhstart)
        self.landhalosrc = toland_index(self.halosource, landindmap)
        #
        if wunit != "None" :
            wunit.write("+++++ Halo receivefrom list : "+str(self.receivefrom)+"\n")
            for ic in range(self.nbreceive) :
                wunit.write("For "+str(self.receivefrom[ic])+" landhalosrc is :"+str(self.landhalosrc[ic])+"\n")
        #
        self.sendto,self.innersend = coresendlist(innersend_map, self.istart, self.ihstart, self.ni, self.jstart, self.jhstart, self.nj)
        self.nbsend = len(self.sendto)
        # To global indexing
        self.innersend_g = toglobal_index(self.innersend, self.ihstart, self.jhstart)
        self.landinnersend = toland_index(self.innersend, landindmap)
        if wunit != "None" :
            wunit.write("+++++ Halo core send list : "+str(self.receivefrom)+"\n")
            wunit.write("+++++ Core send list : "+str(self.sendto)+"\n")
            for ic in range(self.nbsend) :
                wunit.write("To "+str(self.sendto[ic])+" send innersend :"+str(self.innersend[ic])+"\n")
                wunit.write("To "+str(self.sendto[ic])+" send landinnersend :"+str(self.landinnersend[ic])+"\n")
        #
        self.landcorelist = listcoreland(self.istart, self.ihstart, self.ni, self.jstart, self.jhstart, self.nj, landindmap)
        #
        return
    #
    # Send halo to the other procs
    #
    def sendtohalo(self, x) :
        #
        for i in range(max(self.nbreceive,self.nbsend)) :
            if i < self.nbsend :
                if len(x.shape) == 2 :
                    self.comm.send(x[self.innersend[i][0],self.innersend[i][1]], dest=self.sendto[i])
                else :
                    ERROR("Unforessen rank of the variable to be sent to halo")
                    sys.exit()
            if i < self.nbreceive :
                if len(x.shape) == 2 :
                    x[self.halosource[i][0],self.halosource[i][1]] = self.comm.recv(source=self.receivefrom[i])
                else :
                    ERROR("Unforessen rank of the variable to be received in halo")
                    sys.exit()
        return
    #
    # For field gathered in land point send halo to the other procs
    #
    def landsendtohalo(self, x, order='C') :
        #
        for i in range(max(self.nbreceive,self.nbsend)) :
            if i < self.nbsend :
                if len(self.landinnersend[i]) > 0 :
                    if len(x.shape) == 1 :
                        self.comm.send(x[self.landinnersend[i]], dest=self.sendto[i])
                    elif len(x.shape) == 2 :
                        if order == 'C' :
                            self.comm.send(x[:,self.landinnersend[i]], dest=self.sendto[i])
                        elif order == 'F' :
                            self.comm.send(x[self.landinnersend[i],:], dest=self.sendto[i])
                        else :
                            ERROR("Unforessen order of the variable to be sent to halo of land points")
                            sys.exit()
                    else :
                        ERROR("Unforessen rank of the variable to be sent to halo of land points")
                        sys.exit()
            if i < self.nbreceive :
                if len(self.landhalosrc[i]) > 0 :
                    if len(x.shape) == 1 :
                        x[self.landhalosrc[i]] = self.comm.recv(source=self.receivefrom[i])
                    elif len(x.shape) == 2 :
                        if order == 'C' :
                            x[:,self.landhalosrc[i]] = self.comm.recv(source=self.receivefrom[i])
                        elif order == 'F' :
                            x[self.landhalosrc[i],:] = self.comm.recv(source=self.receivefrom[i])
                        else :
                            ERROR("Unforessen order of the variable to be sent to halo of land points")
                            sys.exit()
                    else :
                        ERROR("Unforessen rank of the variable to be received in halo of land points")
                        sys.exit()
        return
    #
    # Gather all fields partitioned in the 2D domain onto the root proc
    #
    def gather(self, x) :
        #
        indim = x.shape
        outdim = indim[:-2]+(self.njg,self.nig)
        if self.rank == 0 :
            xout = np.zeros(outdim, dtype=x.dtype)
            if len(outdim) == 2 :
                xout[:,:] = np.nan
            elif len(outdim) == 3 :
                xout[:,:,:] = np.nan
            elif len(outdim) == 4 :
                xout[:,:,:,:] = np.nan
            else :
                ERROR("Unforessen rank of field to be gathered")
                sys.exit()
                
        else :
            xout = None
        #
	# Offset for the core region
	#
        ihoff=self.istart-self.ihstart
        jhoff=self.jstart-self.jhstart
        #
        if len(outdim) == 2 :
            xtmp = self.comm.gather(x[jhoff:jhoff+self.nj,ihoff:ihoff+self.ni], root=0)
        elif len(outdim) == 3 :
            xtmp = self.comm.gather(x[:,jhoff:jhoff+self.nj,ihoff:ihoff+self.ni], root=0)
        elif len(outdim) == 4 :
            xtmp = self.comm.gather(x[:,:,jhoff:jhoff+self.nj,ihoff:ihoff+self.ni], root=0)
        else :
            ERROR("Unforessen rank of field to be gathered")
            sys.exit()
        #
        if self.rank == 0 :
            for i,z in enumerate(xtmp) :
                zdim=z.shape
                if len(outdim) == 2 :
                    xout[self.alljstart[i]:self.alljstart[i]+zdim[-2],self.allistart[i]:self.allistart[i]+zdim[-1]] = z[:,:]
                elif len(outdim) == 3 :
                    xout[:,self.alljstart[i]:self.alljstart[i]+zdim[-2],self.allistart[i]:self.allistart[i]+zdim[-1]] = z[:,:,:]
                elif len(outdim) == 4 :
                    xout[:,:,self.alljstart[i]:self.alljstart[i]+zdim[-2],self.allistart[i]:self.allistart[i]+zdim[-1]] = z[:,:,:,:]
                else :
                    ERROR("Unforessen rank of field to be gathered")
                    sys.exit()
        #
        return xout
    #
    # For land vectors gathered over land and local to the proc, we gather them on root proc.
    #
    def landscatgat(self, modelgrid, xl, order='C') :
        xf = modelgrid.landscatter(xl, order=order)
        xout = self.gather(xf)
        return xout
    #
    # Convert local index of land points to global index
    #
    def l2glandindex(self, x) :
        nl,nh = x.shape
        y = np.zeros(x.shape, dtype=x.dtype)
        if nl == self.nbland :
            for i in range(nl) :
                for j in range(nh) :
                    # Land indices are in FORTRAN !!
                    y[i,j] = self.l2glandind[x[i,j]-1]+1
        else :
            ERROR("The first dimension does not have the length of the number of land points")
            sys.exit()
        return y
    #
    # Set to zero all points in the core
    #
    def zerocore(self, x, order='C') :
        if len(x.shape) == 1 :
            x[self.landcorelist] = 0.0
        elif len(x.shape) == 2 :
            if order == 'C' :
                x[:,self.landcorelist] = 0.0
            elif order == 'F' :
                x[self.landcorelist,:] = 0.0
            else :
                ERROR("Unforessen order for variable to zero")
                sys.exit()
        
        return
    #
    # Function to sum over the core regions
    #
    def sumcore(self, x, order='C') :
        if len(x.shape) == 1 :
            y = np.nansum(x[self.landcorelist])
        elif len(x.shape) == 2 :
            if order == 'C' :
               y =  np.nansum(x[:,self.landcorelist], axis=1)
            elif order == 'F' :
               y = np.nansum(x[self.landcorelist,:], axis=0)
            else :
                ERROR("Unforessen order for variable to zero")
                sys.exit()
        return y