#!/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 list of QAR ini configuration files
---
- Read from download lists
- Generate ini file to be used by the
"""

import os
import json
import glob
import csv

from dqc_wrapper.utils.utils_forms import get_forms_dates


def generate_ini(rootd, pattern, variable_e, datatype, variable, units, res, gt, gmn,
				 ts, origin, system, dataset, is_s, dr, tg, com, lsm, lnd, vmin, vmax, cmsname):

	ini='''[general]
input = /shared/cds_downloads/{r}/{ds}/{v}
fpattern = {p}
log_dir = /shared/QAR_dqc/logs 
res_dir = /shared/QAR_dqc/results
forms_dir = /shared/cds-forms-c3s
masks_dir = /shared/cds_downloads/masks
cms_name = {cms}

[dataset]
variable = {ve}
datatype = {dt}
cds_dataset = {ds}
cds_variable = {v}
cds_units = {un}
complete = {com}
origin = {o}
system = {s}

[file_format]

[standard_compliance_grib]

[standard_compliance_netcdf]

[spatial_completeness]
land_sea_mask = {lsm}
land_variable = {lnd}

[temporal_completeness]
is_seasonal = {is_s}

[spatial_consistency]
grid_interval = {res}
grid_type = {gt}
grid_mapping_name = {gmn}
dual_resolution = {dr}

[temporal_consistency]
time_step = {ts}
time_granularity = {tg}

[valid_ranges]
valid_min = {vmin}
valid_max = {vmax}
'''.format(
	p=pattern, ve=variable_e, dt=datatype, v=variable, is_s=is_s, dr=dr, gt=gt, com=com, gmn=gmn, un=units, cms=cmsname,
	o=origin, s=system, ds=dataset, ts=ts, r=rootd, res=res, tg=tg, lsm=lsm, lnd=lnd, vmin=vmin, vmax=vmax)

	return ini

"""
Get tolerance limits
"""
def get_tolerance_limits(tolerance_file):
	# Read and skip comments
	var_min_max = dict()
	with open(tolerance_file, 'r') as fp:
		rdr = csv.reader(filter(lambda row: row[0]!='#', fp))
		for row in rdr:
			if len(row):
				var_min_max[row[0]] = row[1:3]

	fp.close()
	return var_min_max

"""
Get temporal resolution (if needed) + land_sea_mask
"""
def get_temporal_res(var, temp_resolution):
	var = var.replace('_anomaly', '') # postprocessed
	for v in temp_resolution:
		variable, s, g, l = v.split(',')
		if variable.strip() == var.strip():
			return s, g, l

	return None, None, 'no'

"""
Get mask file (if needed)
"""
def get_resolution(var, mask_variables):
	for v in mask_variables:
		rr, vv = v.split(',')
		if vv.rstrip() == var.rstrip():
			return rr
	return -1

"""
Init parameters for a reanalysis dataset
"""
def init_params_reanalysis(dl, d, v, vinfo, mask_variables, tolerance_limits, daily, uerra, land):
	root = 'reanalysis'

	# Get variable info
	ve = vinfo['shortName']
	u = vinfo['c3sunit']

	# Some units we know how to set limits
	if u == '%':
		vmin = 0.0
		vmax = 100.0
	elif u == '(0 - 1)':
		vmin = 0.0
		vmax = 1.0
	elif u.lower() == 'degrees':
		vmin = 0.0
		vmax = 360.0
	else:
		vmin = ''
		vmax = ''

	# Try to get tolerance limits (WMO or based on world records)
	try:
		vmin = float(tolerance_limits[v][0].strip())
	except:
		vmin = ''
	try:
		vmax = float(tolerance_limits[v][1].strip())
	except:
		vmax = ''

	# Model distinction
	try:
		o = dl['form']['origin']  # uerra has model distinction
	except:
		o = ''

	# Time resolution granularity (complete or not, split dates or not)
	if daily:
		ts = 3
		com = 'no'
		tg = 'hour'
	else:
		ts = 1
		com = 'yes'
		tg = 'month'
	if uerra:
		if 'mescan' in o:
			ts = 12
		else:
			ts = 6
	if land:
		ts = 1
		if daily:
			tg = 'day'
			com = 'no'

	# Grid resolution (two resolutions allowed) - 0.25 and 0.5
	gmn = ''
	if uerra:
		model = 'uerra'
		res = 11000
		if 'mescan' in o:
			res = 5500
		dr = 'no'
		gt = 'projection'
		gmn = 'lambert_conformal_conic'

		patt = '*.grib'
		if o != '':
			patt = '{}-*.grib'.format(o)
	else:
		model = 'era5'
		res = 0.25
		dr = 'yes'
		gt = 'lonlat'
		if land:
			model = 'era5_land'
			res = 0.1
			dr = 'no'
		patt = '{}-*.grib'.format(v)

	# Product type
	dt = 'an, fc, em'

	# Update resolution if ocean/wave variable (era5-reanalysis vs era5-land)
	rres = get_resolution(v, mask_variables)
	if land:
		lsm = 'yes'
		lnd = 'yes'
	else:
		lsm = 'no'
		lnd = 'no'
		if rres != -1:
			lsm = 'yes'
			res = rres
	if uerra:
		lsm = 'no'
		lnd = 'no'
		if 'snow' in v or 'albedo' in v:
			lsm = 'yes'
			lnd = 'yes'

	# Unique identifier for CMS
	qar_name = d + '-' + v + '.ini'
	cmsname = '{}-{}-{}'.format(d.replace('-', '_'), model, v.replace('-', '_'))
	if len(o):
		qar_name = d + '-' + o + '-' + v + '.ini'
		cmsname = '{}-{}-{}'.format(d.replace('-', '_'), o, v.replace('-', '_'))

	return qar_name, generate_ini(root, patt, ve, dt, v, u, res, gt, gmn, ts, o, '', d, 'no', dr, tg, com, lsm, lnd, vmin, vmax, cmsname)

"""
Init parameters for a seasonal dataset
"""
def init_params_seasonal(dl, d, v, vinfo, mask_variables, temp_resolution, tolerance_limits, daily):

	# Get variable info
	ve = vinfo['shortName']
	u = vinfo['c3sunit']

	# Try to get tolerance limits
	try:
		vmin = float(tolerance_limits[v][0].strip())
	except:
		vmin = ''
	try:
		vmax = float(tolerance_limits[v][1].strip())
	except:
		vmax = ''

	# Time resolution + Data type
	com = 'yes'
	lnd = 'no'
	if daily:
		com = 'no'
		is_s = 'daily'
		dt = 'fc, hc, em'
		if 'pressure' in d:
			ts, tg, lsm = '12', 'hour', 'no'
		else:
			ts, tg, lsm = get_temporal_res(v, temp_resolution)
	else:
		ts, tg = '1', 'month'
		is_s = 'monthly'
		dt = 'fcmean, hcmean, em'
		lsm = 'no'
		if not('pressure' in d):
			try:
				_, _, lsm = get_temporal_res(v, temp_resolution)
			except:
				print(d,v)
				exit()

	# Grid resolution
	res = 1
	dr = 'no'
	gt = 'lonlat'

	# Origin/System
	o = dl['form']['originating_centre']
	s = dl['form']['system']

	# File pattern
	patt = o + '-' + s + '*.grib'

	# Unique identifier
	qar_name = d + '-' + v + '-' + o + '-' + s + '.ini'

	# CMS name (seasonal)
	cmsname = '{}-{}-{}-{}'.format(d.replace('-','_'), o, s, v.replace('-','_'))

	return qar_name, generate_ini('seasonal', patt, ve, dt, v, u, res, gt, '', ts, o, s, d, is_s, dr, tg, com, lsm, lnd, vmin, vmax, cmsname)

"""
Generate a ini file (datachecker configuration)
for each of dataset/variable combination
"""
def get_qar_names(forms_dir, jsonf, codes_var, mask_variables, temp_resolution, tolerance_limits):
	qar_names = []
	inis = []

	with open(jsonf, 'r') as f:
		dd=json.load(f)
		
		for dl in dd:
			for v in dl['variables']:
				try:
					# General
					d = dl['dataset']
					ve = codes_var[d]['variables'][v]
				except:
					print('ERROR: dataset={} variable={} not found in cds_datasets.json'.format(d, v))
					continue

				# SEASONAL
				if 'originating_centre' in dl['form'] and 'system' in dl['form']:
					mandatory_dates = get_forms_dates(d, v, dl['form']['originating_centre'], dl['form']['system'], forms_dir)
					qar_name, ini = init_params_seasonal(dl, d, v, ve, mask_variables, temp_resolution, tolerance_limits, 'daily' in jsonf)
				# REANALYSIS
				else:
					# UERRA-SL / ORIGIN
					origin = ''
					try:
						origin = dl['form']['origin']
					except:
						pass

					mandatory_dates = get_forms_dates(d, v, origin, '', forms_dir)
					qar_name, ini = init_params_reanalysis(dl, d, v, ve, mask_variables, tolerance_limits,
														   'daily' in jsonf, 'uerra' in jsonf, 'land' in jsonf)

				# Next + skip empty datasets
				if len(mandatory_dates):
					qar_names.append(qar_name)
					inis.append(ini)
				else:
					print(qar_name)

	return zip(qar_names, inis)

def write_file(outf, string):
	with open(outf, 'w') as f:
		f.write(string)

def main():

	# Config paths
	conf_dir="../../dqc_downloader/conf"
	qar_dir="./conf"
	datasets = "../grib/cds_datasets.json"
	forms_dir = "/data/cds-forms-c3s"

	# Output directory
	out_dir="/tmp/conf_generate"
	if not os.path.exists(out_dir):
		os.makedirs(out_dir)
	else:
		os.system('rm /tmp/conf_generate/*')

	# Get codes
	codes_var = json.load(open(datasets,'r'))

	# Variables that require a land_sea_mask
	mask_variables = []
	with open('./mask_variables_reanalysis') as f:
		mask_variables = f.read().splitlines()

	# Variables that require a land_sea_mask
	temp_resolution = []
	with open('./temp_resolution_seasonal') as f:
		temp_resolution = f.read().splitlines()

	# Variables that have tolerance limits
	tolerance_limits = get_tolerance_limits('./tolerance_limits_wmo')

	# For downloaded datasets (daily + monthly)
	for json_conf in glob.glob(conf_dir+'/*ly*.json'):

		# Get file names
		qar_names = get_qar_names(forms_dir, json_conf, codes_var, mask_variables, temp_resolution, tolerance_limits)
		for data in qar_names:
			q,i = data
			qar=os.path.join(qar_dir, q)
			qar_out=os.path.join(out_dir, q)

			# Skip if file exists
			if os.path.exists(qar_out):
				print('Skipping ' + qar_out)
			else:
				write_file(qar_out, i)

			# Show differences
			if os.path.exists(qar):
				os.system('diff -b {} {}'.format(qar_out, qar))
			else:
				print('NEW: {}'.format(qar_out))

if __name__ == "__main__":
    main()
