#!/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 bokeh.plotting import figure
from bokeh.layouts import gridplot, column
from bokeh.io import save, output_file
from bokeh.models import Span, Label, Title

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, outfname, csvs, csvh, vv, dd, vmin, vmax, mm):
        self.outfile_hist = outfname
        self.csv_hist = csvh
        self.csv_stat = csvs
        self.valid_min = vmin
        self.valid_max = vmax
        self.varname = vv
        self.dataset = dd
        self.metadata = mm
        self.img_explanation = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'img/stats_explained.png')
        self.regions = {
          'total' : '[all latitudes]', 
          'N_polar' : '[60 to 90 deg]', 
          'N_midlat' : '[30 to 60 deg]', 
          'N_hadley' : '[0 to 30 deg]', 
          'S_hadley' : '[0 to -30 deg]', 
          'S_midlat' : '[-30 to -60 deg]', 
          'S_polar' : '[-60 to -90 deg]'
        }

    '''
    Plot a simplified boxplot diagram (only lines)
    '''
    def plot_boxplot_lines(self, col='File', 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.metadata['unit']),
                 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.metadata['varname']), text_font_style="italic"), 'above')
      p.add_layout(Title(text=self.dataset, 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/Max lines
      p.line(df['date'], df['Max'], color='#8D021F')
      p.circle(df['date'], df['Max'], color='#8D021F', legend_label='Maximum')
      p.line(df['date'], df['Med'], color='#800000', legend_label='Median')
      p.line(df['date'], df['Min'], color='#FF2800')
      p.circle(df['date'], df['Min'], color='#FF2800', legend_label='Minimum')

      # 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_instructions(self, df_stats, pw=1000, ph=300):

      # Load explanatory image
      im = Image.open(self.img_explanation)
      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

      # 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.metadata['unit']),
        'Observed data max = {} {}'.format(maxv, self.metadata['unit']),
        'Observed q1 min = {} {}'.format(minq1, self.metadata['unit']),
        'Observed q3 max = {} {}'.format(maxq3, self.metadata['unit']),
        'Observed IQD min = {} {}'.format(miniqd, self.metadata['unit']),
        'Observed IQD max = {} {}'.format(maxiqd, self.metadata['unit']),
      ]
      p=self.draw_citations(p, c, 0.9, 1.0, 'Observed statistics')

      return p

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

        vert_lines = []

        if not(self.valid_max is None):
          vert_lines.append(Span(location=self.valid_max, dimension='height', line_color='red'))
          vmax = np.where(l[nz] > self.valid_max)
          if len(vmax):
            fig.circle((l[nz])[vmax], y[vmax], color='red', fill_color='white', legend_label='outliers')

        if not(self.valid_min is None):
          vert_lines.append(Span(location=self.valid_min, dimension='height', line_color='red'))
          vmins = np.where(l[nz] < self.valid_min)
          if len(vmins):
            fig.circle((l[nz])[vmins], y[vmins], color='red', fill_color='white', legend_label='outliers')
        
        # Add vertical lines
        if vert_lines != []:
          fig.renderers.extend(vert_lines)

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

        # Get bins (left and right) + remove first element
        l = np.delete(df.columns.values, 0).astype(float) 
        xdist = abs(l[0]-l[1])
        r = [ l+xdist for v in l ]
        
        # Get 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])

        # Set up the figure same as before
        fig = figure(plot_width = pw, plot_height = ph,
                   title = '',
                   x_axis_label = '{} ({})'.format(self.varname, self.metadata['unit']), 
                   y_axis_label = 'Normalized Frequency',
                   toolbar_location='above',
                   y_range = [min(freq) - 0.2*(max(freq)-min(freq)), max(freq) + 0.1*(max(freq)-min(freq))],
                   tools='pan,wheel_zoom,box_zoom,reset')

        # Title/Subtitle
        fig.add_layout(Title(text='{} distribution {}'.format(self.metadata['varname'], self.regions[tag]), text_font_style="italic"), 'above')
        if tag == 'total':
            fig.add_layout(Title(text=self.dataset, text_font_size="13pt"), 'above')

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

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

        # Add mini-boxplot at the bottom
        low = min(freq) - 0.1 * (max(freq)-min(freq))
        lowh = min(freq) - 0.05 * (max(freq)-min(freq))
        lowl = min(freq) - 0.15 * (max(freq)-min(freq))
        fig.line([l[nz][0], l[nz][0]], [2, -2], color='#941333', line_dash=[5, 5])
        fig.line([l[nz][-1], l[nz][-1]], [2, -2], color='#941333', line_dash=[5, 5])

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

        # whiskers / IQD / Median
        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')

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

        return fig

    '''
    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)

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

        # Latitude regions
        p1 = None
        try:
          figs = []
          for t in [ 'N_polar', 'S_polar', 'N_midlat', 'S_midlat', 'N_hadley', 'S_hadley' ]:
            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:
          pass # Dataset has no regions

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

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

        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(p0, p1, p2, p3))

'''
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

    '''
    Calcuate time difference between two dates, return array
    '''
    def timediff(self, dates, gr='hours'):
        # Calculate diff in hours
        d0 = dates[0]
        tdiff = []
        for d1 in dates:
            if gr == 'hours':
                tdiff.append(float((d1 - d0).total_seconds()) / 3600.0)
            else:
                tdiff.append((d1 - d0).total_seconds())
            d0 = d1

        return tdiff

    '''
    Plot completeness (dates according to c3s, versus dates found in the data files)
    '''
    def plot_completeness(self, dates_available, dates_mandatory, pw=1500, ph=300):

        # Calculate minimum diff in seconds
        mindiff = min(self.timediff(dates_mandatory, gr='seconds')[1:])
        d = dates_mandatory[0]
        total_dates = []
        while d <= dates_mandatory[-1]:
            total_dates.append(d)
            d += timedelta(0, mindiff)

        # Mask where data is available
        mask_available = np.zeros(len(total_dates))
        mask_mandatory = np.zeros(len(total_dates))
        ind = 0
        for t in total_dates:
            if t in dates_available:
                mask_available[ind] = 1
            if t in dates_mandatory:
                mask_mandatory[ind] = 1
            ind+=1

        # Set up the figure same as before
        fig = figure(plot_width = pw, plot_height = ph,
                   title = '', y_range=[0,2],
                   x_axis_label = 'Date',
                   y_axis_label = 'Time Difference (hours)', x_axis_type='datetime',
                   toolbar_location='above',
                   tools='pan,wheel_zoom,box_zoom,reset')

        # Title/Subtitle
        fig.add_layout(Title(text='Temporal completeness', text_font_style="italic"), 'above')
        fig.add_layout(Title(text=self.dataset, text_font_size="13pt"), 'above')

        # xaxis sizes
        fig.xaxis.axis_label_text_font_size = '8pt'
        fig.yaxis.axis_label_text_font_size = '8pt'
        fig.title.text_font_size = '8pt'

        # Histograms and points where we have values
        fig.varea(y1=np.zeros(len(mask_mandatory)), y2=mask_mandatory, x=total_dates, legend_label='Dates in C3S catalogue', color='#949494')
        fig.varea(y1=np.zeros(len(mask_available)), y2=mask_available, x=total_dates, legend_label='Dates found in data files', color='#941333')

        output_file(self.outfile)
        save(fig)

    '''
    Plot consistency (timediff between dates)
    '''
    def plot_consistency(self, dates, pw=1500, ph=600):

        # Calculate diff in seconds
        tdiff = self.timediff(dates)

        # Set up the figure same as before
        fig = figure(plot_width = pw, plot_height = ph,
                   title = '',
                   x_axis_label = 'Date',
                   y_axis_label = 'Time Difference (hours)', x_axis_type='datetime',
                   toolbar_location='above',
                   tools='pan,wheel_zoom,box_zoom,reset')

        # Title/Subtitle
        fig.add_layout(Title(text='Temporal consistency', text_font_style="italic"), 'above')
        fig.add_layout(Title(text=self.dataset, text_font_size="13pt"), 'above')

        # xaxis sizes
        fig.xaxis.axis_label_text_font_size = '8pt'
        fig.yaxis.axis_label_text_font_size = '8pt'
        fig.title.text_font_size = '8pt'

        # Histograms and points where we have values
        fig.varea(y1=np.zeros(len(dates)-1), y2=tdiff[1:], x=dates[1:], name='Time difference', color='#941333', fill_alpha=0.6)
        fig.line(x=dates[1:], y=tdiff[1:], color='#941333', line_width=3)
        fig.square(x=dates[1:], y=tdiff[1:], color='#941333')

        output_file(self.outfile)
        save(fig)
