#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import print_function
from netCDF4 import Dataset as nc
import numpy as np
from types import FunctionType
import os.path
import sys
import subprocess
import shlex


# variables not to be checked
NO_CHECK_VARS = {
        'longitude': (),
        'latitude':  (),
        'lon_bnds':  (),
        'lat_bnds':  (),
        'lon':       (),
        'lat':       (),
        'level':     (),
        'time':      ('timesteps_number',),
}

VALID_RANGES = {
        'od550_dust'   : { 'min': 0, 'max': 20. },
        'sconc_dust'   : { 'min': 0, 'max': .00005 },
        'dust_load'    : { 'min': 0, },
        'dust_ext_sfc' : { 'min': 0, },
        'dust_depd'    : { 'min': 0, },
        'dust_depw'    : { 'min': 0, },
}

# methods with relative variables no to be evaluated
NO_METHOD_VAR = {
    'a_check_contains_nan'       : (),
    # 'b_check_contains_negatives' : (),
    'c_check_contains_masked'    : (),
    'd_check_contains_zeros'     : (),
}

MSG_NANS = "NaNs found! Exit."
MSG_NONANS = "NO NaNs found."
# MSG_NEGATIVES = "Negatives found! Exit."
# MSG_NONEGATIVES = "NO Negatives found."
MSG_MASKED = "Timestep {} with all MASKED found! Exit."
MSG_MISSING = "Timestep {} with all MISSING found! Exit."
MSG_NOMASKED = "NO timesteps with all MASKED or MISSING found."
MSG_ZEROS = "Timestep {} with all ZEROS found! Exit."
MSG_NOZEROS = "NO timesteps with all ZEROS found."
MSG_NO_TIMESTEPS = "Timesteps are {}, should be >= 21! Exit."
MSG_OK_TIMESTEPS = "Timesteps number {}, OK"


class NCChecker(object):
    """ Main class for NetCDF Checker """

    def __init__(self, filename):
        self.filename = filename
        self.nc_file = None
        self.methods = sorted([x for x, y in self.__class__.__dict__.items() \
                               if isinstance(y, FunctionType) and 'check' in x])

    def a_check_contains_nan(self, var):
        """ Check if NaNs are present """
        if np.isnan(var[:]).any():
            print(MSG_NANS)
            return 1

        print(MSG_NONANS)
        return 0

#     def b_check_contains_negatives(self, var):
#         """ Check if negatives are present """
#         val = var[:]
#         mval = (hasattr(var, 'missing_value') and \
#                 getattr(var, 'missing_value')) or \
#                (hasattr(var, '_FillValue') and \
#                 getattr(var, '_FillValue', None))
# 
#         if isinstance(val, np.ma.MaskedArray):
#             if val[(~val.mask) & (val != mval) & (val < 0)].size != 0:
#                 print(MSG_NEGATIVES)
#                 return 1
#         elif val[(val != mval) & (val < 0)].size != 0:
#             print(MSG_NEGATIVES)
#             return 1
# 
#         print(MSG_NONEGATIVES)
#         return 0

    def c_check_contains_masked(self, var):
        """ Check if at least one timestep has all masked, EXCEPT the first one """
        val = var[:]
        mval = (hasattr(var, 'missing_value') and \
                getattr(var, 'missing_value')) or \
               (hasattr(var, '_FillValue') and \
                getattr(var, '_FillValue', None))

        for tstep in range(1, val.shape[0]):
            tval = val[tstep]
            if isinstance(tval, np.ma.MaskedArray):
                if tval[~tval.mask].size == 0:
                    print(MSG_MASKED.format(tstep))
                    return 1
            elif tval[(tval != mval)].size == 0:
                print(MSG_MISSING.format(tstep))
                return 1

        print(MSG_NOMASKED)
        return 0

    def d_check_contains_zeros(self, var):
        """ Check if at least one timestep has all zeros, EXCEPT the first one """
        val = var[:]
        mval = (hasattr(var, 'missing_value') and \
                getattr(var, 'missing_value')) or \
               (hasattr(var, '_FillValue') and \
                getattr(var, '_FillValue', None))

        for tstep in range(1, val.shape[0]):
            tval = val[tstep]
            if isinstance(tval, np.ma.MaskedArray):
                if tval[(~tval.mask) & (tval != mval) & (tval != 0)].size == 0:
                    print(MSG_ZEROS.format(tstep))
                    return 1
            elif tval[(tval != mval) & (tval != 0)].size == 0:
                print(MSG_ZEROS.format(tstep))
                return 1

        print(MSG_NOZEROS)
        return 0

    def timesteps_number(self, var):
        """ Check if timesteps are less than 21 """
        if var.size < 21:
            print(MSG_NO_TIMESTEPS.format(var.size))
            return 1
        
        print(MSG_OK_TIMESTEPS.format(var.size))
        return 0

    def var_proc(self, varname):
        """ Processing variable """
        var = self.nc_file.variables[varname]
        if varname in NO_CHECK_VARS:
            print("Checking AUX variable '{}'".format(varname))
            func_list = NO_CHECK_VARS[varname]
            if func_list:
                for func in func_list:
                    print("* STEP 1: {} ... ".format(func), end='')
                    if getattr(self, func)(var) != 0:
                        return 1

        else:
            print("Checking MAIN variable '{}'".format(varname))
            for method in self.methods:
                if varname in NO_METHOD_VAR[method]:
                    continue
                print("* STEP 2: {} ... ".format(method), end='')
                if getattr(self, method)(var) != 0:
                    return 1

        return 0

    def run(self):
        """ Execute the checker """
        self.nc_file = nc(self.filename)
        varnames = [v for v in self.nc_file.variables]
        # empty files
        if not varnames:
            return 1
        # loop on variables
        for varname in varnames:
            if varname not in NO_CHECK_VARS:
                var = self.nc_file.variables[varname]
                mvalname = hasattr(var, 'missing_value') and 'missing_value' or '_FillValue'
                VALID_RANGES[varname.lower()]['mvalname'] = mvalname

            if self.var_proc(varname) != 0:
                self.nc_file.close()
                return 1

        self.nc_file.close()
        # set valid ranges
        # ncap2 -O -v -s 'where(OD550_DUST < 0) OD550_DUST=OD550_DUST@_FillValue; where(OD550_DUST > 0.9) OD550_DUST=OD550_DUST@_FillValue;' 20220831_3H_MEDIAN.nc 20220831_3H_MEDIAN.nc
        cmd = ""
        for varname in varnames:
            if varname in NO_CHECK_VARS:
                continue
            ranges = VALID_RANGES[varname.lower()]
            mvalname = ranges['mvalname']
            if 'min' in ranges:
                vmin = ranges['min']
                cmd += "where({varname} < {vmin}) {varname}={varname}@{mvalname};".format(varname=varname, vmin=vmin, mvalname=mvalname)
            if 'max' in ranges:
                vmax = ranges['max']
                cmd += "where({varname} > {vmax}) {varname}={varname}@{mvalname};".format(varname=varname, vmax=vmax, mvalname=mvalname)
        if cmd:
            cmd = shlex.split("ncap2 -O -v -s '{cmd}' {fname} {fname}".format(cmd=cmd, fname=self.filename))
            print("VALID RANGES COMMAND:", cmd)
            try:
                subprocess.run(cmd)
            except Exception as err:
                print('ERROR', str(err))
                return 1
        return 0


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: {} [filename.nc]".format(sys.argv[0]))
        sys.exit(1)
    if not os.path.isfile(sys.argv[1]):
        print("Error: filename {} doesn't exists or is not a regular file.".format(sys.argv[1]))
        sys.exit(1)

    NC_VALID = NCChecker(sys.argv[1]).run()
    print("RET: {}".format(NC_VALID))
    sys.exit(NC_VALID)
