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

import os
import json
import subprocess
import eccodes as ecc

"""
Function to load GRIB metadata
"""
def load_grib_metadata(version=2):
	curr_dir = os.path.dirname(os.path.realpath(__file__))
	grib_file = os.path.join(curr_dir.replace('utils', 'grib/grib{}.json'.format(version)))
	with open(grib_file, 'r') as fp:
		return json.load(fp)

"""
Function used to quickly get a key from all grib messages
Return unique key values (ex: resolution, stepType, ..)
"""
def get_grib_key(grib_key, grib_file):

	# Get grib key
	cmd = "grib_get -p {k} {f}".format(k=grib_key, f=grib_file)
	proc = subprocess.Popen(cmd,
		shell=True,
	    stdout=subprocess.PIPE,
	    stderr=subprocess.STDOUT
	)

	# Read stdout
	lines = []
	while True:
		line = proc.stdout.readline()
		if not line:
			break
		lines.append(line.strip())

	return list(set(lines))

"""
Function used to get the maximum resolution within a grib file
Useful when opening dataset with xarray/cfgrib (key: numberOfDataPoints)
"""
def grib_get_max_resolution(grib_file):
	# Get available resolutions
	resolutions = get_grib_key('numberOfDataPoints', grib_file)

	# Get maximum resolution (if available)
	maxres = 0
	for r in resolutions:
		maxres = max(maxres, int(r))

	return [ maxres ]

"""
Function used to get all the resolutions in a grib file
"""
def grib_get_resolutions(grib_file):
	# Get available resolutions
	resolutions = get_grib_key('numberOfDataPoints', grib_file)

	# Get maximum resolution (if available)
	res = []
	for r in resolutions:
		try:
			res.append(int(r))
		except:
			pass # text messages skip

	return list(set(res))

"""
Function used to get all the resolutions in a grib file
"""
def grib_get_stepType(grib_file):
	# Get available resolutions
	resolutions = get_grib_key('stepType', grib_file)

	# Get all unique stepTypes (string)
	st = []
	for r in resolutions:
		ln = r.decode("utf-8")
		if not('error' in ln.lower()):
			st.append(ln)

	return list(set(st))

"""
Class: Wrapper with utils around the GRIB1 and GRIB2 file formats
---
- Uses eccodes
"""
class GribFileCheck:

	def __init__(self, grib_file):        
		self.ecc_data = ecc.GribFile(grib_file)
		self.version = self.get_grib_version()

	"""
	Get all messages stored on a grib file
	"""
	def get_message_list(self):
		# This funtion returns all messages within a GRIB file
		msgs = []

		for i in range(len(self.ecc_data)):
			msgs.append(ecc.GribMessage(self.ecc_data))
		
		return msgs

	"""
	This functions loops over a GRIB2 file and checks for certain key,value pairs
	"""
	def check_mandatory_keys(self, check_keys):

		err_msgs = []
		msg = ecc.GribMessage(self.ecc_data)
		for ckey,cval in check_keys.items():  
			for i in range(len(self.ecc_data)):
				try:
					if int(msg[ckey]) != int(cval.strip()):
						err_msgs.append('GRIB Key \'{}\' has value \'{}\' instead of value \'{}\''.format(ckey, msg[ckey], cval))
						break # no need to proceed
				except ValueError:
					pass

		# All messages had the appropiate keys
		return err_msgs

	"""
	Get GRIB version from the first message. Returns integer
	"""
	def get_grib_version(self):

		first_message = self.ecc_data.next()
		keys = first_message.keys()

		if 'GRIBEditionNumber' in keys:
			return first_message['GRIBEditionNumber']

		# Not a valid GRIB file
		return None
