import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
import yaml
import datetime
from dateutil.relativedelta import relativedelta



def get_bardata(var,date,arch,location, coords, region, fcst_type):
    """ Function that loads tercile,skill and probs data
        from .nc files.
        
        var: str with the varible name to be load (tas,prlr,rsds...)
        date: str with the sdate in YYYYMMDD format
        arch: str with the archive dir
        location: str with the location
        return: df with the tercile,skill and probs
    """

    if fcst_type == "seasonal":
        skill_file = f"{arch}/{var}/{date}/{var}-skill_month{date[4:6]}.nc"
        terciles_file = f"{arch}/{var}/{date}/{var}-percentiles_month{date[4:6]}.nc"
        probs_file = f"{arch}/{var}/{date}/{var}-prob_{date}_{date[4:6]}.nc"
    else:
        week = datetime.datetime.strptime(date, "%Y%m%d").strftime("%V")
        skill_file = f"{arch}/{var}/{date}/{var}-skill_week{week}.nc"
        terciles_file = f"{arch}/{var}/{date}/{var}-percentiles_week{week}.nc"
        probs_file = f"{arch}/{var}/{date}/{var}-prob_{date}_{week}.nc"
       
    
    probs = xr.open_dataset(probs_file, decode_times=False,drop_variables='time_step')
    skill = xr.open_dataset(skill_file, decode_times=False, drop_variables='time_step')
    terciles = xr.open_dataset(terciles_file, decode_times=False, drop_variables='time_step')

    data = xr.merge([probs,skill,terciles])
    del probs,skill,terciles

    data = data.transpose('longitude', 'latitude', 'time')
    
    data = data.sel(dict(latitude=coords[1]),method='nearest')
    data = data.sel(dict(longitude=coords[0]),method='nearest')
    data.attrs['location'] = location
    data.attrs['Variable'] = location
    
    data['prob_bp10'] = data['prob_bp10'].where(data.rpss > 0,0)
    data['prob_ap90'] = data['prob_ap90'].where(data.rpss > 0,0)
    
    data['prob_bp10'] = data['prob_bp10'].where(data.bsp10 > 0,0)
    data['prob_ap90'] = data['prob_ap90'].where(data.bsp90 > 0,0)
  
    data['prob_bp10'] = data['prob_bp10'].where(data.prob_bp10 >= 40,0)
    data['prob_ap90'] = data['prob_ap90'].where(data.prob_ap90 >= 40,0)
       
    return(data)


def df_to_pd(df):
    
    """ Function that converts from xarray dataset to pandas object
        
        df: xarray dataset from get_data
        return: pandas dataframe for bar plot
    """
    
    location = df.attrs['location']
    df = df.to_dataframe()
    df = df.drop(['latitude','longitude','rpss','bsp10','bsp90'],axis=1)
    df = df[['prob_bp10','prob_bn','prob_n','prob_an','prob_ap90']]
    df = df.rename(columns={'prob_bp10':"percentile 10",
                            'prob_bn':"below normal",
                            'prob_n':"normal",
                            'prob_an':"above normal",
                            'prob_ap90':"percentile 90"})
    ## TODO: change hardcoded
    #df['Valid time'] = ['December','January','February'] 
    #df = df.set_index('Valid time')
    df['location'] = location
    
    return(df)

def load_colors(conf_file,var):
    """ loads colors from conf file
        
        conf_file: conf_file string with filename including path
        var: string with variable name (diff colors for precp)
        return: dict with color codes according to conf
    """
    with open(conf_file) as f:    
        conf = yaml.load(f, Loader=yaml.FullLoader)
        colors = conf['colors'][var]
    return(colors)

def load_var_attrs(conf_file,var):
    """ loads var atributes from conf file
        
        conf_file: conf_file string with filename including path
        var: string with variable name (diff colors for precp)
        return: dict with var attributes (units,longname...) according to conf
    """
    with open(conf_file) as f:    
        conf = yaml.load(f, Loader=yaml.FullLoader)
        var_attrs = conf['variables'][var]
    return(var_attrs)


def convert_to_float(frac_str):
    """ converts fraction in string to float
        
        frac_str: string with fraction ex. "2/4"
        return: float resulting from operation
    """
    try:
        return float(frac_str)
    except ValueError:
        num, denom = frac_str.split('/')
        try:
            leading, num = num.split(' ')
            whole = float(leading)
        except ValueError:
            whole = 0
        frac = float(num) / float(denom)
        return whole - frac if whole < 0 else whole + frac

def set_alpha(colors,skills,alpha):
    """ sets alpha channel of the rgba  code according to skill
        
        colors: rgb code in tupple form
        skills: list with the skill value for each lead time
        alpha: int alpha value to set in case skill < 0
        return: color in rgba form for each lead time
    """
    colors = tuple(convert_to_float(itup) for itup in colors)

    return([colors + (1,) if s >= 0 else colors + (alpha,) for s in skills])

def set_edge_color(skills):
    """ sets bars edge color as function of skill
        black if >= 0; grey if not
        
        skills: list with the skill value for each lead time
        return: str with color for each lead time
    """
    return(['k' if s >= 0 else 'grey' for s in skills])


def bar_plot(var,date,location,pd,colors,var_attrs,skills,terciles,extremes,outdir):
    
    """ generates the bar plot for terciles and extremes with
        skill mask
        
        var: string with the variable to plot 
        date: string with sdate to plot YYYYMMDD
        location: string with the location
        pd: pandas dataframe with probs
        var_attrs: dict with variable attributes
        skills: xarray dataset with skills scores
        terciles: xarray dataset with terciles scores
        extremes: xarray dataset with extremes scores
        outdir: string with the path to export the files
        return: saves plot in png format
    """
    
    # inits figure
    fig = plt.figure()
    ax = fig.add_axes([0,0,1,1])

    # sets colors and alpha according to skill values
    colors_bn = set_alpha(colors['bn']['rgba'],skills.rpss.values,0.2)
    colors_n = set_alpha(colors['n']['rgba'],skills.rpss.values,0.2)
    colors_an = set_alpha(colors['an']['rgba'],skills.rpss.values,0.2)
    colors_p10 = set_alpha(colors['p10']['rgba'],skills.rpss.values,1)
    colors_p90 = set_alpha(colors['p90']['rgba'],skills.rpss.values,1)
    
    # sets edge colors according to skill values
    edge_colors = set_edge_color(skills.rpss.values)
    
    # plots bars for terciles
    
    X = np.arange(pd.shape[0])*2
    
    pan = ax.bar(X + 0.3, pd['above normal'].to_numpy(),
                 color = colors_an, width = 0.3, label='Above normal',
                 edgecolor=edge_colors,zorder=3)

    pn = ax.bar(X + 0.0, pd['normal'].to_numpy(),
                color = colors_n, width = 0.3, label='Normal',
                edgecolor=edge_colors,zorder=3)

    pbn = ax.bar(X - 0.3, pd['below normal'].to_numpy(),
                 color = colors_bn, width = 0.3, label='Below normal',
                 edgecolor=edge_colors,zorder=3)
    # plots bars for extremes

    p90 = ax.bar(X + 0.5, pd['percentile 90'].to_numpy(),
                 color = colors_p90, width = 0.1, 
                 label='Upper extreme',edgecolor='k',zorder=3)
    
    p10 = ax.bar(X - 0.5, pd['percentile 10'].to_numpy(),
                 color = colors_p10, width = 0.1, 
                 label='Lower extreme',edgecolor='k',zorder=3)
    
    for cont,rect in enumerate(pn):
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['n']['hex'],fontsize=12)
        if skills.rpss[cont].values < 0:
            plt.text(rect.get_x() + rect.get_width() / 2.0, 20, 
                     'Skill < 0%', ha='center', 
                     va='bottom',c='k',fontsize=22,fontweight='roman')
    
    # if prob > 40 print text with prob value
    for rect in pan:
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['an']['hex'],fontsize=12)      
    for rect in pbn:
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['bn']['hex'],fontsize=12) 
    
    # if extreme exists plots text with extreme value
    for cont,rect in enumerate(p10):
        height = rect.get_height()
        if height != 0:
            plt.text(rect.get_x() - rect.get_width(), height+1, 
                     '{:.0f}%\n(< {:.1f}{})'.format(height,float(extremes.p10[cont]),var_attrs['units']), 
                     ha='right', va='center',color=colors['p10']['hex'],fontsize=11)
            
    for cont,rect in enumerate(p90):
        height = rect.get_height()
        if height != 0:
            plt.text(rect.get_x() + rect.get_width(), height+1,
                     '{:.0f}%\n(> {:.1f}{})'.format(height,float(extremes.p90[cont]),var_attrs['units']), 
                     ha='left', va='center',color=colors['p90']['hex'],fontsize=11)

    location = pd.location.values[0]  
    sdate = datetime.datetime.strptime(date, "%Y%m%d")

    title = '{} ({})'.format(var_attrs['name'],location)
    subtitle = f'Seasonal forecast issued on {sdate.strftime("%b %Y" )}'

    plt.text(x=-0.8, y=115, s=title, fontsize=20, weight='bold')
    plt.text(x=-0.8, y=107, s=subtitle, fontsize=14)

    ax.set_ylabel('Probability (%)',fontsize=14)
    #ax.set_xlabel('Valid time')

    # sets terciles ticks 
    tcks = [-0.15,0.15]
    ax.set_xticks(tcks + [i+2 for i in tcks] + [i+4 for i in tcks])
    terciles_vals = list(terciles.isel(time=0).to_array().values) + list(terciles.isel(time=1).to_array().values) + list(terciles.isel(time=2).to_array().values)
   
    if var == 'acprec':
        terciles_vals = [0 if ter < 0 else ter for ter in terciles_vals]
        #terciles_vals = [np.nan for ter in terciles_vals]
    
    plt.text(x=5, y=-10, s='[{}]'.format(var_attrs['units']), fontsize=14,rotation=0, ha='center',
             color='grey')

    ax.set_xticklabels([f'{i:.1f}' if i is not np.nan else '' for i in terciles_vals],
                       fontdict={'fontsize':'11'},rotation=45,color='grey',
                       rotation_mode='anchor',ha='right',va='top')
    #ax.tick_params(length=8)

    x_labels = [(sdate + relativedelta(months=+leadtime+1)).strftime("%B")
                for leadtime in range(3)]## TODO: hardcoded leadtimes
    x_positions = [0,2,4]
    for ind,label in enumerate(x_labels):
        plt.text(x=x_positions[ind], y=-22, s=label, fontsize=14,rotation=0, ha='center')
  
    #ax.legend(bbox_to_anchor=(1, 1.35),frameon=False,fontsize=11)
    ax.legend(bbox_to_anchor=(0.2,0.83,0.6,0.2), 
              loc="upper center",
              mode="expand", borderaxespad=1.5, ncol=2,
              frameon=False)

    leg = ax.get_legend()

    # sets legend colors, to avoid colors with low opacity
    hex_colors = [tuple(convert_to_float(itup) 
                  for itup in colors[i]['rgba']) 
                  for i in ['an','n','bn','p90','p10']]
    for cont,handle in enumerate(leg.legendHandles):   
        handle.set_color(hex_colors[cont])

    ax.set_ylim(0,100)
    ax.set_xlim(-1,5)
    ax.grid(axis='y',alpha=0.4,zorder=0);

    #plt.show()
    fig.savefig(f'{outdir}/bar_plot-{var}-{date}-{location}.png',dpi=130,bbox_inches='tight')
    #fig.savefig('test.png',dpi=130,bbox_inches='tight')
    plt.close()

    
def bar_plot_subseas(var,date,location,pd,colors,var_attrs,skills,terciles,extremes,outdir):
    
    """ generates the bar plot for terciles and extremes with
        skill mask
        
        var: string with the variable to plot 
        date: string with sdate to plot YYYYMMDD
        location: string with the location
        pd: pandas dataframe with probs
        var_attrs: dict with variable attributes
        skills: xarray dataset with skills scores
        terciles: xarray dataset with terciles scores
        extremes: xarray dataset with extremes scores
        outdir: string with the path to export the files
        return: saves plot in png format
    """
    
    # inits figure
    fig = plt.figure()
    ax = fig.add_axes([0,0,1,1])

    # sets colors and alpha according to skill values
    colors_bn = set_alpha(colors['bn']['rgba'],skills.rpss.values,0.2)
    colors_n = set_alpha(colors['n']['rgba'],skills.rpss.values,0.2)
    colors_an = set_alpha(colors['an']['rgba'],skills.rpss.values,0.2)
    colors_p10 = set_alpha(colors['p10']['rgba'],skills.rpss.values,1)
    colors_p90 = set_alpha(colors['p90']['rgba'],skills.rpss.values,1)
    
    # sets edge colors according to skill values
    edge_colors = set_edge_color(skills.rpss.values)
    
    # plots bars for terciles
    
    X = np.arange(pd.shape[0])*2
    
    pan = ax.bar(X + 0.3, pd['above normal'].to_numpy(),
                 color = colors_an, width = 0.3, label='Above normal',
                 edgecolor=edge_colors,zorder=3)

    pn = ax.bar(X + 0.0, pd['normal'].to_numpy(),
                color = colors_n, width = 0.3, label='Normal',
                edgecolor=edge_colors,zorder=3)

    pbn = ax.bar(X - 0.3, pd['below normal'].to_numpy(),
                 color = colors_bn, width = 0.3, label='Below normal',
                 edgecolor=edge_colors,zorder=3)
    # plots bars for extremes

    p90 = ax.bar(X + 0.5, pd['percentile 90'].to_numpy(),
                 color = colors_p90, width = 0.1, 
                 label='Upper extreme',edgecolor='k',zorder=3)
    
    p10 = ax.bar(X - 0.5, pd['percentile 10'].to_numpy(),
                 color = colors_p10, width = 0.1, 
                 label='Lower extreme',edgecolor='k',zorder=3)
    
    for cont,rect in enumerate(pn):
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['n']['hex'],fontsize=12)
        if skills.rpss[cont].values < 0:
            plt.text(rect.get_x() + rect.get_width() / 2.0, 20, 
                     'Skill < 0%', ha='center', 
                     va='bottom',c='k',fontsize=18,fontweight='roman')
    
    # if prob > 40 print text with prob value
    for rect in pan:
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['an']['hex'],fontsize=12)      
    for rect in pbn:
        height = rect.get_height()
        if height >= 40:
            plt.text(rect.get_x() + rect.get_width() / 2.0, height+1, 
                     f'{height:.0f}%', ha='center', 
                     va='bottom',c=colors['bn']['hex'],fontsize=12) 
    
    # if extreme exists plots text with extreme value
    for cont,rect in enumerate(p10):
        height = rect.get_height()
        if height != 0:
            plt.text(rect.get_x() - rect.get_width(), height+1, 
                     '{:.0f}%\n(< {:.1f}{})'.format(height,float(extremes.p10[cont]),var_attrs['units']), 
                     ha='right', va='center',color=colors['p10']['hex'],fontsize=11)
            
    for cont,rect in enumerate(p90):
        height = rect.get_height()
        if height != 0:
            plt.text(rect.get_x() + rect.get_width(), height+1,
                     '{:.0f}%\n(> {:.1f}{})'.format(height,float(extremes.p90[cont]),var_attrs['units']), 
                     ha='left', va='center',color=colors['p90']['hex'],fontsize=11)

    location = pd.location.values[0]  
    sdate = datetime.datetime.strptime(date, "%Y%m%d")

    title = '{} ({})'.format(var_attrs['name'],location)
    subtitle = f'Subseasonal forecast issued on {sdate.strftime("%d %b %Y" )}'

    plt.text(x=-0.8, y=135, s=title, fontsize=20, weight='bold')
    plt.text(x=-0.8, y=127, s=subtitle, fontsize=14)

    ax.set_ylabel('Probability (%)',fontsize=14)
    #ax.set_xlabel('Valid time')

    # sets terciles ticks 
    tcks = [-0.15,0.15]
    ax.set_xticks(tcks + [i+2 for i in tcks] + [i+4 for i in tcks] + [i+6 for i in tcks])
    terciles_vals = list(terciles.isel(time=0).to_array().values) + list(terciles.isel(time=1).to_array().values) + list(terciles.isel(time=2).to_array().values) + list(terciles.isel(time=3).to_array().values)
   
    if var == 'acprec':
        terciles_vals = [0 if ter < 0 else ter for ter in terciles_vals]
        #terciles_vals = [np.nan for ter in terciles_vals]
    
    plt.text(x=7, y=-10, s='[{}]'.format(var_attrs['units']), fontsize=14,rotation=0, ha='center',
             color='grey')

    ax.set_xticklabels([f'{i:.1f}' if i is not np.nan else '' for i in terciles_vals],
                       fontdict={'fontsize':'11'},rotation=45,color='grey',
                       rotation_mode='anchor',ha='right',va='top')
    #ax.tick_params(length=8)

    x_labels = [f'{(sdate + relativedelta(days=+4) + relativedelta(days=+(leadtime)*7)).strftime("%d/%m")}-{(sdate + relativedelta(days=+4) + relativedelta(days=+(leadtime)*7) + relativedelta(days=+6)).strftime("%d/%m")}'
                for leadtime in range(4)]## TODO: hardcoded leadtimes
    
    x_positions = [0,2,4,6]
    for ind,label in enumerate(x_labels):
        plt.text(x=x_positions[ind], y=-22, s=label, fontsize=10,rotation=0, ha='center')
  
    #ax.legend(bbox_to_anchor=(1, 1.35),frameon=False,fontsize=11)
    ax.legend(bbox_to_anchor=(0.2,0.85,0.6,0.2), 
              loc="upper center",
              mode="expand", borderaxespad=1.5, ncol=2,
              frameon=False)

    leg = ax.get_legend()

    # sets legend colors, to avoid colors with low opacity
    hex_colors = [tuple(convert_to_float(itup) 
                  for itup in colors[i]['rgba']) 
                  for i in ['an','n','bn','p90','p10']]
    for cont,handle in enumerate(leg.legendHandles):   
        handle.set_color(hex_colors[cont])

    ax.set_yticks([0,10,20,30,40,50,60,70,80,90,100])
    ax.set_yticklabels(["0","10","20","30","40","50","60","70","80","90","100",""])
    ax.set_ylim(0,120)
    ax.set_xlim(-1,7)
    ax.grid(axis='y',alpha=0.4,zorder=0);

    #plt.show()
    fig.savefig(f'{outdir}/bar_plot-{var}-{date}-{location}.png',dpi=130,bbox_inches='tight')
    plt.close()
