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

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=4, alpha=alp,
             markeredgecolor=c, markeredgewidth=1, c="None")
   elif ptype == "lowextreme":
      m.plot(x,y, linestyle='none', marker="v", markersize=4, alpha=alp,
             markeredgecolor=c, markeredgewidth=1, c="None")
   elif ptype == "hightercile":
      m.plot(x,y, linestyle='none', marker="o", markersize=1.5, alpha=alp,
                  c="None", markeredgecolor=c, markeredgewidth=1.5)
   else:
      m.plot(x,y, linestyle='none', marker="o", markersize=0.7, alpha=alp,
                  c="None", markeredgecolor=c, markeredgewidth=0.7)


def generate_map(probs_file, skill_file, date, leadtime,var, system, outdir):

   sdate = datetime.datetime.strptime(date, "%Y%m%d")
   delta_days = datetime.timedelta(days=4)
   idelta_weeks = datetime.timedelta(weeks=leadtime)
   fdelta_weeks = datetime.timedelta(weeks=leadtime+1)

   ileadtime = sdate + delta_days + idelta_weeks
   fleadtime = sdate + delta_days + fdelta_weeks - datetime.timedelta(days=1)

   ileadtime = ileadtime.strftime("%b %d" )
   fleadtime = fleadtime.strftime("%b %d" )
   sdate = sdate.strftime("%d %b %Y" )

   filename = '{}/{}-{}-{}-lead{}.png'.format(outdir,var, system, date, leadtime+1)

   title = 'Sub-seasonal forecast of temperature'
   subtitle = 'Issued on {}, valid for {} to {}'.format(sdate, ileadtime,fleadtime)

   #########################
   # SETTINGS ########################
   ########################

   domain = [-27, 51, 68, 23] # lats from bigger to smaller
   extreme_threshold = 40 # %
   bigcircle_threshold = 50 # %
   skill_threshold = 0  # %

   # colors
   blue = '#2BA4B4'
   red ='#FF764D'
   grey = '#ACACAC'
   land = '#DEDEDE'
   ocean = '#F7F7F7'

   #########################
   # DATA FILTERING ########################
   #########################

   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.assign_coords(longitude=(((data.longitude + 180) % 360) - 180))
   data = data.sortby('longitude')

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

   t = data.where(data.rpss >= skill_threshold, drop=False)

   n = t.where((t.prob_n > t.prob_an) &(t.prob_n > t.prob_bn), drop=False)
   sn = n.where(n.prob_n < bigcircle_threshold, drop=False)
   bnn = n.where(n.prob_n >= bigcircle_threshold, drop=False)

   b = t.where((t.prob_bn > t.prob_n) &(t.prob_bn > t.prob_an), drop=False)
   sb = b.where(b.prob_bn < bigcircle_threshold, drop=False)
   bbb = b.where(b.prob_bn >= bigcircle_threshold, drop=False)

   a = t.where((t.prob_an > t.prob_n) &(t.prob_an > t.prob_bn), drop=False)
   sa = a.where(a.prob_an < bigcircle_threshold, drop=False)
   baa = a.where(a.prob_an >= bigcircle_threshold, drop=False)

   p90_extremes = data.where(data.bsp90 > skill_threshold, drop=False)
   p90_extremes = p90_extremes.where(p90_extremes.prob_ap90 > extreme_threshold, drop=False)

   p10_extremes = data.where(data.bsp10 > skill_threshold, drop=False)
   p10_extremes = p10_extremes.where(p10_extremes.prob_bp10 > extreme_threshold, drop=False)

   #########################
   # PLOT  ########################
   #########################

   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=land, alpha=1)
   m.drawcoastlines(linewidth=0,color=land)

   m.drawlsmask(land_color=land,
                ocean_color=ocean)

   plotglyph(m, sa, leadtime, "tercile", red)
   plotglyph(m, sn, leadtime, "tercile", grey)
   plotglyph(m, sb, leadtime, "tercile", blue)

   plotglyph(m, baa, leadtime, "hightercile", red)
   plotglyph(m, bnn, leadtime, "hightercile", grey)
   plotglyph(m, bbb, leadtime, "hightercile", blue)

   plotglyph(m, p90_extremes, leadtime, "upextreme", red)
   plotglyph(m, p10_extremes, leadtime, "lowextreme", blue)

   plt.text(x=4.7, y=7.2*10**6, s=title, fontsize=10, weight='bold')
   plt.text(x=4.7, y=6.9*10**6, s=subtitle, fontsize=8, alpha=0.75)

   m.drawcountries(color=ocean,linewidth=1)
   plt.savefig(filename,dpi=200)
   plt.close()

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:])


