#!/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 logging
import eccodes as ecc

from dqc_wrapper.utils.utils_xarray import *

"""
Function to load GRIB metadata from JSON file
- grib1.json and grib2.json contain formatted information from ECCODES
"""
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)

"""
Class: Wrapper with utils around the GRIB1 and GRIB2 file formats
 Uses eccodes (in a memory efficient way)
"""
class GribFileCheck:

	def __init__(self, grib_file, dt=[]):
		self.grib_file = grib_file
		self.version = self.get_grib_version()
		# Empty vector = all datatypes checked = all messages checked
		self.datatypes = dt

	"""
	Get GRIB version (1 or 2)
	"""
	def get_grib_version(self):

		# Open file and rewind to scan all messages.
		f = open(self.grib_file, "rb")
		f.seek(0)
		version = -1
		while True:
			gid = ecc.codes_grib_new_from_file(f)
			if gid is None:
				break  # EOF

			# Get data array
			version = ecc.codes_get(gid, 'GRIBEditionNumber')

			# This is critical to avoid memory 'leak'
			ecc.codes_release(gid)
			break # First message is enough

		return version

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

		# Memory greedy since we are not reading data values
		err_msgs = []
		ecc_data = ecc.GribFile(self.grib_file)
		msg = ecc.GribMessage(ecc_data)

		for ckey, cval in check_keys.items():
			for i in range(len(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 first lat/lon values useful when:
	 reading a mask (test02)
	 getting a sample for histogram init (test06)
	"""
	def get_grib_first_values(self, without_Nan=True):

		# 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

			# Filter data type (if provided, else continue)
			if len(self.datatypes) and not(ecc.codes_get(gid, 'dataType') in self.datatypes):
				ecc.codes_release(gid)
				continue

			# 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

			# Filter out NaN
			if without_Nan:
				data_arr = data_arr[~np.isnan(data_arr)]

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

			return data_arr

		# Reached EOF without valid data values
		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

			# Filter data type (if provided, else continue)
			if len(self.datatypes) and not(ecc.codes_get(gid, 'dataType') in self.datatypes):
				ecc.codes_release(gid)
				continue

			# 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  # SKIP 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, is_land):

		# 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')
			dataType = ecc.codes_get(gid, 'dataType')
			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
			# If there is no mask or it was not found
			sel_mask = None

			# If there is a mask we need to validate/reshape
			if mask_arrays != {}:
				try:
					nbp = len(data_arr)
					sel_mask = mask_arrays[nbp]
				except:
					# Mask not found
					raise Exception('Could not find an appropiate mask for (dt={} + nbp={}) with the available {}'.format(dataType, nbp, mask_arrays.keys()))

			# Count missing and unnecessary values
			mvd, uvd = count_mv_uv(sel_mask, data_arr, is_land)
			mv += mvd
			uv += uvd

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

		return mv, uv
