import numpy as np
from netCDF4 import Dataset
import warnings
warnings.filterwarnings("ignore")
from scipy.ndimage import uniform_filter as uf
import matplotlib.pyplot as plt
from scipy.signal import argrelextrema as lmax
import datetime as dt
from datetime import datetime
from glob import glob

#============================================================
# Function to convert date number to a date and time string

def time_since_inv(datenumber,reftimeandunits):
  """
  Convert a date number (e.g. time units since a date) to 
  a date in string format

  Input: 
  1) The date number
  2) The units and reference time of the date number:
     <units> since yyyy-mm-dd HH:MM:SS
     <units> can be any of the below:
     * seconds or secs
     * minutes or mins
     * hours or hrs
     * days or dys
     * weeks or wks

  Output:
  A float indicating the current time in the units 
   specified
  Requires datetime
  """

  # Split units string and make date and time object
  splitrtu = reftimeandunits.split(" ")
  unit = splitrtu[0]
  refdtob = dt.datetime(int(splitrtu[2][0:4]),
                        int(splitrtu[2][5:7]),
                        int(splitrtu[2][8:10]),
                        int(splitrtu[3][0:2]),
                        int(splitrtu[3][3:5]),
                        int(splitrtu[3][6:8]))

  # Get date number in seconds
  if unit.lower()=="minutes" or unit=="mins":
    datenumber = float(datenumber)*60
  elif unit.lower()=="hours" or unit=="hrs":
    datenumber = float(datenumber)*3600
  elif unit.lower()=="days" or unit=="dys":
    datenumber = float(datenumber)*86400
  elif unit.lower()=="weeks" or unit=="wks":
    datenumber = float(datenumber)*604800

  # Using timedelta return a string of the current 
  return(str(refdtob+dt.timedelta(seconds=datenumber)))

#============================================================
# Load and process data

def loadproc_data(inputfile):

  # Load data from TIMPS file:
  datasetFin = Dataset(inputfile)
  time = datasetFin.variables["time"][:]
  tunt = datasetFin.variables['time'].units
  area = datasetFin.variables["area"][:]
  vrr  = datasetFin.variables["vrr"][:]

  # Smooth time-series of area and VRR:
  fsize = int(len(area)/5)
  if fsize<5: fsize=5
  areaf = uf(area,size=fsize)
  vrrf = uf(vrr,size=fsize)

  # Get gradient of area and VRR:
  dareaf = np.gradient(areaf,time)
  dvrrf = np.gradient(vrrf,time)

  return(time,tunt,area,vrr,areaf,vrrf,dareaf,dvrrf)

#============================================================
# Function to get mature indices from a time-series

def get_mature_inds(ts,lmxi):
  """
  Get mature stage indices for MCS
  1) Time-series of area or vrr
  2) Indices of local maxima within time-series
  3) Factor of the maxima magnitude for which it can be
      defined as mature stage
    
  Output:
  A list of indices representing the mature stage in the 
   given time-series
  
  """

  # Get magnitude of each maxima
  lmx  = [ts[i] for i in lmxi]
   
  # Get magnitude of minima between each maxima
  sepmin = [min(ts[lmxi[i]:lmxi[i+1]]) 
            for i in range(len(lmxi)-1)]
    
  # Group into sets of maxima
  sep   = [True if ((sepmin[i]<lmx[i]*.75) or 
                    (sepmin[i]<lmx[i+1]*.75)) 
           else False for i in range(len(sepmin))]
  grps = [0]; c=0;
  for i in range(len(sep)): 
    if sep[i]: c=c+1
    grps.extend([c])

  # Select the max in each group
  mxinew = []
  for i in np.unique(grps):
    grpi = np.array(lmxi)[np.where(grps==i)]
    mxinew.extend([grpi[ts[grpi].argmax()]])
  lmx  = [ts[i] for i in mxinew]
    
  # Get minima between each actual maxima
  sepmin = [min(ts[:mxinew[0]])]+\
           [min(ts[mxinew[i]:mxinew[i+1]]) 
            for i in range(len(mxinew)-1)]+\
           [min(ts[mxinew[-1]:])]
  sepmini = [np.where(ts==i)[0].tolist()[0] 
             for i in sepmin]
    
  # Check if each maxima is sufficiently larger than 
  #  the surrounding minima
  mxtf = [True if ((lmx[i]*.75>sepmin[i]) and \
                   (lmx[i]*.75>sepmin[i+1])) 
          else False for i in range(len(lmx))]
  im = [[sepmini[i]+n 
         for n in np.where(ts[
          sepmini[i]:sepmini[i+1]]>lmx[i]*.75)[0].tolist()]
         for i in range(len(mxtf)) if mxtf[i]]

  return([item for sublist in im for item in sublist])

#============================================================
# Function to check if stages are within certain conditions

def maxchecks(ts,gmd):
    
  # Ensure growth/decay isn't spuriously defined when a 
  #  maxima is too close to the start of the time-series
  fper = ts[:ts.argmax()+1]
  if min(fper)>max(fper)*.75 or fper[0]>max(fper)*.75:
    for i in np.arange(0,ts.argmax(),1):
      gmd[i]=0
    
  # Ensure growth/decay isn't spuriously defined when 
  # a maxima is too close to the end of the time-series
  lper = ts[ts.argmax():]
  if min(lper)>max(lper)*.75 or lper[-1]>max(lper)*.75:
    for i in np.arange(ts.argmax(),len(ts)-1,1):
      gmd[i]=0

  # Ensure growth or decay doesn't occur until variable 
  #  at least below .75 of a maxima
  for ti in np.arange(ts.argmax(),len(ts),1):
    if ts[ti]>max(ts)*.9:
      if gmd[ti]==1 or gmd[ti]==3: gmd[ti]=0
    else:
      break

  return(gmd)

#============================================================
# Calculate mature stages

def get_gmd(areaf,vrrf,dareaf,dvrrf):

  # Calculate all local maxima in filtered VRR and area:
  mxia  = lmax(areaf,np.greater)[0].tolist()
  mxiv  = lmax(vrrf,np.greater)[0].tolist()

  # Instance where there are insufficient time-series to 
  #  calculate the stages
  if len(mxia)<1 or len(mxiv)<1 \
                 or max(areaf)*.5<min(areaf) \
                 or max(vrrf)*.5<min(vrrf):
    print("Time-series is insufficient to calculate MCS phases")
    return(False)

  # Separate multiple local maxima into separate mature stages
  ima = get_mature_inds(areaf,mxia)
  imv = get_mature_inds(vrrf,mxiv)

  # Select only the indices that are mature for both
  im = np.intersect1d(ima,imv)

  # Check against area and vrr change for growth/decay
  gmd = [2 if i in im 
    else 1 if (dareaf[i]>0 and dvrrf[i]>0)
    else 3 if (dareaf[i]<0 and dvrrf[i]<0)
    else 0
    for i in np.arange(0,len(areaf),1)]

  # Checks to ensure growth or decay isn't spuriously 
  #  assigned to a maxima that can't be defined as a maxima
  gmd = maxchecks(areaf,gmd)
  gmd = maxchecks(vrrf,gmd)

  # Ensure each phase lasts longer than just two time-steps
  for i,x in enumerate(gmd):
    if i==0: xnow=x; lenx=1; inds=[0]
    elif i==len(gmd)-1:
      lenx = lenx+1; inds = inds+[i]
      if lenx<4: 
        gmd = [0 if m in inds else n 
               for m,n in enumerate(gmd)]
    elif i>0:
      if xnow==x: lenx = lenx+1; inds = inds+[i]
      else:
        if lenx<4: 
          gmd = [0 if m in inds else n 
                 for m,n in enumerate(gmd)]
        xnow=x; lenx=1; inds=[i]


  # Ensure each growth/decay at least doubles or halves
  for i,x in enumerate(gmd):
    if i==0: xnow=x; inds=[0]
    elif i>0:
      if xnow==x: lenx = lenx+1; inds = inds+[i]
      elif i==len(gmd)-1:
        inds = inds+[i]
        if max(areaf[inds])*.5<min(areaf[inds]):
          gmd = [0 if m in inds else n 
                 for m,n in enumerate(gmd)]
        elif max(vrrf[inds])*.5<min(vrrf[inds]):
          gmd = [0 if m in inds else n 
               for m,n in enumerate(gmd)]
      else:
        if xnow==1 or xnow==3:
          if max(areaf[inds])*.5<min(areaf[inds]):
            gmd = [0 if m in inds else n 
                   for m,n in enumerate(gmd)]
          elif max(vrrf[inds])*.5<min(vrrf[inds]):
            gmd = [0 if m in inds else n 
                   for m,n in enumerate(gmd)]
        xnow=x; lenx=1; inds=[i]

  return(gmd)

#============================================================
# Function to plot time-series and stages

def plot_ts(time,timeunits,area,areaf,vrr,vrrf,gmd,infile,indir):
  
  # Get dates and x-tick labels
  datetime = [time_since_inv(t,tunt) for t in time]
  date     = time_since_inv(time[0],tunt)
  xtiks    = time[::4]
  xtiklabs = datetime[::4]

  # Generate figure
  fig = plt.figure(figsize=(6,9), dpi=100)

  # Generate axis and plot for area
  ax1 = fig.add_axes([0.14, 0.56,.85, .43])
  ax1.plot(time,area,"k",alpha=.5)
  ax1.plot(time,areaf,"k--")
  ax1.set_ylabel("area ($km^2$)")
  ax1.set_xticks(xtiks[1:])
  ax1.set_xticklabels([" "]*len(xtiks[1:]))
  ax1.set_xlim([time[0],time[-1]])
  ax1.set_ylim([0,max(area)+max(area)/20])
    
  # Shade mature stages
  for i in range(len(time)):
    if gmd[i]==1:
      ax1.fill_betweenx([0,max(area)+max(area)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="green",linewidth=0.0)
    if gmd[i]==2:
      ax1.fill_betweenx([0,max(area)+max(area)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="orange",linewidth=0.0)
    if gmd[i]==3:
      ax1.fill_betweenx([0,max(area)+max(area)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="red",linewidth=0.0)

  # Generate axis and plot for vrr
  ax2 = fig.add_axes([0.14, 0.12,.85, .43])
  ax2.plot(time,vrr,"k",alpha=.5)
  ax2.plot(time,vrrf,"k--")
  ax2.set_ylabel("VRR ($mm hr^{-1} km^2$)")
  ax2.set_xticks(xtiks[1:])
  ax2.set_xticklabels(xtiklabs[1:],rotation=30,ha='right')
  ax2.set_xlim([time[0],time[-1]])
  ax2.set_ylim([0,max(vrr)+max(vrr)/20])
    
  # Shade mature stages
  for i in range(len(time)):
    if gmd[i]==1:
      ax2.fill_betweenx([0,max(vrr)+max(vrr)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="green",linewidth=0.0)
    if gmd[i]==2:
      ax2.fill_betweenx([0,max(vrr)+max(vrr)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="orange",linewidth=0.0)
    if gmd[i]==3:
      ax2.fill_betweenx([0,max(vrr)+max(vrr)/20],
      [time[i]-.25,time[i]-.25],
      [time[i]+.25,time[i]+.25],
      alpha=0.1,color="red",linewidth=0.0)

  plt.savefig(f"fig_{infile[len(indir):len(indir)+13]}.png")

#============================================================
# Run code

indir = "/uufs/chpc.utah.edu/common/home/u0816744/zips2/TIMPS_online/TIMPS_files/2020/01/"
infid = "TIMPS_"

infiles = glob(f"{indir}{infid}*_20200110*.nc")

for infile in infiles:

  print(f"Working on {infile}")

  # Load required data
  time,tunt,area,vrr,areaf,vrrf,dareaf,dvrrf = loadproc_data(infile)

  # Calculate growth, mature, and decay stages
  gmd = get_gmd(areaf,vrrf,dareaf,dvrrf)
  if not gmd: gmd = [0]*len(areaf)

  # Plot figure
  plot_ts(time,tunt,area,areaf,vrr,vrrf,gmd,infile,indir)

#============================================================
# End