#!/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 (scientific assessment)
---
- 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'
]
checks_seasonal = [
    'bias', 'corr', 'fcrpss', 'frpss'
]


"""
Generate a Bokeh table plot with the legend+summary
"""
def plot_summary(passed, total):

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

	p = figure(title="{} - {} - ({})".format('EQC-figures production',
        '{} out of {} generated'.format(int(passed), int(total)), date.today().strftime('%d %b')),
		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 figure was generated'}, 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 figure is not readable'}, 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

    # 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=title,
        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 = '11pt'
    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, forms_dir, var, c, paramId):

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

    # Check cds forms availability
    mandatory_dates = get_forms_dates(dataset, var, '', '', forms_dir)

    # Append to dataframe
    if len(mandatory_dates):
        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, forms_dir, var, c, paramId):

    # Expected figures
    if 'seasonal' in dataset:
        total_figs = 96
        if 'pressure' in dataset:
            total_figs = total_figs*4
    if 'reanalysis' in dataset:
        total_figs = 1
        if c == 'departure':
            total_figs = 12
        elif c in ['iav', 'mean', 'rspread', 'spread' ]:
            total_figs = 13
        else:
            total_figs = 1
        if 'pressure' in dataset:
            total_figs = total_figs*4

    # 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(paramId, 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 == total_figs:
                res = 1
            elif nfigs > 0:
                res = 2
            else:
                res = 0

            # Check cds forms availability
            mandatory_dates = get_forms_dates(dataset, var, o, s, forms_dir)

            # Append to dataframe
            if len(mandatory_dates):
                df = df.append({'ds': dataset, 'id': '{}-{}-{}'.format(o,s,var), '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, forms_dir, var, c, info['paramId'])
                    else:
                        df = fill_value_reanalysis(df, dataset, datasetdir, forms_dir, var, c, info['paramId'])

            # Number of checks passed (only ok)
            passed = np.where(df.res.values == 1)[0].size
            total = float(df.res.values.size)

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

if __name__ == "__main__":
    main()