#!/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
import numpy as np

from dqc_wrapper.utils.utils_xarray import *

"""
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.grib_file = 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

	"""
	Get GRIB first lat/lon values (useful when reading a mask)
	"""
	def get_grib_first_values(self):

		# Open file and rewind to scan all messages.
		f = open(self.grib_file, "rb")
		f.seek(0)

		while True:
			gid = ecc.codes_grib_new_from_file(f)
			if gid is None:
				break # EOF

			# Get data array
			data_arr = ecc.codes_get_array(gid, 'values')  # GRIB key 'values' array
			data_arr[data_arr == ecc.codes_get(gid, 'missingValue')] = np.nan
			data_arr = data_arr[~np.isnan(data_arr)]

			return data_arr

			# This is critical to avoid memory 'leak'
			ecc.codes_release(gid)

		return None

	"""
	Get GRIB all data values (useful for the valid ranges)
	"""
	def get_grib_all_values(self, subset=None):

		# Open file and rewind to scan all messages.
		f = open(self.grib_file, "rb")
		f.seek(0)

		# Concatenate all values in grib file
		vals = None
		while True:
			gid = ecc.codes_grib_new_from_file(f)
			if gid is None:
				break # EOF

			# Get data array
			if vals is None:
				np.concatenate(vals, ecc.codes_get_array(gid, 'values'))  # GRIB key 'values' array)
			else:
				vals = ecc.codes_get_array(gid, 'values')  # GRIB key 'values' array

			# This is critical to avoid memory 'leak'
			ecc.codes_release(gid)

		return None

	"""
	Get GRIB - Hadley region histogram analysis
	"""
	def get_grib_hadley_regions(self, init, hbins, df_hist, subset=None):

		# Open file and rewind to scan all messages.
		f = open(self.grib_file, "rb")
		f.seek(0)

		# Go through all values in grib file
		count = 0

		while True:
			gid = ecc.codes_grib_new_from_file(f)
			if gid is None:
				break # EOF

			# Get data array + applying missing values
			data_arr = ecc.codes_get_array(gid, 'values')  # GRIB key 'values' array
			data_arr[data_arr == ecc.codes_get(gid, 'missingValue')] = np.nan

			# 6 Hadley cells histograms (more a less)
			beg = 0
			lat_step = int(ecc.codes_get(gid, 'numberOfDataPoints') / 6)
			end = lat_step

			for cell in ['N_polar', 'N_midlat', 'N_hadley', 'S_hadley', 'S_midlat', 'S_polar']:
				# Compute region histogram (only valid values)
				slice = data_arr[beg:end]
				slice = slice[~np.isnan(slice)]
				x, b = np.histogram(slice, bins=hbins)

				# Store into pandas/csv: Initialization of the 6 cells vs sum
				try:
					if count < 6:
						xsum = x
						minv = min(slice)
						maxv = max(slice)
						if init == 0:
							df_hist.loc[cell] = x
					else:
						xsum = np.add(xsum, x)
						minv = min(minv, min(slice))
						maxv = max(maxv, max(slice))
						df_hist.loc[cell] += x
				except:
					pass # Empty region (e.g - all NaN - snowfall)

				# Next Hadley region
				beg = end
				end += lat_step
				count += 1

			# This is critical to avoid memory 'leak'
			ecc.codes_release(gid)

		return xsum, minv, maxv

	"""
	Get GRIB missing values (read message by message and compare with the mask if any)
	"""
	def get_grib_mv(self, mask_arrays):

		# Open file and rewind to scan all messages.
		f = open(self.grib_file, "rb")
		f.seek(0)

		# Counters (missing values [mv] and unnecessary values [uv])
		mv = 0
		uv = 0

		while True:
			gid = ecc.codes_grib_new_from_file(f)
			if gid is None:
				break # EOF

			# Get data array
			data_arr = ecc.codes_get_array(gid, 'values')  # GRIB key 'values' array

			# Check 1: Bitmap present
			bitmapPresent = ecc.codes_get(gid, 'bitmapPresent')
			if bitmapPresent:
				# Get the bitmap array which contains 0s and 1s
				bitmap = ecc.codes_get_array(gid, 'bitmap', int)
				# Do some sanity checking
				assert len(bitmap) == ecc.codes_get_size(gid, 'values')
				assert len(bitmap) == ecc.codes_get(gid, 'numberOfDataPoints')

			# Check 2: Apply missing values
			data_arr[data_arr == ecc.codes_get(gid, 'missingValue')] = np.nan

			# Count and compare with mask
			try:
				# If there is a mask we need to validate
				if mask_arrays != {}:
					mask_arr = mask_arrays[len(data_arr)]
					mvd, uvd = count_mv_uv(mask_arr, data_arr.reshape(mask_arr.shape))

				# Otherwise we just count NaN
				else:
					mvd, uvd = count_mv_uv(None, data_arr)

				mv += mvd
				uv += uvd
			except:
				raise Exception('Resolution miss-match, please check the mask you provided')

			# This is critical to avoid memory 'leak'
			ecc.codes_release(gid)

		return mv, uv
