#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on 01.01.2016

@author: richter bettina

Edited on 01.01.2021

@author: mayer stephanie

Edited on 01.11.2022

@author: richter bettiina

"""

import datetime
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from dateutil.rrule import rrule, DAILY



#%%

###########################################################################################################
#####
#####   Name:       plot_single_profile
#####
#####   Purpose:    plot single profile for a given timestamp
#####               x-axis: Hand hardness
#####               y-axis: snow depth
#####               colors: grain type colorcode
#####               
#####   Optional:   additional variable can be a layer properity which is plotted as a line (seperate x-axis)
#####               
###########################################################################################################

def plot_single_profile(fig,ax, pro, var=None, ax_colorbar=None):
    
    hh = -pro['hand_hardness'] #get postive values for hand hardness 
    height = pro['height']
    
    if len(height) == 0: return
      
    #Define colorcode for grain types
    cgt=['greenyellow','darkgreen','pink','lightblue','blue','magenta','red','cyan','lightblue']#'red','grey']
    gt = divmod(pro['grain_type'],100)[0].astype(int)
    
    ########### contours ##########################
    hs = height[-1]
    ax.plot([0, 0], [0, hs], c='black', linewidth=1)  #y-axis line
    ax.plot([0, hh[0]], [0, 0], c='black', linewidth=1)  #x-axis line

    ################ Plot profile against hand hardness ########################
    
    ybottom=0
    for iy, y in enumerate(height):

        ######################################################        
        #Plot layer contours - horizontal line top of layer
        if iy == len(height)-1:
            ax.plot([0, hh[iy]], [y, y], c='black', linewidth=1)
        else:
            ax.plot([0, np.max([hh[iy], hh[iy + 1]])], [y, y], c='black', linewidth=1)        
        #Plot layer contours - vertical line right of layer
        ax.plot([hh[iy], hh[iy]], [ybottom, y], c='black', linewidth=1)

        ######################################################
        #Fill layer with color of grain type
        ax.fill_betweenx([ybottom, y], 0, hh[iy], color=cgt[int(gt[iy]) - 1])
        
        #cb = ax.pcolormesh([0, hh[iy]],[ybottom, y],np.ones((2,2))*gt[iy],cmap=cmap,vmin=vmin,vmax=vmax,alpha=0.9)
        ybottom = y
        
    ################ colorbar ########################################################
    if ax_colorbar:
        cmapcolorbar = ['greenyellow', 'darkgreen', 'pink', 'lightblue', 'blue', 'magenta', 'red', 'cyan']
        ticklabels = ['PP', 'DF', 'RG', 'FC', 'DH', 'SH', 'MF', 'IF']
        cmapc = mpl.colors.ListedColormap(cmapcolorbar)
        bounds = np.arange(len(cmapcolorbar) + 1)

        norm = mpl.colors.BoundaryNorm(bounds, cmapc.N)
        ticks = bounds + .5

        cb1 = mpl.colorbar.ColorbarBase(ax_colorbar, cmap=cmapc, norm=norm, ticks=ticks, orientation='vertical')
        cb1.ax.set_yticklabels(ticklabels)

    ########## Plot SH again, since layers are very thin not alost not visible ########################
    try:
        ish = int(np.where(( gt == 6 ))[0])
        ax.fill_betweenx([height[ish-1],height[ish]] , 0, hh[ish], color='magenta' )
        print('++++++++++++++ SHFOUND  in profile +++++++++')
    except: pass

    ax.set_xlim(0, 5.5)
    ax.set_xticks(np.arange(0, 5.5, 1))
    ax.set_xticklabels(['', 'F ', '4F', '1F', 'P', 'K'])
    ax.set_ylabel('Snow depth [cm]')
    ax.set_xlabel('Hand hardness')
    
    if hs < 200: 
        ax.set_ylim(0,200)
    else: 
        ax.set_ylim(0, hs+20)

    #^######### Plot Variable as line, seperate x-axis ########################
    if var:
        variable = pro[var]
        ax11 = ax.twiny()
        height_var= np.repeat(np.concatenate((np.array([0]), height)), 2)[1:-1]
        var_repeat = np.repeat(variable, 2)
        ax11.plot(var_repeat, height_var, c='black', linewidth=2.5, label=var)

        ax11.legend(loc=1)
        ax11.set_xlabel(var)
           
        return ax11  

#%%
###########################################################################################################
#####
#####   Name:       plot_seasonal_evolution
#####
#####   Purpose:    plot seasonal evolution of snowpack
#####               x-axis: Time
#####               y-axis: snow depth
#####               colors: grain type by colorcode or colorbar for layer properties, e.g. density
#####               
#####               
###########################################################################################################

def plot_seasonal_evolution(ax, prof, 
                      a = datetime.datetime(2021, 10, 1, 12),b = datetime.datetime(2022, 6, 30, 12), 
                      rule=None, #To make it faster: rule can be e.g. DAILY, then only one profile per day is plotted. Change to: #rule=DAILY,
                      cmap=None, 
                      var='grain_type', colorbar = True,
                      vmin=0,vmax=2):

    ts = sorted( prof['data'].keys() )
    if not a: a=ts[0]
    if not b: b=ts[-1]
    
    if rule:
        timestamps = rrule(rule, dtstart=a, until=b)
    else: 
        timestamps = np.array(ts)[ts.index(a):ts.index(b)]
        
    for dt in timestamps:
        try:
            pro=prof['data'][dt]
        except: continue
        #get layer height
        try:
            depth = pro['height']
        except: 
            continue
        depth_edges = np.concatenate((np.array([0]),depth))
        #Get time delta to plot
        x=[dt,dt+datetime.timedelta(days=1)]
        #Get grain type
        if var=='grain_type':
            cgt=['greenyellow','darkgreen','pink','lightblue','blue','magenta','red','cyan','lightblue']#,'red','grey']
            cmap=mcolors.ListedColormap(cgt)

            gt = divmod(pro['grain_type'],100)[0]
            gt=gt.astype(int)
            if len(gt)==0: 
                continue
            vmin=0.5
            vmax=len(cmap.colors)+0.5
            cb=ax.pcolormesh(x,depth_edges,np.array([gt,gt]).transpose(),cmap=cmap,vmin=vmin,vmax=vmax,alpha=0.9)
        else:
            try: 
                variable = pro[var]
            except: 
                print('No variable named '+var)
                return()
            cb=ax.pcolormesh(x,depth_edges,np.array([variable,variable]).transpose(),vmin=vmin,vmax=vmax)
            
    if colorbar:
        cb1=plt.colorbar(cb,ax=ax)
        if var=='grain_type':
            ticklabels = ['PP', 'DF', 'RG', 'FC', 'DH', 'SH', 'MF', 'IF','FC']
            cb1.ax.set_yticklabels(ticklabels)
            cb1.set_label('Grain type')
        elif var=='lwc':
            cb1.set_label('LWC [%]')
        else: cb1.set_label(var)
    
    ax.set_ylabel('Snow depth [cm]')
    
    return(cb)

