#
# Reads ans manages the model grid
#
import numpy as np
from netCDF4 import Dataset
# Doc : https://spacetelescope.github.io/spherical_geometry/api/spherical_geometry.polygon.SphericalPolygon.html
from spherical_geometry import polygon
import RPPtools as RPP
#
import configparser
config=configparser.ConfigParser()
config.read("run.def")
EarthRadius=config.getfloat("OverAll", "EarthRadius")
#
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 getbox(ncdf, corners) :
    halo=0.1
    ii=[np.argmin(abs(ncdf.variables["nav_lon"][0,:]-(np.min(corners[0][:])-halo))),
        np.argmin(abs(ncdf.variables["nav_lon"][0,:]-(np.max(corners[0][:])+halo)))]
    jj=[np.argmin(abs(ncdf.variables["nav_lat"][:,0]-(np.min(corners[1][:])-halo))),
        np.argmin(abs(ncdf.variables["nav_lat"][:,0]-(np.max(corners[1][:])+halo)))]

    return min(ii), max(ii), min(jj), max(jj)
#
# Get the corners for a regular lat/lon grid
#
def corners(lon, lat) :
    jjm,iim = lon.shape
    cornerspoly = []
    cornersll = []
    centersll=[]
    radiusll = []
    index = []
    hdlon=np.mean(np.abs(np.diff(lon[0,:])))
    hdlat=np.mean(np.abs(np.diff(lat[:,0])))
    #
    for i in range(iim) :
        for j in range(jjm) :
            #
            boxll = RPP.boxit([lon[j,i], lat[j,i]], hdlon, hdlat, 2)
            lons = [p[0] for p in boxll]
            lats = [p[1] for p in boxll]
            #
            cornerspoly.append(polygon.SphericalPolygon.from_lonlat(lons, lats, center=[lon[j,i], lat[j,i]]))
            #
            centersll.append([lon[j,i], lat[j,i]])
            radiusll.append(RPP.maxradius([lon[j,i], lat[j,i]], lons, lats))
            #
            index.append([j,i])
            cornersll.append(boxll)

    return cornersll, cornerspoly, centersll, radiusll, index
#
def gather(x, index) :
    y=[]
    for ia in index :
        # print("XXXX shape of ia :", ia.shape, ia.dtype)
        # print("XXXX ia :", ia)
        y.append(list(x[ia[0,i],ia[1,i]] for i in range(ia.shape[1]) ))
    return y
#
def getattrcontaining(nc, varname, substr) :
    att=[]
    for s in nc.variables[varname].ncattrs() :
        if s.lower().find(substr) >= 0 :
            att.append(nc.variables[varname].getncattr(s))

    return att
#
class HydroGrid :
    def __init__(self, lolacorners) :
        #
        self.source=config.get("OverAll", "HydroFile")
        INFO("Opening in HydroGrid : "+self.source)
        self.ncfile=Dataset(self.source,'r')
        istr, iend, jstr, jend = getbox(self.ncfile, lolacorners)
        self.box=[istr, iend, jstr, jend]
        self.lon=np.copy(self.ncfile.variables["nav_lon"][jstr:jend,istr:iend])
        self.lat=np.copy(self.ncfile.variables["nav_lat"][jstr:jend,istr:iend])
        self.jjm,self.iim = self.lon.shape
        DEBUG("# Range Lon :"+str(np.min(self.lon))+" -- "+str(np.max(self.lon)))
        DEBUG("# Range Lat :"+str(np.min(self.lat))+" -- "+str(np.max(self.lat)))
        self.polyll, self.polylist, self.centers, self.radius, self.index = corners(self.lon, self.lat)

    def select(self, c, r) :
        indices=[]
        for i in range(len(self.centers)) :
            if RPP.loladist(c,self.centers[i]) <= r+self.radius[i] :
                indices.append(i)
        return indices

class HydroData :
    def __init__(self, nf, box, index) :
        istr, iend, jstr, jend = box[:]
        self.trip=gather(nf.variables["trip"][jstr:jend,istr:iend].astype(np.float32), index)
        self.tripdesc=nf.variables["trip"].long_name
        #
        self.basins=gather(nf.variables["basins"][jstr:jend,istr:iend], index)
        self.basinsdesc=nf.variables["basins"].long_name
        att = getattrcontaining(nf, "basins", "max")
        if len(att) > 0 :
            self.basinsmax=att[0]
        else :
            INFO("We need to scan full file to find maximum number of basins")
            self.basinsmax=np.max(np.where(nf.variables["basins"][:,:] <  1.e10))
        #
        self.topoind=gather(nf.variables["topoind"][jstr:jend,istr:iend].astype(np.float32), index)
        self.topoinddesc=nf.variables["topoind"].long_name
        att = getattrcontaining(nf, "topoind", "min")
        if len(att) > 0 :
            self.topoindmin=att[0]
        else :
            INFO("We need to scan full file to find minimum topoind over domain")
            self.topoindmin=np.min(np.where(nf.variables["topoind"][:,:] <  1.e15))
        #
        #
        self.topoindh=gather(nf.variables["topoind_h"][jstr:jend,istr:iend].astype(np.float32), index)
        self.topoindhdesc=nf.variables["topoind_h"].long_name
        att = getattrcontaining(nf, "topoind_h", "min")
        if len(att) > 0 :
            self.topoindhmin=att[0]
        else :
            INFO("We need to scan full file to find minimum topoind_h over domain")
            self.topoindhmin=np.min(np.where(nf.variables["topoind_h"][:,:] <  1.e15))
        #
        self.disto=gather(nf.variables["disto"][jstr:jend,istr:iend].astype(np.float32), index)
        self.distodesc=nf.variables["disto"].long_name
        #
        self.fac=gather(nf.variables["fac"][jstr:jend,istr:iend].astype(np.float32), index)
        self.facdesc=nf.variables["fac"].long_name
        #
        if "orog" in nf.variables.keys():
            self.orog = gather(nf.variables["orog"][jstr:jend,istr:iend].astype(np.float32), index)
            self.orogdesc=nf.variables["orog"].long_name
        else:
            self.orog = gather(np.zeros((jend-jstr,iend-istr)).astype(np.float32), index)
        #
        if "floodplains" in nf.variables.keys():
            self.floodplains = gather(nf.variables["floodplains"][jstr:jend,istr:iend].astype(np.float32), index)
            self.floodplainsdesc=nf.variables["floodplains"].long_name
        else:
            self.floodplains = gather(np.zeros((jend-jstr,iend-istr)).astype(np.float32), index)