import matplotlib.pyplot as plt
import numpy as np
import datetime as dt
import os
import warnings


import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
import cartopy.feature as cfeature
import matplotlib.ticker as mticker
import matplotlib.patches as patches


def _get_unique(dic1,dic2):
    d = {}
    for i in dic1.keys():
        if i not in dic2.keys():
            d.update({i:dic1[i]})
    return d



class map:
    
    def __init__(self, x, y, xticks=None, yticks=None, center=110, **kwargs):
        """Create a geographical map object.

        The map object can be initializated one time, and be drew many times. Method
        calls are used to set paramters or drawing. Some paramters are shared by all
        the methods of this map object.

        **Arguments:**

        *x*
            A list that contains three or two items. The first two iterms are the
            beginning and the ending of the longitudes. The third iterm is optional
            which means the step from the beginning to the end, if it is not given, we
            would set it 1 as default.

        *y*
            The same as x, but was lattitude's.

        **Optional arguments:**

        *xticks*
            A iterable object. If labeled, label the specific longitudes.If is default,
            draw some label lines in the map.

        *yticks*
            The same as xticks, but was lattitude's.

        *center*
            If the projection is the default PlateCarree, center is the center longitude of this projection.

        **Returns:**
            A basemap object.

        *Other parameters:*
        
        *subplot*
            subplot infomation, set to (1,1,1) as default.

        *projection*
            A cartopy.crs object. It is the projection method of map.

        *transform*
            A cartopy.crs object that transform the data projection,
            only used when filling the data.

        *crs*
            A cartopy.crs, was used to define a map.

        *coast*
            Bool. Whether to draw the coast.(True as default)
            
        *coast_linewidth*   : default 1
        *coast_color*   : default "black"
        *coast_linestyle*  : default '-'
        *coast_resolution*  : default "110m"

        *borders*
            Bool. Whether to draw the border of the countries.(False as default)
            
        *border_linewidth*  : 1
        *border_linestyle*  : '-'
        *border_color*  : 'k'
        *border_resolution*  : "110m"

        *rivers*
            Bool. Whether to load the rivers.(False as default)     
        *river_linewidth*  : 1
        *river_linestyle*  : '-'
        *river_color*  : 'blue'
        *river_resolution*  : '110m',

        *lake*:  the same as *rivers*

        *lonlat*
            Bool. Whether to label the longitude and latitude.(True as default)

        *grid*: the same as borders

        *labels*
            A list that defines where to draw the lonlat labels.
            Set ['left','bottom'] as default.



        **Examples:**
            m = map([70,160],[10,60],borders=True)

        
            
            

        """
        self._para = {'projection':ccrs.PlateCarree(center), 'transform':ccrs.PlateCarree(),'crs':ccrs.PlateCarree(),
                      'coast':True,'coast_linewidth':1,'coast_color':'black','coast_linestyle':'-','coast_resolution':'110m',
               'borders':False,'border_linewidth':1,'border_linestyle':'-','border_color':'k','border_resolution':'110m',
                'rivers':False, 'river_linewidth':1,'river_linestyle':'-', 'river_color':'blue', 'river_resolution':'110m',
                'lakes':False, 'lake_linewidth':0.8,'lake_linestyle':'-', 'lake_color':'blue', 'lake_resolution':'110m',
                'lonlat':True, 'grid':False, 'grid_linewidth':0.9,'grid_linestyle':':','grid_color':'k', 'labels':['left','bottom'],

               }
        self._para.update(kwargs)
        self.xlims, self.ylims = x[:2], y[:2]
        x[1], y[1] = x[1]+0.0001, y[1]+0.0001
        self.x, self.y = np.arange(*x), np.arange(*y)
        self.xticks = np.round(np.linspace(*self.xlims,5)) if xticks is None else np.array(xticks)
        self.xticks[self.xticks>180] = self.xticks[self.xticks>180] - 360
        self.xticks = np.sort(self.xticks)
        self.yticks =  np.round(np.linspace(*self.ylims,5)) if yticks is None else np.array(yticks)
        self.__need_subplot, self.__extend, self.__need_add_subplot = True, True, True
        self.__subplot_args, self.__subplot_kwargs = kwargs.get('subplot',(1,1,1)),{}



    def axes(self,*args,projection=None):
        proj = projection if projection else self._para['projection']
        self.ax = plt.axes(*args,projection=proj)
        self.__need_subplot, self.__need_add_subplot = True, False
        return self.ax

    def _extend(self,x,y):
        if self.__extend:
            ext = np.zeros(4)            
            ext[0] = np.min([self.xlims[0],x[0]])
            ext[1] = np.max([self.xlims[1],x[-1]])
            ext[2] = np.min([self.ylims[0],y[0]])
            ext[3] = np.max([self.ylims[1],y[-1]])
            self.ax.set_extent(ext,crs=self._para['crs'])
            self.__extend = False

    def figure(self,*args,**kwargs):
        self.__need_subplot, self.__extend = True, True
        self.fig = plt.figure(*args,**kwargs)
        return self.fig
        
    def subplot(self, *args,**kwargs):
        self.__need_subplot, self.__need_add_subplot = True, True
        self.__subplot_args = args
        self.__subplot_kwargs = kwargs
        self._subplot()

    def _subplot(self):
        
        if self.__need_subplot:
            self._para.update(self.__subplot_kwargs)
            self.fig = plt.gcf()
            if self.__need_add_subplot:
                self.ax = self.fig.add_subplot(*self.__subplot_args,projection=self._para['projection'])
            self.__extend = True
            if self._para['coast']:
                self.ax.coastlines(resolution=self._para['coast_resolution'],lw=float(self._para['coast_linewidth']),
                                   color=self._para['coast_color'],linestyle=self._para['coast_linestyle'])
            if self._para['borders']:
                self.ax.add_feature(cfeature.BORDERS.with_scale(self._para['border_resolution']),linestyle=self._para['border_linestyle'],
                                    lw=float(self._para['border_linewidth']),edgecolor=self._para['border_color'])
            if self._para['rivers']:
                self.ax.add_feature(cfeature.RIVERS.with_scale(self._para['river_resolution']),linestyle=self._para['river_linestyle'],
                                    lw=float(self._para['river_linewidth']),color=self._para['river_color'])
            if self._para['lakes']:
                self.ax.add_feature(cfeature.LAKES.with_scale(self._para['lake_resolution']),linestyle=self._para['lake_linestyle'],
                                    lw=float(self._para['lake_linewidth']),color=self._para['lake_color'])
            
            if self._para['lonlat']:
                draw_labels = False if isinstance(self._para["projection"],ccrs.PlateCarree) else True
                if not self._para['grid']:
                    self._para['grid_linewidth'] = 0
                gl = self.ax.gridlines(draw_labels=draw_labels,linewidth=self._para['grid_linewidth'],
                                       linestyle=self._para['grid_linestyle'], color=self._para['grid_color'], alpha=0.8,
                                       xlocs=self.xticks,ylocs=self.yticks)
                for la in ['left','bottom','top','right']:
                    if la not in self._para["labels"]:
                        exec(f"gl.{la}_labels=False")
                if isinstance(self._para["projection"],ccrs.PlateCarree):
                    gl.xlocator = mticker.FixedLocator(self.xticks) 
                    gl.ylocator = mticker.FixedLocator(self.yticks)
                    
                    self.ax.set_xticks(self.xticks, crs=self._para['crs'])
                    self.ax.set_yticks(self.yticks, crs=self._para['crs'])
                    self.ax.xaxis.set_major_formatter(LongitudeFormatter(zero_direction_label=True))
                    self.ax.yaxis.set_major_formatter(LatitudeFormatter())
            self._extend(self.x,self.y)
            self.__need_subplot = False
        

    def contourf(self, *args , **kwargs):
        if len(args)==1:
            data,x,y = args[0],None,None
        elif len(args)==3:
            x,y,data = args
            
        dic = {'cbar':True,'fmt':'%d','pad':0.08,'fraction':0.04,
               'location':'right','cax':None}
        kw = _get_unique(kwargs,dic)
        dic.update(kwargs)
        self._subplot()
            
        X, Y = self.x if x is None else x, self.y if y is None else y
        self._extend(X,Y)
        self.contourf_map = self.ax.contourf(X, Y, data, transform=self._para['transform'],**kw)
        if dic['cbar']:
            ori = {'bottom':'horizontal','right':'vertical'}
            self.colorbar = plt.colorbar(self.contourf_map,orientation=ori[dic['location']],format = dic['fmt'],pad=dic['pad'],fraction=dic['fraction'],cax=dic['cax'])
        return self.contourf_map
            

    def contour(self, *args, **kwargs):
        if len(args)==1:
            data,x,y = args[0],None,None
        elif len(args)==3:
            x,y,data = args
        
        dic = {'clabel':True,'fmt':'%d','inline':True,'linewidths':0.85,
               'fontsize':8,'color':'k'}
        kw = _get_unique(kwargs,dic)
        if ('cmap' in kw.keys())&('colors' in kw.keys()):
            kw.pop('colors')
        dic.update(kwargs)
        self._subplot()
        
        X, Y = self.x if x is None else x, self.y if y is None else y
        self._extend(X,Y)
        self.contour_map = self.ax.contour(X, Y, data, transform=self._para['transform'],linewidths=dic['linewidths'],**kw)

        if dic['clabel']:
            cb = plt.clabel(self.contour_map,inline=dic['inline'], fontsize=dic['fontsize'],
                                colors=kwargs.get('color',kwargs.get('colors',None)),fmt=dic['fmt'])
        return self.contour_map


    def ttest_plot(self,logic_matrix, x=None, y=None,**kwargs):
        dic = {'color':'k','marker':'o','linestyle':'',
               'alpha':0.8,'markersize':1,}
        
        X, Y = self.x if x is None else x, self.y if y is None else y
        xx, yy = np.meshgrid(X, Y)
        x, y = xx[logic_matrix], yy[logic_matrix]
        self.ax.plot(x,y,color=dic['color'],marker=dic['marker'],linestyle=dic['linestyle'],
                      alpha=dic['alpha'],markersize=dic['markersize'],transform=self._para['transform'])
        return 

    def pcolor(self, *args, **kwargs):
        if len(args)==1:
            data,x,y = args[0],None,None
        elif len(args)==3:
            x,y,data = args
            
        dic = {'cbar':True,'fmt':'%d','pad':0.08,'fraction':0.04,
               'location':'right','cax':None,'extend':'neither'}
        kw = _get_unique(kwargs,dic)
        dic.update(kwargs)
        self._subplot()
            
        X, Y = self.x if x is None else x, self.y if y is None else y
        self._extend(X,Y)
        self.pcolor_map = self.ax.pcolor(X, Y, data,transform=self._para['transform'],**kw)
        if dic['cbar']:
            ori = {'bottom':'horizontal','right':'vertical'}
            cb = plt.colorbar(self.pcolor_map,orientation=ori[dic['location']],format = dic['fmt'],
                              pad=dic['pad'],fraction=dic['fraction'],cax=dic['cax'],extend=dic['extend'])
        return self.pcolor_map
        
    def quiver(self,*args, skip=None,headwidth=2.5,scale=200,width=0.004,color='black',legend=False,**kwargs):
        self._subplot()
        if len(args)==2:
            u,v = args[0]
        elif len(args)==4:
            x,y,u,v = args
        
        X, Y = self.x if x is None else x, self.y if y is None else y
        self._extend(X,Y)
        if skip is not None:
            xs, ys = skip
            self._X, self._Y = self._X[::ys,::xs], self._Y[::ys,::xs]
            u, v = u[::ys,::xs], v[::ys,::xs]
        self.quiver_map = self.ax.quiver(self._X,self._Y,u,v,headwidth=headwidth,scale=scale,width=width,color=color,transform=self._para['transform'])
        if legend:
            plt.quiverkey(self.quiver_map,**kwargs)

    def streamplot(self,u,v,x=None,y=None,skip=None,**kwargs):
        dic = dict(color='k',linewidth=0.6,arrowsize=0.8,arrowstyle='<-')
        dic.update(kwargs)
        self._subplot()
        X, Y = self.x if x is None else x, self.y if y is None else y
        self._extend(X,Y)
        self._X, self._Y = np.meshgrid(X, Y)
        if skip is not None:
            xs, ys = skip
            self._X, self._Y = self._X[::xs,::ys], self._Y[::xs,::ys]
            u, v = u[::xs,::ys], v[::xs,::ys]
        self.quiver_map = self.ax.streamplot(self._X,self._Y,u,v,transform=self._para['transform'],**dic)

    def plot(self,*args,**kwargs):
        return self.ax.plot(*args,**kwargs,transform=self._para['transform'])

    def add_shape(self,filename,linewidth=1,linestyle='-',encoding='gbk',rec=-1,method=None,**kwargs):
        from cartopy.io.shapereader import Reader
        self._subplot()
        dic = dict(edgecolor='k',facecolor=None)
        dic.update(kwargs)
        if not method:
            try:
                reader = Reader(filename)
                dic['facecolor'] =  dic['facecolor'] if dic['facecolor'] else 'none'
                enshicity = cfeature.ShapelyFeature(reader.geometries(),self._para['transform'], **dic)
                self.ax.add_feature(enshicity, linewidth = linewidth,linestyle=linestyle)
                reader.close()
            except:
                method = 1     
        if method:
            import matplotlib.patches as patches
            from shapefile import Reader
            if isinstance(rec,int) and (rec>0):
                rec = [rec]
            sf = Reader(filename,encoding=encoding)
            fill = True if dic['facecolor'] else False
            for i,sh in enumerate(sf.shapes()):
                if (rec==-1)or(i in rec): 
                    points = np.array(sh.points)
                    if len(sh.parts)==2:
                        self.ax.add_patch(patches.Polygon(points[sh.parts[1]:],linestyle=None,transform=self._para['transform'],
                                                      color=dic.get('color',dic['edgecolor']),fill=fill,facecolor=dic['facecolor']))
                        points = points[:sh.parts[1]]
                    self.ax.add_patch(patches.Polygon(points,linewidth=linewidth,linestyle=linestyle,transform=self._para['transform'],
                                                      color=dic.get('color',dic['edgecolor']),fill=fill,facecolor=dic['facecolor']))
        return

    def load_province(self,**kwargs):
        file = os.path.dirname(__file__)+'/shapefiles/cnhimap.shp'
        self.add_shape(file,**kwargs)

    def load_china(self,**kwargs):
        file = os.path.dirname(__file__)+'/shapefiles/china_country.shp'
        self.add_shape(file,**kwargs)

    def load_river(self,**kwargs):
        file = os.path.dirname(__file__)+'/shapefiles/rivers.shp'
        self.add_shape(file,**kwargs)

    def add_box(self,x,y,fill=False,**kwargs):
        tmp = [x[0],y[0],x[1]-x[0],y[1]-y[0]]
        self.ax.add_patch(patches.Rectangle(tmp[:2],tmp[2],tmp[3],transform=self._para['transform'],fill=fill,**kwargs))
        return

    def text(self,*args,**kwargs):
        self.ax.text(*args,**kwargs,transform=self._para['transform'])

    def title(self,label,*args,**kwargs):
        plt.title(label,*args,**kwargs)

    def show(self):
        if self.__need_subplot:
            self._subplot()
            self._extend(self.x,self.y)
        plt.show()

    def savefig(self,*arg,**kwargs):
        plt.savefig(*arg,**kwargs)

              
def set_xy_delta(step=None,axis='x'):
    if step==None:
        return
    else:
        ax = plt.gca()
    if (axis!='y')|(axis!='Y'):
        ax.xaxis.set_major_locator(plt.MultipleLocator(step))
    else:
        ax.yaxis.set_major_locator(plt.MultipleLocator(step))

def font(font='simsun'):
    plt.rcParams['font.sans-serif']=[font]
    plt.rcParams['axes.unicode_minus']=False


def mat_bar(matrix,**kwargs):
    para = {'number':True, 'width':0.2, 'legend':None,'fmt':'.1f',
            'fontsize':8,
            'color':['blue','deepskyblue','orange','red'],
        }
    para.update(kwargs)
    data = np.array([matrix]) if matrix.ndim<2 else np.array(matrix)
    shp = data.shape
    
    ax = plt.gca()
    ax.xaxis.set_major_locator(plt.MultipleLocator(1))
    for j in range(shp[0]):
        idx = j if j<len(para['color']) else j%len(para['color'])       
        w = 0.5*para['width']-para['width']*shp[0]/2+j*para['width']
        x, y = np.arange(1+w,shp[1]+w+1), data[j,:]

        ax.bar(x,y,width=para['width'],color=para['color'][idx])

        if para['number']:
            for a,b in zip(x,y):
                txt = '{:'+para['fmt']+'}'
                ax.text(a-0.5*para['width']+0.01,b+0.015*np.max(data),txt.format(b),fontsize=para['fontsize'])
        
        if para['legend'] is not None:
            ax.legend(para['legend'],framealpha=0.4)
            
        

        
        
        




if __name__=='__main__':
##    import matplotlib.patches as patches
##    m = map([70,130,2],[15,55,2])
##    d = np.arange(31*21).reshape(21,31)
##    c = m.contour(d,cmap=plt.cm.flag)
##    m.add_shape('./shapefiles/cnhimap.shp')
##    m.show()
    pass
