#!/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 QAR production status
---
- Read QAR directory
- Generate bokeh table plot
"""

import os
import shutil
import pandas as pd
import configparser
from datetime import date
from math import pi

from bokeh.io import save, output_file
from bokeh.models import LinearColorMapper
from bokeh.plotting import figure
from bokeh.layouts import column

# Available tests
tests = [
	'file_format', 'standard_compliance', 
	'temporal_completeness', 'temporal_consistency',
	'spatial_completeness', 'spatial_consistency',
	'valid_ranges', 'netcdf_converter'
]

"""
Generate a Bokeh table plot with the legend+summary
"""
def plot_summary(title, data, width):

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

	p = figure(title="{} {} ({})".format(title, 'QAR production legend ', date.today().strftime('%d %b')),
		x_axis_location="above", plot_width=width, 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, filt, width):

	# Filter data
	data = data[data.ds == filt]

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

	# Sort by id
	data = data.sort_values('id', ascending=False)

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

	# length of the plot
	yy = data.id.unique()
	if len(yy) < 10:
		ph = len(yy)*40
	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(
		       filt, completed, available),
	     	   x_range=tests, y_range=yy,
	           x_axis_location="above", plot_width=width, 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

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

	return p

"""
OK/NOK to integers
"""
def ok_nok(val):
	if val == 'ok':
		return 1
	elif val == 'warn':
		return 2
	else:
		return -1 # N/A

"""
Get metadata for each QAR
"""
def get_metadata(ff):

	ini = configparser.ConfigParser()
	ini.read(ff)
	ds = ini['dataset']['cds_dataset']
	id = ini['dataset']['cds_variable']

	if 'seasonal' in ds:
		id = '{}-{}-{}'.format(ini['temporal_completeness']['origin'], ini['temporal_completeness']['system'], id)

	return ds, id

"""
Get information from QAR/ini result
"""
def parse_ini(ff):

	# There is empty
	res = {}
	metadata = {}

	if os.path.exists(ff):
		ini = configparser.ConfigParser()
		ini.read(ff)
		res = dict.fromkeys(tests)
		metadata = {}

		for s in ini.sections():
			if 'result' in s:
				res[s.replace('_result', '')] = ok_nok(ini[s]['res'])	
				metadata[s.replace('_result', '')] = '{}: {}'.format(ini[s]['res'], ini[s]['msg'])

	return res, metadata

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

	# Get all ini files
	script_dir = os.path.dirname(os.path.realpath(__file__))
	conf_dir = os.path.join(script_dir, 'conf')
	
	# Results dir
	res_dir = '/shared/QAR_dqc/results'
	os.chdir(res_dir)
	df = pd.DataFrame(columns=['ds', 'id', 'test', 'res'])
	
	for f in os.listdir(conf_dir):
		ds, var = get_metadata(os.path.join(conf_dir, f))
		ini_res = os.path.join(res_dir, f)
		rs, mdata = parse_ini(ini_res)
		completed = 0

		# Result available
		if rs != {}:
			tt = 0
			for k,v in rs.items():
				if v is None:
					v = 0
					mess = 'Data unavailable'
				else:
					mess = mdata[k]

				df=df.append({'ds': ds, 'id': var, 'test': k, 'res': v, 'msg': mess }, ignore_index=True)
				tt+=1

			# All checks passed
			completed = (len(tests) == tt)
		else:			
			for c in tests:		
				df=df.append({'ds': ds, 'id': var, 'test': c, 'res': 0, 'msg': 'Data unavailable'},
					ignore_index=True)

		# Add completed status
		df=df.append({'ds': ds, 'id': var, 'test': 'all', 'res': completed, 'msg': ''},
						ignore_index=True)

	# Save plot (Seasonal)
	ws = 800
	sumS = plot_summary('Seasonal', df, ws)
	ps1 = table_plot(df, 'seasonal-monthly-pressure-levels', ws)
	ps2 = table_plot(df, 'seasonal-monthly-single-levels', ws)
	ps3 = table_plot(df, 'seasonal-postprocessed-pressure-levels', ws)
	ps4 = table_plot(df, 'seasonal-postprocessed-single-levels', ws)
	ps5 = table_plot(df, 'seasonal-original-pressure-levels', ws)
	ps6 = table_plot(df, 'seasonal-original-single-levels', ws)
	output_file(os.path.join(res_dir, 'qar_status_seasonal.html'), title='Seasonal')
	save(column(sumS, ps1, ps2, ps3, ps4, ps5, ps6))

	# Save plot (Reanalysis)
	wr = 700
	sumR = plot_summary('Reanalysis', df, wr)
	pr1 = table_plot(df, 'reanalysis-era5-pressure-levels-monthly-means', wr)
	pr2 = table_plot(df, 'reanalysis-era5-single-levels-monthly-means', wr)
	pr3 = table_plot(df, 'reanalysis-era5-pressure-levels', wr)
	pr4 = table_plot(df, 'reanalysis-era5-single-levels', wr)
	pr5 = table_plot(df, 'reanalysis-uerra-europe-pressure-levels', wr)
	pr6 = table_plot(df, 'reanalysis-uerra-europe-single-levels', wr)
	output_file(os.path.join(res_dir, 'qar_status_reanalysis.html'), title='Reanalysis')
	save(column(sumR, pr1, pr2, pr3, pr4, pr5, pr6))

if __name__ == "__main__":
    main()