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

def generate_ini(rootd, pattern, variable_e, datatype, variable, res, gt, ts, origin, system, dataset, is_s, dr, tg, com, lsm):

	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

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

[file_format]

[standard_compliance_grib]

[standard_compliance_netcdf]

[spatial_completeness]
land_sea_mask = {lsm}

[temporal_completeness]
is_seasonal = {is_s}

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

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

[valid_ranges]
valid_min =
valid_max =	

'''.format(
	p=pattern, ve=variable_e, dt=datatype, v=variable, is_s=is_s, dr=dr, gt=gt, com=com,
	o=origin, s=system, ds=dataset, ts=ts, r=rootd, res=res, tg=tg, lsm=lsm)

	return ini

"""
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, ve, mask_variables, daily, uerra):
	root = 'reanalysis'

	# 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:
		ts = 6

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

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

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

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

	# Update resolution if ocean/wave variable (reanalysis)
	rres = get_resolution(v, mask_variables)
	lsm = 'no'
	if rres != -1:
		lsm = 'yes'
		res = rres

	# Unique identifier
	qar_name = d + '-' + v + '.ini'
	if len(o):
		qar_name = d + '-' + o + '-' + v + '.ini'

	return qar_name, generate_ini(root, patt, ve, dt, v, res, gt, ts, o, '', d, 'no', dr, tg, com, lsm)

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

	# Time resolution + Data type
	com = 'yes'
	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'

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

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

	with open(jsonf, 'r') as f:
		dd=json.load(f)
		
		for dl in dd:
			for v in dl['variables']:				
				# General
				d = dl['dataset']
				ve = codes_var[dl['dataset']]['variables'][v]	

				# SEASONAL
				if 'originating_centre' in dl['form']:
					qar_name, ini = init_params_seasonal(dl, d, v, ve, mask_variables, temp_resolution, 'daily' in jsonf)
						
				# REANALYSIS
				else:
					qar_name, ini = init_params_reanalysis(dl, d, v, ve, mask_variables, 'daily' in jsonf, 'uerra' in jsonf)

				# Next
				qar_names.append(qar_name)
				inis.append(ini)

	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/datasets.json"

	# Output directory
	out_dir="/tmp/conf_generate"
	if not os.path.exists(out_dir):
		os.makedirs(out_dir)

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

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

		# Get file names
		qar_names = get_qar_names(json_conf, codes_var, mask_variables, temp_resolution)
		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:
				print('Writting ' + qar_out)
				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()
