import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, cm
import numpy as np

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

inc=np.zeros((8,2), dtype=np.int8)
inc[0,0] = 1; inc[0,1] =0
inc[1,0] = 1; inc[1,1] =1
inc[2,0] = 0; inc[2,1] =1
inc[3,0] = -1; inc[3,1] =1
inc[4,0] = -1; inc[4,1] =0
inc[5,0] = -1; inc[5,1] =-1
inc[6,0] = 0; inc[6,1] =-1
inc[7,0] = 1; inc[7,1] =-1

epsilon=0.00001
undef_int=999999999
#
def buildfullgrid(lon, lat, lonint, latint) :
    nbpts, nlon, nlat = lon.shape
    # LON order
    for i in range(nlat) :
        if np.count_nonzero(~np.isnan(lon[0,:,i])) > 0 :
            dlon = np.sign(np.gradient(lon[0,:,i]))[~np.isnan(np.gradient(lon[0,:,i]))][0]
    # LAT order 
    for i in range(nlon) :
        if np.count_nonzero(~np.isnan(lat[0,i,:])) > 0 :
            dlat = np.sign(np.gradient(lat[0,i,:]))[~np.isnan(np.gradient(lat[0,i,:]))][0]
    #
    xlon=np.unique(np.hstack(lon))[~np.isnan(np.unique(np.hstack(lon)))]
    xlat=np.unique(np.hstack(lat))[~np.isnan(np.unique(np.hstack(lat)))]
    # Zomm in
    xlon=xlon[np.argmin(np.abs(xlon-min(lonint))):np.argmin(np.abs(xlon-max(lonint)))]
    xlat=xlat[np.argmin(np.abs(xlat-min(latint))):np.argmin(np.abs(xlat-max(latint)))]
    #
    if np.abs(np.min(np.gradient(xlon))-np.min(np.gradient(xlon))) < epsilon :
        res_lon = np.mean(np.gradient(xlon))
    else :
        ERROR("Could not determine the longitude resolution of the hydrological grid")

    if np.abs(np.min(np.gradient(xlat))-np.min(np.gradient(xlat))) < epsilon :
        res_lat = np.mean(np.gradient(xlat))
    else :
        ERROR("Could not determine the latitude resolution of the hydrological grid")
        
    if np.sign(np.gradient(xlon))[0] != dlon :
        lonv=np.concatenate(([np.max(xlon)+res_lon],xlon[::-1],[np.min(xlon)-res_lon]))
    else :
        lonv=np.concatenate(([np.min(xlon)-res_lon],xlon[:],[np.max(xlon)+res_lon]))
    if np.sign(np.gradient(xlat))[0] != dlat :
        latv=np.concatenate(([np.max(xlat)+res_lat],xlat[::-1],[np.min(xlat)-res_lat]))
    else :
        latv=np.concatenate(([np.min(xlat)-res_lat],xlat[:],[np.max(xlat)+res_lat]))
    #
    # return the full grid where longitude and latitude are 2D arrays
    #
    return np.meshgrid(lonv, latv, indexing='ij')
#
#
#
def enlargegrid(lon, lat, trip, data, lon_full, lat_full) :
    #
    nbpts, nlon, nlat = lon.shape
    nlonfull,nlatfull= lon_full.shape
    fulldata=np.zeros((nlonfull,nlatfull), dtype=np.float32)
    fulldata[:,:]=np.nan
    fulltrip=np.zeros((nlonfull,nlatfull), dtype=np.float32)
    fulltrip[:,:]=undef_int
    for ib in range(nbpts) :
        for i in range(nlon) :
            for j in range(nlat) :
                if trip[ib,i,j] < 1000 :
                    dist=np.sqrt((lon_full-lon[ib,i,j])**2 + (lat_full-lat[ib,i,j])**2)
                    il,jl=np.unravel_index(np.argmin(dist, axis=None), dist.shape)
                    fulltrip[il,jl]=trip[ib,i,j]
                    fulldata[il,jl]=data[ib,i,j]
    return fulltrip, fulldata
#
#
#
def enlargebasins(lon, lat, datasz, data, lon_full, lat_full) :
    nbpts, nlon, nlat = lon.shape
    nlonfull, nlatfull = lon_full.shape
    nbpts, nbvmax = datasz.shape
    fullbasins=np.zeros((nlonfull,nlatfull), dtype=np.float32)
    fullbasins[:,:]=np.nan
    for ib in range(nbpts) :
        for iz in range(nbvmax) :
            if datasz[ib,iz] > 0 :
                for isz in range(datasz[ib,iz]) :
                    # Fortran to C indexing
                    i=data[ib,iz,isz,0]-1
                    j=data[ib,iz,isz,1]-1
                    dist=np.sqrt((lon_full-lon[ib,i,j])**2 + (lat_full-lat[ib,i,j])**2)
                    il,jl=np.unravel_index(np.argmin(dist, axis=None), dist.shape)
                    fullbasins[il,jl]=iz
    return fullbasins
#
#
#
def closestpoint(lonout, latout, nbptb, basin_pts, lonbx, latbx) :
    lonin = 3.5
    latin = 39.0
    dist = undef_int
    for ib in range(nbptb) :
        # Fortran to C indexing
        ix=basin_pts[ib,0]-1
        iy=basin_pts[ib,1]-1
        if not np.isnan(lonbx[ix,iy]) :
            distnew = np.sqrt((lonout-lonbx[ix,iy])**2+(latout-latbx[ix,iy])**2)
            if distnew < dist :
                lonin = lonbx[ix,iy]
                latin = latbx[ix,iy]
                dist = distnew
    return lonin, latin, dist
#
#################################################################################
#            
def MAPPLOT(filename, loninterval, latinterval, hoverlap, data, atmpolys,
            label=' ', title=' ', lat_int = 5.0, lon_int=10, lat_min = 0, lon_min = 0, alldir=0):
    #
    # Test if we have to do a plot here
    #
    lon = hoverlap.lon_bx
    lat = hoverlap.lat_bx
    trip = hoverlap.trip_bx
    #
    if np.nanmin(lon) <= max(loninterval) and np.nanmax(lon) > min(loninterval) and \
       np.nanmin(lat) <= max(latinterval) and np.nanmax(lat) > min(latinterval) :
        #
        lon_full, lat_full = buildfullgrid(lon, lat, loninterval, latinterval)
        nblon, nblat = lon_full.shape
        res_lon=lon_full[1,0]-lon_full[0,0]
        res_lat=lat_full[0,1]-lat_full[0,0]
        #
        trip_full, data_full = enlargegrid(lon, lat, trip, data, lon_full, lat_full)
        #
        fig = plt.figure(1)
        ax = fig.add_subplot(111)
        ax.set_title(title)
        ax.set_xlabel('Longitude',labelpad=20)
        ax.set_ylabel('Latitude',labelpad=20)
        ax.axis([0, 10, 0, 10])
        #
        m = Basemap(projection='cyl', resolution='l', \
                    llcrnrlon=np.min(lon_full)-res_lon, llcrnrlat=np.min(lat_full)-0.5*res_lat, \
                    urcrnrlon=np.max(lon_full)+0.5*res_lon, urcrnrlat=np.max(lat_full)+res_lat)
        
        m.drawcoastlines(linewidth=0.5)
        m.drawcountries(linewidth=0.5)
        m.drawparallels(np.arange(-180.,180.,0.1),labels=[1,0,0,0]) # draw parallels
        m.drawmeridians(np.arange(-80.,80.,0.1),labels=[0,0,0,1]) # draw meridians
        #
        cmap = 'RdBu'
        cmap = 'jet_r'
        cmap = 'Spectral'
        #
        x, y = m(lon_full-0.5*res_lon, lat_full+0.5*res_lat)
        z = np.ma.array(data_full, mask=np.isnan(data_full)) 
        datamap = m.pcolor(x, y, z, zorder=0)
        #
        cNorm = matplotlib.colors.Normalize(vmin=np.nanmin(data_full), vmax=np.nanmax(data_full))
        datamap.set_norm(cNorm)
        plt.colorbar(datamap, cmap=cmap, norm=cNorm, orientation="vertical")
        #
        # Add polygone of atmospheric grid
        #
        for poly in atmpolys :
            xl,yl=m(np.array(poly)[:,1], np.array(poly)[:,0])
            plt.plot(xl,yl, color="m")
        #
        #
        #
        for i in range(1,nblon-1):
            for j in range(1,nblat-1):
                x0,y0 = m(lon_full[i,j], lat_full[i,j])
                t = int(trip_full[i,j]-1)
                if ( t >=0 and t < 10 ) :
                    x1,y1 = m(lon_full[i+inc[t,0],j+inc[t,1]], lat_full[i+inc[t,0],j+inc[t,1]])
                    plt.arrow(x0, y0, (x1-x0)*0.75, (y1-y0)*0.75, width=0.002, head_width=0.004, head_length=0.004, fc='w',ec=None)
                elif(t >= 100 and t < 110) :
                    iinc=inc[t-100,0]
                    jinc=inc[t-100,1]
                    x1,y1 = m(lon_full[i+iinc,j+jinc], lat_full[i+iinc,j+jinc])
                    plt.arrow(x0, y0, (x1-x0)*0.75, (y1-y0)*0.75, width=0.002, head_width=0.004, head_length=0.004, fc='r',ec=None)
                elif (t > 90 and t < 100):
                    # ocean point or inland point
                    plt.scatter(x0, y0, marker='o', facecolor='w', edgecolor='k', s=4)
                else :
                    plt.scatter(x0, y0, marker='x', facecolor='b', edgecolor='k', s=4)
        #
        plt.savefig(filename+'.png', dpi=450)
        plt.close()
#
#
#
def SUPERMESHPLOT(filename, loninterval, latinterval, hoverlap, hsuper, atmpolys, title=' ') :
    #
    # Test if we have to do a plot here
    #
    lon = hoverlap.lon_bx
    lat = hoverlap.lat_bx
    basin_sz = hsuper.basin_sz
    basin_pts = hsuper.basin_pts
    outlet = hsuper.basin_outcoor
    #
    if np.nanmin(lon) <= max(loninterval) and np.nanmax(lon) > min(loninterval) and \
       np.nanmin(lat) <= max(latinterval) and np.nanmax(lat) > min(latinterval) :
        #
        lon_full, lat_full = buildfullgrid(lon, lat, loninterval, latinterval)
        nblon, nblat = lon_full.shape
        res_lon=lon_full[1,0]-lon_full[0,0]
        res_lat=lat_full[0,1]-lat_full[0,0]
        #
        basins_full = enlargebasins(lon, lat, basin_sz, basin_pts, lon_full, lat_full)
        #
        #
        #
        fig = plt.figure(1)
        ax = fig.add_subplot(111)
        ax.set_title(title)
        ax.set_xlabel('Longitude',labelpad=20)
        ax.set_ylabel('Latitude',labelpad=20)
        ax.axis([0, 10, 0, 10])
        #
        m = Basemap(projection='cyl', resolution='l', \
                    llcrnrlon=np.min(lon_full)-res_lon, llcrnrlat=np.min(lat_full)-0.5*res_lat, \
                    urcrnrlon=np.max(lon_full)+0.5*res_lon, urcrnrlat=np.max(lat_full)+res_lat)
        
        m.drawcoastlines(linewidth=0.5)
        m.drawcountries(linewidth=0.5)
        m.drawparallels(np.arange(-180.,180.,0.1),labels=[1,0,0,0]) # draw parallels
        m.drawmeridians(np.arange(-80.,80.,0.1),labels=[0,0,0,1]) # draw meridians
        #
        cmap = 'RdBu'
        cmap = 'Pastel1'
        cmap = 'prism'
        cNorm = matplotlib.colors.Normalize(vmin=np.nanmin(basins_full), vmax=np.nanmax(basins_full))
        #
        x, y = m(lon_full-0.5*res_lon, lat_full+0.5*res_lat)
        z = np.ma.array(basins_full, mask=np.isnan(basins_full)) 
        datamap = m.pcolor(x, y, z, zorder=0, cmap=cmap)
        #
        plt.colorbar(datamap, cmap=cmap, norm=cNorm, orientation="vertical")
        #
        # Add polygone of atmospheric grid
        #
        for poly in atmpolys :
            xl,yl=m(np.array(poly)[:,0], np.array(poly)[:,1])
            plt.plot(xl,yl, color="m")
        #
        nbpts, nbvmax = basin_sz.shape
        for ib in range(nbpts) :
            for iz in range(nbvmax) :
                if basin_sz[ib,iz] > 0 :
                    x0,y0 = m(outlet[ib,iz,1], outlet[ib,iz,0])
                    # basin_lshead decides if this is an outlow point or not
                    if hsuper.basin_lshead[ib,iz] < undef_int :
                        #
                        if hsuper.outflow_grid[ib,iz] >= 0 :
                            # We know we flow into another box
                            ibo = hsuper.outflow_grid[ib,iz]-1
                            if ibo == ib :
                                # Same grid box
                                if hsuper.outflow_basin[ib,iz] < undef_int :
                                    INFO("hsuper.outflow_basin[ib,iz] = "+str(hsuper.outflow_basin[ib,iz]))
                                    ERROR("We cannot be here as we flow into an HTU of the same grid but we do not know the HTU index")
                                else :
                                    izo = hsuper.outflow_basin[ib,iz]-1
                                    lonout, latout, dist = closestpoint(outlet[ib,iz,1], outlet[ib,iz,0], hsuper.basin_sz[ib,izo], \
                                                                  hsuper.basin_pts[ib,izo,:,:], lon[ib,:,:], lat[ib,:,:])
                                    x1,y1 = m(lonout,latout)
                                    plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.0005, head_width=0.001, head_length=0.001, fc='r',ec='r')
                            else :
                                # We do not yet know if we flow into another grid
                                print("Flow into another grid :", hsuper.outflow_grid[ib,iz]-1, hsuper.outflow_basin[ib,iz]-1)
                                if hsuper.outflow_basin[ib,iz] < undef_int :
                                    izo = hsuper.outflow_basin[ib,iz]-1
                                    for ii in range(hsuper.basin_sz[ibo,izo]) :
                                        ix = hsuper.basin_pts[ibo,izo,ii,0]-1
                                        iy = hsuper.basin_pts[ibo,izo,ii,1]-1
                                        print("Possible points : ", lon[ibo,ix,iy], lat[ibo,ix,iy])
                                    lonout, latout, dist = closestpoint(outlet[ib,iz,1], outlet[ib,iz,0], hsuper.basin_sz[ibo,izo], \
                                                                  hsuper.basin_pts[ibo,izo,:,:], lon[ibo,:,:], lat[ibo,:,:])
                                    print("neighbour ==", dist, "--", outlet[ib,iz,1], outlet[ib,iz,0], lonout, latout)
                                    x1,y1 = m(lonout,latout)
                                    plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.001, head_width=0.002, head_length=0.002, fc='w',ec=None)
                                else :
                                    # We only know the general direction
                                    lonout = outlet[ib,iz,1]+res_lon*np.sin(np.radians(hsuper.basin_lshead[ib,iz]))
                                    latout = outlet[ib,iz,0]+res_lat*np.cos(np.radians(hsuper.basin_lshead[ib,iz]))
                                    x1,y1 = m(lonout,latout)
                                    plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.002, head_width=0.004, head_length=0.004, fc='r',ec=None)
                        else :
                            # We only know the general direction
                            lonout = outlet[ib,iz,1]+res_lon*np.sin(np.radians(hsuper.basin_lshead[ib,iz]))
                            latout = outlet[ib,iz,0]+res_lat*np.cos(np.radians(hsuper.basin_lshead[ib,iz]))
                            x1,y1 = m(lonout,latout)
                            plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.002, head_width=0.004, head_length=0.004, fc='r',ec=None)
                    else :
                        if hsuper.outflow_grid[ib,iz] == -1 :
                            # Outflow of a large river
                            plt.scatter(x0, y0, marker='o', facecolor='b', edgecolor='b', s=4)
                        elif hsuper.outflow_grid[ib,iz] == -2 :
                            # Outflow of a small river
                            plt.scatter(x0, y0, marker='o', facecolor='b', edgecolor='b', s=4)
                        elif hsuper.outflow_grid[ib,iz] == -3 :
                            # Outflow to lake or local pond (return flow)
                            plt.scatter(x0, y0, marker='s', facecolor='g', edgecolor='g', s=4)
                        elif hsuper.outflow_grid[ib,iz] == -4 :
                            # Flows into basin of same grid but not yet determines
                            if hsuper.outflow_basin[ib,iz] < undef_int :
                                izo = hsuper.outflow_basin[ib,iz]-1
                                lonout, latout, dist = closestpoint(outlet[ib,iz,1], outlet[ib,iz,0], hsuper.basin_sz[ib,izo], \
                                                              hsuper.basin_pts[ib,izo,:,:], lon[ib,:,:], lat[ib,:,:])
                                x1,y1 = m(lonout,latout)
                                plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.0005, head_width=0.001, head_length=0.001, fc='w',ec='w')
                            else :
                                plt.scatter(x0, y0, marker='s', facecolor='w', edgecolor='w', s=4)
                            
                        else :
                            # Flows into another same grid
                            izo = hsuper.outflow_basin[ib,iz]-1
                            lonout, latout, dist = closestpoint(outlet[ib,iz,1], outlet[ib,iz,0], hsuper.basin_sz[ib,izo], \
                                                                hsuper.basin_pts[ib,izo,:,:], lon[ib,:,:], lat[ib,:,:])
                            x1,y1 = m(lonout,latout)
                            plt.arrow(x0, y0, (x1-x0), (y1-y0), width=0.0005, head_width=0.001, head_length=0.001, fc='w',ec='w')
                        
        plt.savefig(filename+'.png', dpi=450)
        plt.close()
    return