#!/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, ts, origin, system, dataset, is_s, mf, mv, dr, tg):
	input_dir = ''.format(ds=dataset, v=variable)

	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

[dataset]
variable = {ve}
datatype = {dt}
cds_dataset = {ds}
cds_variable = {v}

[file_format]

[standard_compliance]

[spatial_completeness]
mask_file = {mf}
mask_var = {mv}

[temporal_completeness]
origin = {o}
system = {s}
is_seasonal = {is_s}

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

[temporal_consistency]
time_step = {t}
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,
	o=origin, s=system, ds=dataset, t=ts, r=rootd, res=res, mf=mf, mv=mv, tg=tg)

	return ini

"""
Get mask file (if needed)
"""
def get_mask_file(root, var, ds, o, s, mask_variables):

	m = ''
	mv = ''
	res = -1

	for v in mask_variables:

		if var in v:
			res, var = v.split(',')

			# SEASONAL: Only original resolution has a variable land_sea_mask
			if root == 'seasonal' and 'original' in ds:
				m='/shared/cds_downloads/masks/{d}-{o}-{s}-land_sea_mask.grib'.format(d=ds, o=o, s=s)
				mv='lsm'

			# REANALYSIS: Only single-levels
			if root == 'reanalysis' and 'single-levels' in ds:
				m='/shared/cds_downloads/masks/{d}-{r}-land_sea_mask.grib'.format(d=ds, r=res)
				mv='lsm'

	return m, mv, res

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

	# Time resolution
	if daily:
		tg = 'day'
	else:
		tg = 'month'

	# Grid resolution (two resolutions allowed) - 0.25 and 0.5
	res = 0.25
	dr = 'yes'

	# File pattern
	patt = '{}-*.grib'.format(v)

	# Product type
	dt = 'an, fc'

	# Mask file
	mf, mv, rres = get_mask_file(root, v, d, '', '', mask_variables)
	if mf != '' and mv != '':
		res = rres

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

	return qar_name, generate_ini(root, patt, ve, dt, v, res, 1, '', '', d, 'no', mf, mv, dr, tg)

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

	# Time resolution + Data type
	if daily:
		tg = 'day'
		is_s = 'daily'
		dt = 'fc, hc, em'
	else:
		tg = 'month'
		is_s = 'monthly'
		dt = 'fcmean, hcmean, em'

	# Grid resolution
	res = 1
	dr = 'no'

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

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

	# Only original has land_sea_mask
	mf, mv, rres = get_mask_file('seasonal', v, d, o, s, mask_variables)

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

	return qar_name, generate_ini('seasonal', patt, ve, dt, v, res, 1, o, s, d, is_s, mf, mv, dr, tg)

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

	with open(jsonf, 'r') as f:
		dd=json.load(f)
		
		for dl in dd:
			for v in dl['variables']:				

				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, 'daily' in jsonf)
						
				# REANALYSIS
				else:	
					qar_name, ini = init_params_reanalysis(dl, d, v, ve, mask_variables, 'daily' 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') as f:
		mask_variables = 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)
		for data in qar_names:
			q,i = data
			qar=os.path.join(qar_dir, q)
			qar_out=os.path.join(out_dir, q)
			
			if os.path.exists(qar_out):
				print('Skipping ' + q)				
			else:
				print('Writting ' + q)
				write_file(qar_out, i)

			if os.path.exists(qar):
				os.system('diff -b {} {}'.format(qar_out, qar))

if __name__ == "__main__":
    main()
