#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#     
#     Authors: 
#       Joan Sala Calero (joan.sala@bsc.es)
#       Jesus Pena-Izquierdo (jesus.pena@bsc.es)

import pandas as pd
from PIL import Image
from datetime import date, datetime, timedelta
from math import pi

from bokeh.plotting import figure
from bokeh.layouts import gridplot, column
from bokeh.io import save, output_file
from bokeh.models import Label, Title, Div, Button, CustomJS, Paragraph
from bokeh.models import ColorBar, LinearColorMapper, PrintfTickFormatter, ColumnDataSource

from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_grib import *
from dqc_wrapper.utils.utils_stats import *

'''
Class: Bokeh plots (valid_ranges)
---
- General plot functions for histograms 
- Statistical formulas needed to produce the plots
'''
class BokehPlot:

    def __init__(self, title, outfname, csvs, csvh, vv, dd, vmin, vmax, control_limits, mm, uu, msg, res):
        self.title = title
        self.outfile_hist = outfname
        self.csv_hist = csvh
        self.csv_stat = csvs
        self.valid_min = vmin
        self.valid_max = vmax
        self.control_limits = control_limits
        self.varname = vv
        self.dataset = dd
        self.varname = mm['varname']
        self.units = uu
        self.img_explanation = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'img/stats_explained.png')
        self.img_header = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'img/header.png')
        self.img_footer = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'img/footer.png')
        self.regions = {
          'total' : '[all latitudes, all timesteps]',
          'N_polar' : 'North/Polar [60 to 90 deg]',
          'N_midlat' : 'North/Mid-latitude [30 to 60 deg]',
          'N_hadley' : 'North/Equator [0 to 30 deg]',
          'S_hadley' : 'South/Equator [0 to -30 deg]',
          'S_midlat' : 'South/Mid-latitude [-30 to -60 deg]',
          'S_polar' : 'South/Polar [-60 to -90 deg]'
        }

        # Res msg
        if res == 'ok':
            self.final_msg = 'OK: {}'.format(msg)
            self.final_msg_html = '<b style="color:green">OK:</b> {}'.format(msg.replace('[', '<br/>['))
        elif res == 'warn':
            self.final_msg = 'WARNING: {}'.format(msg)
            self.final_msg_html = '<b style="color:orange">WARNING:</b> {}'.format(msg.replace('[', '<br/>['))
        else:
            self.final_msg = 'ERROR: {}'.format(msg)
            self.final_msg_html = '<b style="color:red">ERROR:</b> {}'.format(msg.replace('[', '<br/>['))

        # Specific text for each dataset family
        family = 'reanalysis'
        if 'seasonal' in self.dataset:
            family = 'seasonal'

        # Load text definitions
        self.overview_result = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/overview_result.txt'))
        self.summary_text = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/plot_summary.txt'))
        self.summary_text_uerra = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/plot_summary_uerra.txt'))
        self.method_text = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/caption_methodology.txt'))
        self.caption_hist = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/caption_histogram_{}.txt'.format(family)))
        self.caption_tseries = load_text_file(
            os.path.join(os.path.dirname(os.path.realpath(__file__)), 'texts/caption_tseries.txt'))

    '''
    Plot a simplified boxplot diagram (only lines)
    '''
    def plot_boxplot_lines(self, reg, col='File_dt', pw=1000, ph=500):

        # generate some synthetic time series for six different categories
        df = pd.read_csv(self.csv_stat)

        # Add date column to dataframe (take date from the filenames)
        tt = get_dates_from_files(df[col])
        df['date'] = tt
        df = df.sort_values(by=['date'])

        # Define figure
        pw=1000
        ph=600

        # Set up the figure same as before
        p = figure(plot_width = pw, plot_height = ph,
                 title = '',
                 x_axis_label = 'Time',
                 y_axis_label = '{} ({})'.format(self.varname, self.units),
                 x_axis_type='datetime',
                 toolbar_location='above',
                 tools='pan,wheel_zoom,box_zoom,reset')
        p.toolbar.logo = None

        # Title/Subtitle
        p.add_layout(Title(text='{} variability [all latitudes, all timesteps]'.format(self.varname)), 'above')
        p.add_layout(Title(text=self.title, text_font_size="13pt"), 'above')

        # whiskers / IQD
        p.line(df['date'], df['q1']-1.5*df['IQR'], color='#ED2939', line_dash=[1, 2])
        p.line(df['date'], df['q3']+1.5*df['IQR'], color='#ED2939', line_dash=[1, 2])
        p.varea(x=df['date'], y1=df['q1']-1.5*df['IQR'], y2=df['q3']+1.5*df['IQR'], color='#A45A52', alpha=0.1, legend_label='1.5*IQR')

        # Boxes (Q1-Q3)
        p.line(df['date'], df['q1'], color='#800000', line_dash=[1, 2])
        p.line(df['date'], df['q3'], color='#800000', line_dash=[1, 2])
        p.varea(x=df['date'], y1=df['q1'], y2=df['q3'], color='#800000', alpha=0.25, legend_label='IQR (Q3-Q1)')

        # Min/Med/Max lines
        mx = df.groupby(col).max()
        mn = df.groupby(col).min()
        p.line(mx['date'], mx['Max'], color='#8D021F')
        p.line(df['date'], df['Med'], color='#800000', legend_label='Median', line_width=2)
        p.line(mn['date'], mn['Min'], color='#8D021F')

        # Min/Max circles + tooltips
        p.circle(mn['date'], mn['Min'], color='#8D021F', legend_label='Minimum/Maximum')
        p.circle(mx['date'], mx['Max'], color='#8D021F')

        # Tolerance
        if not(self.valid_min is None):
            p.line([df['date'].iloc[0], df['date'].iloc[-1]], [self.valid_min, self.valid_min], line_color='magenta',
               legend_label='Tolerance limits', line_dash=[5, 5], line_width=2)
        if not(self.valid_max is None):
            p.line([df['date'].iloc[0], df['date'].iloc[-1]], [self.valid_max, self.valid_max], line_color='magenta',
               legend_label='Tolerance limits', line_dash=[5, 5], line_width=2)

        # Outliers (tolerance limits)
        vmax = df.loc[df['Max'] > self.valid_max]
        vmin = df.loc[df['Min'] < self.valid_min]
        if len(vmax.values):
            p.circle(vmax.date.values, vmax.Max.values, color='magenta', fill_color='white', legend_label='Detected outliers')
        if len(vmin.values):
            p.circle(vmin.date.values, vmin.Min.values, color='magenta', fill_color='white', legend_label='Detected outliers')

        # Outliers (control limits)
        if not(self.control_limits[reg]['vmin_nmax'] is None):
            p.line([df['date'].iloc[0], df['date'].iloc[-1]], [self.control_limits[reg]['vmin_nmax'], self.control_limits[reg]['vmin_nmax']], line_color='purple',
               legend_label='Control limits', line_dash=[5, 5], line_width=2)
        if not(self.control_limits[reg]['vmax_nmax'] is None):
            p.line([df['date'].iloc[0], df['date'].iloc[-1]], [self.control_limits[reg]['vmax_nmax'], self.control_limits[reg]['vmax_nmax']], line_color='purple',
               legend_label='Control limits', line_dash=[5, 5], line_width=2)

        # Plot summary figure
        info =  self.plot_instructions(df)

        return p, info

    '''
    Draw a list of text lines into a graph
    '''
    def draw_citations(self, p, citations, ystart, xstart, title):

        # For every citation a row
        c = Label(x=xstart, y=ystart,
                text=title, render_mode='css', text_font_size='8pt', text_color='firebrick',
                border_line_color='white', border_line_alpha=1.0, text_font_style='bold',
                background_fill_color='white', background_fill_alpha=1.0)
        p.add_layout(c)
        r=1
        for text in citations:
            c = Label(x=xstart, y=ystart - r*0.05,
                      text=text, render_mode='css', text_font_size='8pt',
                      border_line_color='white', border_line_alpha=1.0,
                      background_fill_color='white', background_fill_alpha=1.0)
            p.add_layout(c)
            r+=1

        return p

    '''
    Plot instructions, description of parameters (all in one png image)
    '''
    def plot_image(self, image_path, pw=1000, ph=150, full=False):

        # Load explanatory image
        im = Image.open(image_path)
        im = im.convert("RGBA")
        imarray = np.array(im)[::-1]

        p = figure(plot_width = pw, plot_height = ph, x_range=(0, 1.5), y_range=(0, 1), tools="")
        p.image_rgba(image=[imarray], x=0, y=0, dw=0.6, dh=1.0)

        p.xaxis.visible = False
        p.xgrid.visible = False
        p.yaxis.visible = False
        p.ygrid.visible = False
        p.xgrid.grid_line_color = None
        p.ygrid.grid_line_color = None
        p.toolbar.logo = None
        p.toolbar_location = None

        return p

    '''
    Plot instructions, description of parameters (all in one png image)
    '''
    def plot_instructions(self, df_stats):

        p = self.plot_image(self.img_explanation, pw=1000, ph=300)

        # Plot help
        fd = df_stats['date'].iloc[0].strftime('%d %b %Y')
        ld = df_stats['date'].iloc[-1].strftime('%d %b %Y')
        minv = min(df_stats['Min'])
        maxv = max(df_stats['Max'])
        c = [
            'Q1 = 25% of values (lower)',
            'Q3 = 75% of values (upper)',
            'IQR = Interquartile Range (Q3 - Q1)',
            'IQD = Interquartile Deviation (±1.5 * IQR)',
            'σ = standard deviation'
        ]
        p=self.draw_citations(p, c, 0.9, 0.65, 'Definitions')

        # Plot summary table
        fd = df_stats['date'].iloc[0].strftime('%d %b %Y')
        ld = df_stats['date'].iloc[-1].strftime('%d %b %Y')
        minv = format_float(min(df_stats['Min']))
        maxv = format_float(max(df_stats['Max']))
        minq1 = format_float(min(df_stats['q1']))
        maxq3 = format_float(max(df_stats['q3']))
        miniqd = format_float(min(df_stats['q1']-1.5*df_stats['IQR']))
        maxiqd = format_float(max(df_stats['q3']+1.5*df_stats['IQR']))
        c = [
            'Time coverage [{} to {}]'.format(fd, ld),
            'Observed data min = {} {}'.format(minv, self.units),
            'Observed data max = {} {}'.format(maxv, self.units),
            'Observed q1 min = {} {}'.format(minq1, self.units),
            'Observed q3 max = {} {}'.format(maxq3, self.units),
            'Observed IQD min = {} {}'.format(miniqd, self.units),
            'Observed IQD max = {} {}'.format(maxiqd, self.units),
        ]
        p=self.draw_citations(p, c, 0.9, 1.0, 'Observed statistics')

        return p

    '''
    Create vertical span and circle outliers
    '''
    def create_vertical_span(self, value, fig, color, l, y, nz, legend, outliers_draw=False, max=False):
        if not(value is None):
            fig.line([value, value], [0, 1], line_color=color, legend_label=legend, line_dash=[5, 5], line_width=2)
            if max:
                vmax = np.where(l[nz] > value)
                if len(vmax) and outliers_draw:
                    fig.circle((l[nz])[vmax], y[vmax], color=color, fill_color='white', legend_label='Detected outliers')
            else:
                vmin = np.where(l[nz] < value)
                if len(vmin) and outliers_draw:
                    fig.circle((l[nz])[vmin], y[vmin], color=color, fill_color='white', legend_label='Detected outliers')

    '''
    Add valid min/max to plot (vertical lines + highlights)
    '''
    def add_valid_minmax(self, reg, fig, l, y, nz, vertical=True):

        if vertical:
            self.create_vertical_span(self.valid_max, fig, 'magenta', l, y, nz,
                                      'Tolerance limits', outliers_draw=True, max=True)
            self.create_vertical_span(self.valid_min, fig, 'magenta', l, y, nz,
                                      'Tolerance limits', outliers_draw=True, max=False)
            self.create_vertical_span(self.control_limits[reg]['vmin_nmax'], fig, 'purple', l, y, nz,
                                      'Control limits', outliers_draw=False, max=True)
            self.create_vertical_span(self.control_limits[reg]['vmax_nmax'], fig, 'purple', l, y, nz,
                                      'Control limits', outliers_draw=False, max=False)

    '''
    Plot a single histogram, return figure
    '''
    def plot_histogram_figure(self, df, tag='total', col='File', pw=500, ph=250, log=False):

        # Get bins (left and right) + remove first element ['File']
        left = np.delete(df.columns.values, 0).astype(float)
        xdist = abs(left[0]-left[1])
        middle = left + xdist/2.0
        right = left + xdist
        
        # Get bins count + remove first element
        y = np.delete(df.loc[df[col] == tag].values, 0).astype(float)

        # Get non-zero bins
        nz = np.nonzero(y)

        # Percentages instead of number of bins
        freq = y[nz]/np.sum(y[nz])

        # Add mini-boxplot at the bottom
        low = -0.1 * (max(freq)-min(freq))
        lowh = -0.05 * (max(freq)-min(freq))
        lowl = -0.15 * (max(freq)-min(freq))

        # Yrange adjusted to fix also boxplot
        yrange = [lowl+lowh, max(freq) + 0.1*(max(freq)-min(freq))]
        if yrange[0] == yrange[1]:
            yrange = [ yrange[0], yrange[0]+0.1 ]

        # Set up the figure same as before
        if log:
            fig = figure(plot_width = pw, plot_height = ph,
                       title = '',
                       x_axis_label = '{} ({})'.format(self.varname, self.units),
                       y_axis_label = 'Normalized Frequency (log)',
                       toolbar_location='above',
                       y_axis_type = 'log',
                       tools='pan,wheel_zoom,box_zoom,reset')
        else:
            fig = figure(plot_width = pw, plot_height = ph,
                       title = '',
                       x_axis_label = '{} ({})'.format(self.varname, self.units),
                       y_axis_label = 'Normalized Frequency',
                       toolbar_location='above',
                       y_range = yrange,
                       tools='pan,wheel_zoom,box_zoom,reset')

        # Title/Subtitle
        if tag == 'total':
            fig.add_layout(Title(text='{} {}'.format(self.varname, self.regions[tag]), text_font_style="bold"), 'above')
            fig.add_layout(Title(text=self.title, text_font_size="12pt"), 'above')
        else:
            fig.add_layout(Title(text='{}'.format(self.regions[tag])), 'above')

        # xaxis sizes
        font_size = '9pt'
        if tag != 'total':
            font_size = '8pt'
        fig.xaxis.axis_label_text_font_size = font_size
        fig.yaxis.axis_label_text_font_size = font_size
        fig.title.text_font_size = font_size

        # Histograms and points where we have values
        fig.vbar(x=middle[nz], width=xdist, bottom=0, top=freq, color='#941333', legend_label='Histogram of data values')
        
        # Dash where there are values
        fig.dash(x=middle[nz], y=freq, color='#941333')
        fig.toolbar.logo = None
        
        # Vertical lines indicating valid_min/valid_max if available
        self.add_valid_minmax(tag, fig, middle, freq, nz)

        # Compute approximate stats from the histogram
        q1, med, q3 =  stats_from_hist(middle, y)
        iqr = q3 - q1

        # whiskers / IQD / Median / MinMax
        fig.line([q1 - 1.5 * iqr, q1], [low, low], color='#941333')
        fig.line([q3, q3 + 1.5 * iqr], [low, low], color='#941333')
        fig.line([q1 - 1.5 * iqr, q1 - 1.5 * iqr], [lowl, lowh], color='#941333')
        fig.line([q3 + 1.5 * iqr, q3 + 1.5 * iqr], [lowl, lowh], color='#941333')
        fig.line([q1, q1], [lowl, lowh], color='#941333')
        fig.line([med, med], [lowl, lowh], color='#941333')
        fig.line([q3, q3], [lowl, lowh], color='#941333')
        fig.line([min(left[nz]), min(left[nz])], [0, lowl], color='#941333', line_dash=[2, 2])
        fig.line([max(right[nz]), max(right[nz])], [0, lowl], color='#941333', line_dash=[2, 2])

        # Boxes (Q1-Q3) / MinMax
        fig.hbar(y=low, height=lowl-lowh, left=q1, right=q3, color='#800000', fill_alpha=0.25, line_color='#941333')

        # Legend
        fig.legend.visible = (tag=='total')

        return fig

    '''
    Create a button that displays html text information within the plots
    '''
    def create_button_information(self, text, pw=1000, ph=330, label_txt="Show/Hide methods applied and outputs description", show=True):
        div = Div(text="""{}""".format(text), width=pw, height=ph)
        div.visible = show
        showHide = """
            o1.visible = !(o1.visible);
        """
        butt = Button(label=label_txt, callback=CustomJS(args={'o1': div}, code=showHide))
        return butt, div

    '''
    Plot multiple histograms into a grid and return a grid of the plots (file) 
    '''
    def plot_histograms(self, just_figures=False):

        # Read histograms into pandas
        df_hist = pd.read_csv(self.csv_hist)

        # Summary text
        buttS, divS = self.create_button_information(self.summary_text, ph=340)
        buttSu, divSu = self.create_button_information(self.summary_text_uerra, ph=340)
        _, divO = self.create_button_information(self.overview_result.replace('####', self.final_msg_html), ph=130)
        butth, divH = self.create_button_information(self.caption_hist, ph=100, label_txt='Show/Hide plot description')
        buttT, divT = self.create_button_information(self.caption_tseries, ph=100, label_txt='Show/Hide plot description')
        buttM, divM = self.create_button_information(self.method_text, ph=490, label_txt='Show/Hide methodology')

        # Total histogram
        p0 = self.plot_histogram_figure(df_hist, tag='total', pw=1000, ph=500)

        # Latitude regions
        p1 = None

        # Lattitude regions are shown for all datasets except for UERRA
        if not 'uerra' in self.dataset:
            try:
                figs = []
                for t in [ 'N_polar', 'S_polar', 'N_midlat', 'S_midlat', 'N_hadley', 'S_hadley' ]: # Order of hadley cells
                    figs.append(self.plot_histogram_figure(df_hist, tag=t))
                # Group plots into grid
                p1 = gridplot(figs, ncols=2, toolbar_options=dict(logo=None))
            except Exception as e:
                pass # Dataset has no regions
        

        # Boxplot dates / Summary figure
        p2, p3 = self.plot_boxplot_lines('total')

        # Save to file
        print('INFO: Saving plot {}'.format(self.outfile_hist))
        output_file(self.outfile_hist)

        # For UERRA, p1 is never shown
        if 'uerra' in self.dataset:
            if just_figures:
                if p2 is None:
                    return [ p0, p3 ]
                else:
                    return [p0, p2, p3]
            else:
                if p2 is None:
                    save(column(p0, p3))
                else:
                    save(column(buttSu, divSu, divO, p0, butth, divH, p2, buttT, divT, p3, buttM, divM))
        else:
            if just_figures:
                if p1 is None or p2 is None:
                    return [ p0 ]
                else:
                    return [p0, p1, p2, p3]
            else:
                if p1 is None or p2 is None:
                    save(p0)
                else:
                    save(column(buttS, divS, divO, p0, butth, divH, p1, p2, buttT, divT, p3, buttM, divM))

'''
Class: Temporal plots (temporal completeness and consistency)
---
- General plot functions for timeseries 
'''
class TemporalPlot:

    def __init__(self, outfname, ds):
        self.outfile = outfname
        self.dataset = ds
        self.months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']

    '''
    Generate square figure (months, years, rate)
    '''
    def generate_figure(self, title, df, years, pw, ph):
        
        # For xrange
        if len(years) == 1:
            years.append(years[0])

        if 'consistency' in title:
            # Blue to yellow
            mapper = LinearColorMapper(palette="Viridis256", low=df.rate.min(), high=df.rate.max()+1)
            tp = [('File (date)', '@month @year'), ('#Timesteps', '@rate')]
        else:
            # Legend colors (not available=0, error=1, ok=2)
            colors = ["#dddddd", "#8b0000", "#75968f"]
            mapper = LinearColorMapper(palette=colors, low=0, high=2)
            tp = [('File (date)', '@month @year')]

        TOOLS = "hover,save"
        p = figure(title="{} - {}".format(title, self.dataset),
                   x_range=sorted(list(set(years))), y_range=self.months.reverse(),
                   x_axis_location="above", plot_width=pw, plot_height=ph,
                   tools=TOOLS, toolbar_location='below', tooltips=tp)

        # Axis settings
        p.grid.grid_line_color = None
        p.axis.axis_line_color = None
        p.axis.major_tick_line_color = None
        p.axis.major_label_text_font_size = "7pt"
        p.axis.major_label_standoff = 0
        p.xaxis.major_label_orientation = pi / 3

        # Plot rectangles
        p.rect(x="year", y="month", width=1, height=1,
               source=df,
               fill_color={'field': 'rate', 'transform': mapper},
               line_color='gray')

        # Legend
        if 'consistency' in title:
            color_bar = ColorBar(color_mapper=mapper, major_label_text_font_size="8pt",
                                 formatter=PrintfTickFormatter(format="%d"),
                                 label_standoff=6, border_line_color=None, location=(0, 0))
            p.add_layout(color_bar, 'right')

        return p

    '''
    Plot completeness (#dates for every month/year)
    '''
    def plot_completeness(self, dates_av, dates_mn, pw=1280, ph=768):

        # Init dataframe
        df = pd.DataFrame(columns=['year', 'month', 'rate', 'ind'])

        # Loop over dates (init)
        years = []
        for yy in range(dates_mn[0].year, dates_mn[-1].year + 1):
            years.append(str(yy))
            for m in self.months:
                df = df.append({'year': str(yy), 'month': m, 'rate': 0, 'ind': '{}{}'.format(yy,m) }, ignore_index=True)

        # Loop over dates (init)
        for d in dates_mn:
            df.loc[df.ind == d.strftime("%Y%b"), 'rate'] += 1
        for d in dates_av:
            df.loc[df.ind == d.strftime("%Y%b"), 'rate'] += 1

        # Figure define
        p = self.generate_figure('Temporal completeness', df, years, pw, ph)

        # Save to file
        output_file(self.outfile)
        save(p)

    '''
    Plot consistency (#dates for every month/year)
    '''
    def plot_consistency(self, dates, pw=1280, ph=768):

        # Init dataframe
        df = pd.DataFrame(columns=['year', 'month', 'rate', 'ind'])
        dts = sorted(dates)

        # Loop over dates (init)
        years = []
        for yy in range(dts[0].year, dts[-1].year + 1):
            years.append(str(yy))
            for m in self.months:
                df = df.append({'year': str(yy), 'month': m, 'rate': 0, 'ind': '{}{}'.format(yy,m) }, ignore_index=True)

        # Loop over dates (init)
        for d in dts:
            df.loc[df.ind == d.strftime("%Y%b"), 'rate'] += 1

        # Figure define
        p = self.generate_figure('Temporal consistency', df, years, pw, ph)

        # Save to file
        output_file(self.outfile)
        save(p)
