from __future__ import (absolute_import, division, print_function)

from .util import (get_iterable, is_standard_wrf_var, extract_vars, viewkeys,
                   get_id)
from .g_cape import (get_2dcape, get_3dcape, get_cape2d_only,
                     get_cin2d_only, get_lcl, get_lfc, get_3dcape_only,
                     get_3dcin_only)
from .g_ctt import get_ctt
from .g_dbz import get_dbz, get_max_dbz
from .g_dewpoint import get_dp, get_dp_2m
from .g_geoht import (get_geopt, get_height, get_stag_geopt, get_stag_height,
                      get_height_agl)
from .g_helicity import get_srh, get_uh
from .g_latlon import get_lat, get_lon
from .g_omega import get_omega
from .g_pressure import get_pressure, get_pressure_hpa
from .g_pw import get_pw
from .g_rh import get_rh, get_rh_2m
from .g_slp import get_slp
from .g_temp import (get_tc, get_eth, get_temp, get_theta, get_tk, get_tv,
                     get_tw)
from .g_terrain import get_terrain
from .g_uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir,
                      get_uvmet_wspd_wdir, get_uvmet_wspd, get_uvmet_wdir,
                      get_uvmet10_wspd, get_uvmet10_wdir)
from .g_vorticity import get_avo, get_pvo
from .g_wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10,
                     get_u_destag, get_v_destag, get_w_destag,
                     get_destag_wspd, get_destag_wdir, get_destag_wspd10,
                     get_destag_wdir10)
from .g_times import get_times, get_xtimes
from .g_cloudfrac import (get_cloudfrac, get_low_cloudfrac, get_mid_cloudfrac,
                          get_high_cloudfrac)


# func is the function to call.  kargs are required arguments that should
# not be altered by the user
_FUNC_MAP = {"cape2d": get_2dcape,
             "cape3d": get_3dcape,
             "dbz": get_dbz,
             "maxdbz": get_max_dbz,
             "dp": get_dp,
             "dp2m": get_dp_2m,
             "height": get_height,
             "geopt": get_geopt,
             "srh": get_srh,
             "uhel": get_uh,
             "omega": get_omega,
             "pw": get_pw,
             "rh": get_rh,
             "rh2m": get_rh_2m,
             "slp": get_slp,
             "theta": get_theta,
             "temp": get_temp,
             "tk": get_tk,
             "tc": get_tc,
             "theta_e": get_eth,
             "tv": get_tv,
             "twb": get_tw,
             "terrain": get_terrain,
             "times": get_times,
             "xtimes": get_xtimes,
             "uvmet": get_uvmet,
             "uvmet10": get_uvmet10,
             "avo": get_avo,
             "pvo": get_pvo,
             "ua": get_u_destag,
             "va": get_v_destag,
             "wa": get_w_destag,
             "lat": get_lat,
             "lon": get_lon,
             "pressure": get_pressure_hpa,
             "pres": get_pressure,
             "wspd_wdir": get_destag_wspd_wdir,
             "wspd_wdir10": get_destag_wspd_wdir10,
             "uvmet_wspd_wdir": get_uvmet_wspd_wdir,
             "uvmet10_wspd_wdir": get_uvmet10_wspd_wdir,
             "ctt": get_ctt,
             "cloudfrac": get_cloudfrac,
             "geopt_stag": get_stag_geopt,
             "zstag": get_stag_height,
             "height_agl": get_height_agl,
             # Diagnostics below are extracted from multi-product diagnostics
             "cape2d_only": get_cape2d_only,
             "cin2d_only": get_cin2d_only,
             "lcl": get_lcl,
             "lfc": get_lfc,
             "cape3d_only": get_3dcape_only,
             "cin3d_only": get_3dcin_only,
             "uvmet_wspd": get_uvmet_wspd,
             "uvmet_wdir": get_uvmet_wdir,
             "uvmet10_wspd": get_uvmet10_wspd,
             "uvmet10_wdir": get_uvmet10_wdir,
             "wspd": get_destag_wspd,
             "wdir": get_destag_wdir,
             "wspd10": get_destag_wspd10,
             "wdir10": get_destag_wdir10,
             "low_cloudfrac": get_low_cloudfrac,
             "mid_cloudfrac": get_mid_cloudfrac,
             "high_cloudfrac": get_high_cloudfrac
             }

_VALID_KARGS = {"cape2d": ["missing"],
                "cape3d": ["missing"],
                "dbz": ["do_variant", "do_liqskin"],
                "maxdbz": ["do_variant", "do_liqskin"],
                "dp": ["units"],
                "dp2m": ["units"],
                "height": ["msl", "units"],
                "geopt": [],
                "srh": ["top"],
                "uhel": ["bottom", "top"],
                "omega": [],
                "pw": [],
                "rh": [],
                "rh2m": [],
                "slp": ["units"],
                "temp": ["units"],
                "tk": [],
                "tc": [],
                "theta": ["units"],
                "theta_e": ["units"],
                "tv": ["units"],
                "twb": ["units"],
                "terrain": ["units"],
                "times": [],
                "xtimes": [],
                "uvmet": ["units"],
                "uvmet10": ["units"],
                "avo": [],
                "pvo": [],
                "ua": ["units"],
                "va": ["units"],
                "wa": ["units"],
                "lat": [],
                "lon": [],
                "pres": ["units"],
                "pressure": ["units"],
                "wspd_wdir": ["units"],
                "wspd_wdir10": ["units"],
                "uvmet_wspd_wdir": ["units"],
                "uvmet10_wspd_wdir": ["units"],
                "ctt": ["fill_nocloud", "missing", "opt_thresh", "units"],
                "cloudfrac": ["vert_type", "low_thresh",
                              "mid_thresh", "high_thresh"],
                "geopt_stag": [],
                "zstag": ["msl", "units"],
                "height_agl": ["units"],
                "cape2d_only": ["missing"],
                "cin2d_only": ["missing"],
                "lcl": ["missing"],
                "lfc": ["missing"],
                "cape3d_only": ["missing"],
                "cin3d_only": ["missing"],
                "uvmet_wspd": ["units"],
                "uvmet_wdir": ["units"],
                "uvmet10_wspd": ["units"],
                "uvmet10_wdir": ["units"],
                "wspd": ["units"],
                "wdir": ["units"],
                "wspd10": ["units"],
                "wdir10": ["units"],
                "low_cloudfrac": ["vert_type", "low_thresh",
                                  "mid_thresh", "high_thresh"],
                "mid_cloudfrac": ["vert_type", "low_thresh",
                                  "mid_thresh", "high_thresh"],
                "high_cloudfrac": ["vert_type", "low_thresh",
                                   "mid_thresh", "high_thresh"],
                "default": []
                }

_ALIASES = {"cape_2d": "cape2d",
            "cape_3d": "cape3d",
            "eth": "theta_e",
            "mdbz": "maxdbz",
            "geopotential": "geopt",
            "helicity": "srh",
            "latitude": "lat",
            "longitude": "lon",
            "omg": "omega",
            "p": "pres",
            "rh2": "rh2m",
            "z": "height",
            "ter": "terrain",
            "updraft_helicity": "uhel",
            "td": "dp",
            "td2": "dp2m",
            "cfrac": "cloudfrac",
            "wspd_wdir_uvmet": "uvmet_wspd_wdir",
            "wspd_wdir_uvmet10": "uvmet10_wspd_wdir",
            "th": "theta",
            "low_cfrac": "low_cloudfrac",
            "mid_cfrac": "mid_cloudfrac",
            "high_cfrac": "high_cloudfrac",
            "wspd_uvmet": "uvmet_wspd",
            "wdir_uvmet": "uvmet_wdir",
            "wspd_uvmet10": "uvmet10_wspd",
            "wdir_uvmet10": "uvmet10_wdir",
            }


class ArgumentError(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg


def _undo_alias(alias):
    actual = _ALIASES.get(alias, None)
    if actual is None:
        return alias
    else:
        return actual


def _check_kargs(var, kargs):
    for arg in viewkeys(kargs):
        if arg not in _VALID_KARGS[var]:
            if var != "default":
                raise ValueError("'{}' is an invalid keyword "
                                 "argument for '{}'".format(arg, var))
            else:
                raise ValueError("'{}' is an invalid keyword "
                                 "argument".format(arg))


def getvar(wrfin, varname, timeidx=0,
           method="cat", squeeze=True, cache=None, meta=True,
           **kwargs):

    """Returns basic diagnostics from the WRF ARW model output.

    A table of all available diagnostics is below.

    .. include:: ../../_templates/product_table.txt

    Args:
        wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
            iterable): WRF-ARW NetCDF
            data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
            or an iterable sequence of the aforementioned types.

        varname (:obj:`str`) : The variable name.

        timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
            desired time index. This value can be a positive integer,
            negative integer, or
            :data:`wrf.ALL_TIMES` (an alias for None) to return
            all times in the file or sequence. The default is 0.

        method (:obj:`str`, optional): The aggregation method to use for
            sequences.  Must be either 'cat' or 'join'.
            'cat' combines the data along the Time dimension.
            'join' creates a new dimension for the file index.
            The default is 'cat'.

        squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
            with a size of 1 from being automatically removed from the shape
            of the output. Default is True.

        cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
            that can be used to supply pre-extracted NetCDF variables to the
            computational routines.  It is primarily used for internal
            purposes, but can also be used to improve performance by
            eliminating the need to repeatedly extract the same variables
            used in multiple diagnostics calculations, particularly when using
            large sequences of files.
            Default is None.

        meta (:obj:`bool`, optional): Set to False to disable metadata and
            return :class:`numpy.ndarray` instead of
            :class:`xarray.DataArray`.  Default is True.

        **kwargs: Optional keyword arguments for certain diagnostics.
            See table above.


    Returns:

        :class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
        enabled and the *meta* parameter is True, then the result will be a
        :class:`xarray.DataArray` object.  Otherwise, the result will be a
        :class:`numpy.ndarray` object with no metadata.


    Raises:
        :class:`ValueError`: Raised when an invalid diagnostic type or
            keyword argument is passed to the routine.
        :class:`FortranError`: Raised when a problem occurs during a Fortran
            calculation.

    See Also:

        :class:`numpy.ndarray`, :class:`xarray.DataArray`


    Examples:
        Using netCDF4

        .. code-block:: python

            from netCDF4 import Dataset
            from wrf import getvar

            wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00")
            slp = getvar(wrfnc, "slp")

        Using PyNIO

        .. code-block:: python

            from Nio import open_file
            from wrf import getvar

            wrfnc = open_file("wrfout_d02_2010-06-13_21:00:00"+".nc", "r")
            slp = getvar(wrfnc, "slp")

        Using Iterables:

        .. code-block:: python

            import os
            from netCDF4 import Dataset
            from wrf import getvar

            filedir = "/path/to/wrf/files"
            wrfin = [Dataset(f) for f in os.listdir(filedir)
                    if f.startswith("wrfout_d02_")]

            uvmet = getvar(wrfin, "uvmet", timeidx=3, units="kt")


    """

    _key = get_id(wrfin)

    wrfin = get_iterable(wrfin)

    if is_standard_wrf_var(wrfin, varname) and varname != "Times":
        _check_kargs("default", kwargs)
        return extract_vars(wrfin, timeidx, varname,
                            method, squeeze, cache, meta, _key)[varname]
    elif varname == "Times":
        varname = "times"  # Diverting to the get_times routine

    actual_var = _undo_alias(varname)
    if actual_var not in _VALID_KARGS:
        raise ValueError("'{}' is not a valid variable name".format(varname))

    _check_kargs(actual_var, kwargs)

    return _FUNC_MAP[actual_var](wrfin, timeidx, method, squeeze, cache,
                                 meta, _key, **kwargs)
