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

# Available tests
from dqc_wrapper.utils.utils_files import *

tests = [
	'file_format',
	'standard_compliance_grib', 'standard_compliance_netcdf',
	'temporal_completeness', 'temporal_consistency',
	'spatial_completeness', 'spatial_consistency',
	'valid_ranges', 'download_status'
]

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

	# 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, 'DQC-4-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 (dataset level)
	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']

	# 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(
		       filt, completed, available),
	     	   x_range=tests, y_range=sorted(yy, reverse=True),
	           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

	# 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

"""
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.RawConfigParser()
	ini.read(ff)
	ds = ini['dataset']['cds_dataset']
	id = ini['dataset']['cds_variable']

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

	if 'uerra' in ds and 'single' in ds:
		id = '{}-{}'.format(ini['dataset']['origin'], id)

	return ds, id

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

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

	if os.path.exists(ff):
		ini = configparser.RawConfigParser()
		ini.read(ff)
		res = dict.fromkeys(tests)
		msg = dict.fromkeys(tests)

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

	return res, msg

"""
Generate download status
"""
def download_status(config_path):
	# Read configuration
	conf = CheckerConfig(config_path)
	conf.read_config(execute=False)

	# List files and get last date
	files = sorted(glob.glob(os.path.join(conf.path_files, conf.fpattern)))
	dates = sorted(get_dates_from_files(files))

	# Dates available (complete vs Thist - 2019)
	if len(files):
		if dates[-1] < datetime(2019, 6, 1, 0, 0, 0, 0):
			return 2, len(files), dates[-1].strftime('%Y%m%d')
		else:
			# More than 12 files seems reasonable (6 years daily, 1 year monthly)
			if len(files) >= 12:
				return 1, len(files), dates[-1].strftime('%Y%m%d')
			else:
				return 2, len(files), dates[-1].strftime('%Y%m%d')

	# Nothing available
	else:
		return 0, 0, 'not available'

"""
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, 'check_tests')
	
	# 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):
		# Parse ini
		ds, var = get_metadata(os.path.join(conf_dir, f))
		ini_conf = os.path.join(conf_dir, f)
		ini_res = os.path.join(res_dir, f)
		rs, mdata = parse_ini(ini_res)
		completed = 0

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

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

			# All checks passed
			completed = (len(tests) == tt)
		else:			
			for t in tests:
				if t != 'download_status':
					df=df.append({'ds': ds, 'id': var, 'test': t, '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)

		# Add downloaded status
		status, downloaded, lastdate = download_status(ini_conf)
		df=df.append({'ds': ds, 'id': var, 'test': 'download_status', 'res': status,
					  'msg': '{} downloaded, last date is {}'.format(downloaded, lastdate) },
						ignore_index=True)

	# Root dir (plots)
	root_dir = '/shared/QAR_dqc'

	# Save plot (Seasonal)
	ws = 800
	seasonal_path = os.path.join(root_dir, 'qar_status_seasonal.html')
	sumS = plot_summary('Seasonal', 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(seasonal_path, title='Seasonal')
	print('Written: {}'.format(seasonal_path))
	save(column(sumS, ps1, ps2, ps3, ps4, ps5, ps6))

	# Save plot (Reanalysis)
	wr = 700
	reanalysis_path = os.path.join(root_dir, 'qar_status_reanalysis.html')
	sumR = plot_summary('Reanalysis', 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-era5-land', wr)
	pr6 = table_plot(df, 'reanalysis-era5-land-monthly-means', wr)
	pr7 = table_plot(df, 'reanalysis-uerra-europe-pressure-levels', wr)
	pr8 = table_plot(df, 'reanalysis-uerra-europe-single-levels', wr)
	output_file(reanalysis_path, title='Reanalysis')
	print('Written: {}'.format(reanalysis_path))
	save(column(sumR, pr1, pr2, pr3, pr4, pr5, pr6, pr7, pr8))

	# Global statistics
	all_test = df.loc[df['test'] == 'all']
	completed = all_test.res.sum()
	available = all_test.res.count()
	print('{} of {} completed'.format(completed, available))

if __name__ == "__main__":
    main()