import numpy as np
import datetime
import xarray as xr
#import plotly.express as px
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
from mpl_toolkits.basemap import Basemap
import pandas as pd
import os
import sys, getopt
from dateutil.relativedelta import relativedelta

def get_domain(region):
    
    if region == "cat":
        domain = [0.2034, 3.3046, 42.87,40.54] # lats from bigger to smaller
    else:
        print(f'Region <{region}> NOT FOUND')
        
    return(domain)
        
def get_data(probs_file,skill_file,region):
    
    domain = get_domain(region)
    
    probs = xr.open_dataset(probs_file, decode_times=False)
    skill = xr.open_dataset(skill_file, decode_times=False)
    data = xr.merge([probs,skill])
    del probs,skill

    data = data.transpose('longitude', 'latitude', 'time', 'time_step')

    data = data.sel(dict(latitude=slice(domain[3],domain[2])))
    data = data.sel(dict(longitude=slice(domain[0],domain[1])))
    
    return(data)


def generate_plot(data,region,locations,leadtime,thrs,colors,subtitle,outfile):
    
    fig = plt.figure(figsize=(3, 3), dpi=300)
    
    domain = get_domain(region)

    m = Basemap(llcrnrlon=domain[0],llcrnrlat=domain[3],urcrnrlon=domain[1],
             urcrnrlat=domain[2],projection='mill', resolution='h')

    m.drawmapboundary(linewidth=0)
    m.fillcontinents(color=colors['land'], alpha=1)
    m.drawcoastlines(linewidth=0,color=colors['land'])

    m.drawlsmask(land_color=colors['land'],
                 ocean_color=colors['ocean'])

    for category in ["above","normal","below"]:

        c = {"above":colors['red'],"normal":colors['grey'],"below":colors['blue']}

        plotglyph(m,filter_tercile(data,category,thrs['bigcircle'],
                                  thrs['skill'],"small"),
                  leadtime,"lowcircle", c[category])

        plotglyph(m,filter_tercile(data,category,thrs['bigcircle'],
                                   thrs['skill'],"big"),
                  leadtime,"bigcircle", c[category])

    plotglyph(m, filter_extreme(data,"prob_ap90",thrs['extreme'],thrs['skill']),
              leadtime, "upextreme", colors['red'])

    plotglyph(m, filter_extreme(data,"prob_bp10",thrs['extreme'],thrs['skill']),
              leadtime, "lowextreme", colors['blue'])
    
    # plot locations:
    for location,point in locations.items():
        x,y = m(point[0],point[1])
        m.plot(x,y,linestyle='none', marker="+", markersize=5,
               markeredgecolor='k', markeredgewidth=0.5, c="None",alpha=0.4)

    if region == "cat":
        m.readshapefile('/esarchive/scratch/ptrascas/R/sunset/comarques/bm5mv21sh0tpc1_20200601_0', 'Comarca',
                        default_encoding='ISO8859-1',color=colors['ocean'],linewidth=0.5)
    else:
        m.drawcountries(color=colors['ocean'], linewidth=1)
        m.drawrivers(color=colors['ocean'], linewidth=1)
        
    #plt.suptitle(title, fontsize=10, weight='bold')
    plt.title(subtitle, fontsize=7, alpha=0.75)

    fig.tight_layout()
    plt.savefig(outfile,dpi=300)
    #plt.show()
    plt.close()

def plotglyph(m, data, time, ptype, c):
   lons = data.longitude.to_pandas().values
   lats = data.latitude.to_pandas().values
   lons, lats = np.meshgrid(lons,lats)

   ext = data.rpss.isel(time=time).to_masked_array().mask.transpose()
   lons = np.where(ext == True, False, lons)
   lats = np.where(ext == True, False, lats)

   alp = 0.8

   x,y = m(lons,lats)
   if ptype == "upextreme":
      m.plot(x,y, linestyle='none', marker="^", markersize=5, alpha=alp,
             markeredgecolor=c, markeredgewidth=1, c="None")
   elif ptype == "lowextreme":
      m.plot(x,y, linestyle='none', marker="v", markersize=5, alpha=alp,
             markeredgecolor=c, markeredgewidth=1, c="None")
   elif ptype == "bigcircle":
      m.plot(x,y, linestyle='none', marker="o", markersize=1.6, alpha=alp,
                  c="None", markeredgecolor=c, markeredgewidth=1.6)
   elif ptype == "lowcircle": 
      m.plot(x,y, linestyle='none', marker="o", markersize=1, alpha=alp,
                  c="None", markeredgecolor=c, markeredgewidth=1)
        
def filter_tercile(df,category,prob_thrs,skill_thrs,circle):
    
    if skill_thrs is not None:
        df = df.where(df.rpss >= skill_thrs, drop=False)
    
    # most-likely-tercile
    if category == "normal": 
        df = df.where((df.prob_n > df.prob_an) &(df.prob_n > df.prob_bn), drop=False)
        if circle == "big":
            df = df.where(df.prob_n >= prob_thrs, drop=False)
        elif circle == "small":
            df = df.where(df.prob_n < prob_thrs, drop=False)
    elif category == "below":
        df = df.where((df.prob_bn > df.prob_n) &(df.prob_bn > df.prob_an), drop=False)
        if circle == "big":
            df = df.where(df.prob_bn >= prob_thrs, drop=False)
        elif circle == "small":
            df = df.where(df.prob_bn < prob_thrs, drop=False)
    elif category == "above":
        df = df.where((df.prob_an > df.prob_n) &(df.prob_an > df.prob_bn), drop=False)    
        if circle == "big":
            df = df.where(df.prob_an >= prob_thrs, drop=False)
        elif circle == "small":
            df = df.where(df.prob_an < prob_thrs, drop=False)

    return(df)

def filter_extreme(df,category,prob_thrs,skill_thrs):
    
    skill_var = {"prob_bp10":"bsp10", "prob_ap90":"bsp90"}
    
    if skill_thrs is not None:
        df = df.where(df[skill_var[category]] > skill_thrs, drop=False)
    if prob_thrs is not None:
        df = df.where(df[category] > prob_thrs, drop=False)
    
    return(df)

def main(argv):

   usage = 'DST-plots.py -v <tas,tasmin,...> -d <date> -s <s2s,cfs> -o <outdir>'
   try:
      opts, args = getopt.getopt(argv,"hv:d:s:o:",["help","variable=","sdate=", "system="])
   except getopt.GetoptError:
      print(usage)
      sys.exit(2)

   for opt, arg in opts:
      if opt == '-h':
         print(usage)
         sys.exit()
      elif opt in ("-s", "--system"):
         system = str(arg)
      elif opt in ("-d", "--sdate"):
         date = arg
      elif opt in ("-v", "--variable"):
         var = arg
      elif opt in ("-o", "--outdir"):
         outdir = arg

   week = datetime.datetime.strptime(date, "%Y%m%d").strftime("%V")

   if system == "s2s":
      src_dir = '/esarchive/oper/S2S4E-data/weekly_statistics/S2S/'
      probs_file = src_dir + '/' + var + '/' + date + '/' + var + \
      '-prob_' + date + '_' + week + '.nc'
      skill_file = src_dir + '/' + var + '/' + date + '/' + var + \
      '-skill_week' + week + '.nc'
   else:
      src_dir = '/esarchive/oper/S2S4E-data/weekly_statistics/'
      probs_file = src_dir + '/' + var + '/' + var + \
      '-prob_' + date + '_' + week + '.nc'
      skill_file = src_dir + '/' + var + '/' + var + \
      '-skill_week' + week + '.nc'

   for leadtime in range(4):
      generate_map(probs_file, skill_file, date, leadtime, var, system, outdir)

if __name__ == "__main__":
   main(sys.argv[1:])


