#!/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)

"""
Script: Generate a summary of the EQC production status
---
- Sync c3sforms
- Read QAR directory
- Generate bokeh table plot
"""

import pandas as pd
from datetime import date
from math import pi
from bokeh.io import save, output_file
from bokeh.models import LinearColorMapper, ColumnDataSource
from bokeh.plotting import figure
from bokeh.layouts import column

from dqc_wrapper.utils.utils_files import *
from dqc_wrapper.utils.utils_forms import *

# Available columns
checks_reanalysis = [
    'departure', 'g_anom', 'iav', 'mean', 'rspread', 'spread', 'spread_sc', 'spread_ts', 'trend_map', 'STATUS'
]
checks_seasonal = [
    'bias', 'corr', 'fcrpss', 'frpss'
]

"""
Generate a Bokeh table plot with the legend+summary
"""
def plot_summary():
    # Legend colors (err, not processed, ok, warn)
    colors = ["#8b0000", "#dddddd", "#75968f", "#ffa500"]
    mapper = LinearColorMapper(palette=colors, low=-1, high=2)

    p = figure(title="{} (updated: {})".format('EQC-figures production legend ', datetime.now().strftime('%d %b %H:%M')),
               x_axis_location="above", plot_width=1000, plot_height=75,
               tools='', toolbar_location='below', x_range=[-1, 4],
               tooltips=[('msg', '@msg')])

    p.title.align = 'center'
    p.title.text_font_size = '9pt'
    p.toolbar.logo = None
    p.toolbar_location = None
    p.xaxis.visible = False
    p.xgrid.visible = False
    p.yaxis.visible = False
    p.ygrid.visible = False

    legend_data = pd.DataFrame(columns=['res'])
    legend_data = legend_data.append(
        {'col': 0, 'row': 0, 'res': 0, 'msg': 'Unknown or data unavailable'}, ignore_index=True)
    legend_data = legend_data.append(
        {'col': 1, 'row': 0, 'res': 1, 'msg': 'Ok or test passed'}, ignore_index=True)
    legend_data = legend_data.append(
        {'col': 2, 'row': 0, 'res': 2, 'msg': 'Warning or missing information'}, ignore_index=True)
    legend_data = legend_data.append(
        {'col': 3, 'row': 0, 'res': -1, 'msg': 'Error or test not passed'}, ignore_index=True)

    p.rect(x="col", y="row", width=1, height=1,
           source=legend_data,
           fill_color={'field': 'res', 'transform': mapper},
           line_color='#444444')

    return p

"""
Generate a Bokeh table plot with tooltips
Each row is a QAR
Filter by keyword
"""
def table_plot(data, title):

    # Checks/Columns
    checks = checks_reanalysis
    if 'seasonal' in title:
        checks = checks_seasonal

    # Count totals and drop column
    all_test = data.loc[data['test'] == 'all']
    completed = all_test.res.sum()
    available = all_test.res.count()
    data = data.loc[data['test'] != 'all']

    # Legend colors (err, not processed, ok, warn)
    colors = ["#8b0000", "#dddddd", "#75968f", "#ffa500"]
    mapper = LinearColorMapper(palette=colors, low=-1, high=2)

    # Height of the plot (pixels)
    yy = data.id.unique()
    if len(yy) < 10:
        ph = len(yy) * 50
    elif len(yy) < 20:
        ph = len(yy) * 30
    elif len(yy) < 40:
        ph = len(yy) * 25
    else:
        ph = len(yy) * 18

    p = figure(title="{} ({} of {} done)".format(
        title, completed, available),
        x_range=checks, y_range=sorted(yy, reverse=True),
        x_axis_location="above", plot_width=1000, plot_height=ph,
        tools='', toolbar_location='below',
        tooltips=[('dataset', '@ds'), ('id', '@id'), ('test', '@test'), ('msg', '@msg')])

    p.title.align = 'center'
    p.title.text_font_size = '9pt'
    p.toolbar.logo = None
    p.toolbar_location = None
    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 / 2

    # Plot values
    p.rect(x="test", y="id", width=1, height=1,
           source=data,
           fill_color={'field': 'res', 'transform': mapper},
           line_color='#444444')

    return p

"""
Fill value on reanalysis table (variable)
"""
def fill_value_reanalysis(df, dataset, datasetdir, var, c):

    check_dir = os.path.join(datasetdir, '{}/{}'.format(var, c))
    figures_found = glob.glob(os.path.join(check_dir, '*.png'))
    nfigs = len(figures_found)

    df = df.append({'ds': dataset, 'id': var, 'test': c, 'res': int(nfigs > 0), 'msg': '{} figures found'.format(nfigs)}, ignore_index=True)
    return df

"""
Fill value on seasonal table (variable/origin/system)
"""
def fill_value_seasonal(df, dataset, datasetdir, var, c):

    # Available systems (chosen for analysis)
    origin_system = {
        'meteo_france': [5, 6, 7],
        'dwd': [2],
        'ecmwf': [4, 5],
        'cmcc': [3],
        'ukmo': [12, 13, 14]
    }
    for o,sys in origin_system.items():
        for s in sys:
            check_dir = os.path.join(datasetdir, '{}/{}_s{}/{}'.format(var, o, s, c))
            figures_found = glob.glob(os.path.join(check_dir, '*.png'))
            nfigs = len(figures_found)

            # Define results (number of figures)
            res = -1
            if nfigs == 96:
                res = 1
            elif nfigs > 0:
                res = 2
            else:
                res = 0
            df = df.append({'ds': dataset, 'id': '{}-{}_{}'.format(var,o,s), 'test': c, 'res': res, 'msg': '{} figures found'.format(nfigs)}, ignore_index=True)

    return df

"""
Generate EQC report main
"""
def main():

    # Get all ini files
    script_dir = os.path.dirname(os.path.realpath(__file__))
    grib1_def = os.path.join(script_dir, 'grib/grib1.json')
    grib2_def = os.path.join(script_dir, 'grib/grib2.json')

    # Results dir
    forms_dir = '/data/cds-forms-c3s'
    res_dir = '/shared/plots_scientific_assessment'
    os.chdir(res_dir)

    # For every dataset on the forms dir
    for f in os.listdir(forms_dir):

        # See if we find plots
        group = os.path.basename(f).split('-')[0]
        dataset = os.path.basename(f)
        rootdir = os.path.join(res_dir, group)
        datasetdir = os.path.join(rootdir, dataset)

        # Checks/Columns
        checks = checks_reanalysis
        if 'seasonal' in group:
            checks = checks_seasonal

        # Result available (loop per dataset/variable/check)
        if os.path.exists(datasetdir):

            # Init empty dataframe
            df = pd.DataFrame(columns=['ds', 'id', 'test', 'res', 'msg'])

            # Gather info from c3sforms
            ds_info, _ = get_dataset_information(dataset, forms_dir, grib1_def, grib2_def)

            # Fill tables
            for var,info in ds_info.items():
                for c in checks:
                    check_dir = os.path.join(datasetdir, '{}/{}'.format(var, c))
                    if 'seasonal' in dataset:
                        df = fill_value_seasonal(df, dataset, datasetdir, var, c)
                    else:
                        df = fill_value_reanalysis(df, dataset, datasetdir, var, c)

            # Plot dataframe
            outfile = os.path.join(res_dir, '{}.html'.format(dataset))
            s = plot_summary()
            p = table_plot(df, dataset)
            output_file(outfile, title=dataset)
            save(column(s,p))
            print('Written: {}'.format(outfile))

if __name__ == "__main__":
    main()