from __future__ import (absolute_import, division, print_function)

from collections.abc import Iterable

import numpy as np

from .constants import Constants, ProjectionTypes
from .extension import _lltoxy, _xytoll
from .util import (extract_vars, extract_global_attrs,
                   either, is_moving_domain, iter_left_indexes,
                   is_mapping, is_multi_file)
from .py3compat import viewkeys, viewitems
from .projutils import dict_keys_to_upper


def _lat_varname(wrfin, stagger):
    """Return the latitude variable name for the specified stagger type.

    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.

        stagger (:obj:`str`): The staggered grid type which is one of the
            following:

                - 'm': Use the mass grid (default).
                - 'u': Use the same staggered grid as the u wind component,
                  which has a staggered west_east (x) dimension.
                - 'v': Use the same staggered grid as the v wind component,
                  which has a staggered south_north (y) dimension.

    Returns:

        :obj:`str`: The latitude variable name.

    """
    if stagger is None or stagger.lower() == "m":
        varname = either("XLAT", "XLAT_M")(wrfin)
    elif stagger.lower() == "u" or stagger.lower() == "v":
        varname = "XLAT_{}".format(stagger.upper())
    else:
        raise ValueError("invalid 'stagger' value")

    return varname


def _lon_varname(wrfin, stagger):
    """Return the longitude variable name for the specified stagger type.

    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.

        stagger (:obj:`str`): The staggered grid type, which is one of the
            following:

                - 'm': Use the mass grid (default).
                - 'u': Use the same staggered grid as the u wind component,
                  which has a staggered west_east (x) dimension.
                - 'v': Use the same staggered grid as the v wind component,
                  which has a staggered south_north (y) dimension.

    Returns:

        :obj:`str`: The latitude variable name.

    """
    if stagger is None or stagger.lower() == "m":
        varname = either("XLONG", "XLONG_M")(wrfin)
    elif stagger.lower() == "u" or stagger.lower() == "v":
        varname = "XLONG_{}".format(stagger.upper())
    else:
        raise ValueError("invalid 'stagger' value")

    return varname


def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key):
    """Return the map projection parameters.

    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.

        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.

        stagger (:obj:`str`): The staggered grid type, which is one of the
            following:

                - 'm': Use the mass grid (default).
                - 'u': Use the same staggered grid as the u wind component,
                  which has a staggered west_east (x) dimension.
                - 'v': Use the same staggered grid as the v wind component,
                  which has a staggered south_north (y) dimension.

        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.

        _key (:obj:`int`, optional): A caching key. This is used for internal
            purposes only.  Default is None.

    Returns:


    """
    if timeidx is not None:
        if timeidx < 0:
            raise ValueError("'timeidx' must be greater than 0")

    attrs = extract_global_attrs(wrfin, attrs=("MAP_PROJ", "TRUELAT1",
                                               "TRUELAT2", "STAND_LON",
                                               "DX", "DY"))
    map_proj = attrs["MAP_PROJ"]
    truelat1 = attrs["TRUELAT1"]
    truelat2 = attrs["TRUELAT2"]
    stdlon = attrs["STAND_LON"]
    dx = attrs["DX"]
    dy = attrs["DY"]

    if map_proj == ProjectionTypes.LAT_LON:
        pole_attrs = extract_global_attrs(wrfin, attrs=("POLE_LAT",
                                                        "POLE_LON"))
        pole_lat = pole_attrs["POLE_LAT"]
        pole_lon = pole_attrs["POLE_LON"]
        latinc = (dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS
        loninc = (dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS
    else:
        pole_lat = 90.0
        pole_lon = 0.0
        latinc = 0.0
        loninc = 0.0

    latvar = _lat_varname(wrfin, stagger)
    lonvar = _lon_varname(wrfin, stagger)

    lat_timeidx = timeidx

    is_moving = is_moving_domain(wrfin, latvar=latvar, lonvar=lonvar,
                                 _key=_key)

    # Only need one file and one time if the domain is not moving
    if not is_moving:
        # Always use the 0th time for non-moving domains to avoid problems
        lat_timeidx = 0

        if is_multi_file(wrfin):
            if not is_mapping(wrfin):
                wrfin = next(iter(wrfin))  # only need one file
            else:
                first_entry = next(iter(viewkeys(wrfin)))
                wrfin = wrfin[first_entry]
                key = _key[first_entry]
                return _get_proj_params(wrfin, timeidx, stagger,
                                        method, squeeze, cache, key)

    xlat = extract_vars(wrfin, lat_timeidx, (latvar,), method, squeeze, cache,
                        meta=False, _key=_key)[latvar]
    xlon = extract_vars(wrfin, lat_timeidx, (lonvar,), method, squeeze, cache,
                        meta=False, _key=_key)[lonvar]

    ref_lat = np.ravel(xlat[..., 0, 0])
    ref_lon = np.ravel(xlon[..., 0, 0])

    # Note: fortran index
    known_x = 1.0
    known_y = 1.0

    return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
            pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, loninc)


# known_x and known_y are 0-based
def _kwarg_proj_params(**projparams):
    """Return the map projection parameters.

    This function aggregates the projection parameter keyword
    arguments and also performs sanity checking on them.

    Args:

        **projparams: Projection parameter keyword arguments.

    Returns:

        :obj:`tuple`: The map projection parameters.

    """
    projparams = dict_keys_to_upper(projparams)

    map_proj = projparams.get("MAP_PROJ")
    truelat1 = projparams.get("TRUELAT1")
    truelat2 = projparams.get("TRUELAT2")
    stdlon = projparams.get("STAND_LON")
    ref_lat = projparams.get("REF_LAT")
    ref_lon = projparams.get("REF_LON")
    pole_lat = projparams.get("POLE_LAT", 90.0)
    pole_lon = projparams.get("POLE_LON", 0.0)
    known_x = projparams.get("KNOWN_X")  # Use 0-based
    known_y = projparams.get("KNOWN_Y")  # Use 0-based

    dx = projparams.get("DX")
    dy = projparams.get("DY")
    latinc = projparams.get("LATINC")
    loninc = projparams.get("LONINC")

    # Sanity checks
    # Required args for all projections
    for name, var in viewitems({"MAP_PROJ": map_proj,
                                "REF_LAT": ref_lat,
                                "REF_LON": ref_lon,
                                "KNOWN_X": known_x,
                                "KNOWN_Y": known_y,
                                "DX": dx}):
        if var is None:
            raise ValueError("'{}' argument required".format(name))

    # ref_lat and ref_lon are expected to be lists
    ref_lat = np.ravel(np.asarray([ref_lat]))
    ref_lon = np.ravel(np.asarray([ref_lon]))

    # Fortran wants 1-based indexing
    known_x = known_x + 1
    known_y = known_y + 1

    if map_proj in (ProjectionTypes.LAMBERT_CONFORMAL,
                    ProjectionTypes.POLAR_STEREOGRAPHIC,
                    ProjectionTypes.MERCATOR):
        if truelat1 is None:
            raise ValueError("'TRUELAT1' argument required")
    else:
        if truelat1 is None:
            truelat1 = 0.0

    # Map projection 6 (lat lon) required latinc, loninc, and dy
    if map_proj == ProjectionTypes.LAT_LON:
        if latinc is None:
            raise ValueError("'LATINC' argument required")

        if loninc is None:
            raise ValueError("'LONINC' argument required")

        if dy is None:
            raise ValueError("'DY' argument required")
    else:
        latinc = 0.0
        loninc = 0.0
        dy = 0.0

    return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, pole_lat,
            pole_lon, known_x, known_y, dx, dy, latinc, loninc)


# Will return 0-based indexes
def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0,
              stagger=None, method="cat", squeeze=True, cache=None,
              _key=None, as_int=True, **projparams):
    """Return the x,y coordinates for a specified latitude and longitude.

    The *latitude* and *longitude* arguments can be a single value or a
    sequence of values.

    The leftmost dimension of the returned array represents two different
    quantities:

        - return_val[0,...] will contain the X (west_east) values.
        - return_val[1,...] will contain the Y (south_north) values.

    Args:

        latitude (:obj:`float` or sequence): A single latitude or a sequence
            of latitude values to be converted.

        longitude (:obj:`float` or sequence): A single longitude or a sequence
            of latitude values to be converted.

        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.

        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.

        stagger (:obj:`str`): By default, the latitude and longitude are
            returned on the mass grid, but a staggered grid can be chosen
            with the following options:

                - 'm': Use the mass grid (default).
                - 'u': Use the same staggered grid as the u wind component,
                  which has a staggered west_east (x) dimension.
                - 'v': Use the same staggered grid as the v wind component,
                  which has a staggered south_north (y) dimension.

        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.

        _key (:obj:`int`, optional): A caching key. This is used for internal
            purposes only.  Default is None.

        as_int (:obj:`bool`): Set to True to return the x,y values as
            :obj:`int`, otherwise they will be returned as :obj:`float`.

        **projparams: Map projection keyword arguments to set manually.

    Returns:
        :class:`xarray.DataArray` or :class:`numpy.ndarray`: The
        x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y).
        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.

    """

    if wrfin is not None:
        (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
         pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
         loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze,
                                    cache, _key)
    else:
        (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
         pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
         loninc) = _kwarg_proj_params(**projparams)

    if isinstance(latitude, Iterable):
        lats = np.asarray(latitude)
        lons = np.asarray(longitude)

        # Note:  For scalars, this will make a single element array
        lats = lats.ravel()

        lons = lons.ravel()

        if (lats.size != lons.size):
            raise ValueError("'latitude' and 'longitude' "
                             "must be the same length")

        if ref_lat.size == 1:
            outdim = [2, lats.size]
            extra_dims = [outdim[1]]
        else:
            # Moving domain will have moving ref_lats/ref_lons
            outdim = [2, ref_lat.size, lats.size]
            extra_dims = outdim[1:]

        result = np.empty(outdim, np.float64)

        for left_idxs in iter_left_indexes(extra_dims):
            # Left indexes is a misnomer, since these will be on the right
            x_idxs = (0,) + left_idxs
            y_idxs = (1,) + left_idxs
            if ref_lat.size == 1:
                ref_lat_val = ref_lat[0]
                ref_lon_val = ref_lon[0]
            else:
                ref_lat_val = ref_lat[left_idxs[-2]]
                ref_lon_val = ref_lon[left_idxs[-2]]

            lat = lats[left_idxs[-1]]
            lon = lons[left_idxs[-1]]

            xy = _lltoxy(map_proj, truelat1, truelat2, stdlon,
                         ref_lat_val, ref_lon_val, pole_lat, pole_lon,
                         known_x, known_y, dx, dy, latinc, loninc,
                         lat, lon)

            # Note:  comes back from fortran as y,x
            result[x_idxs] = xy[1]
            result[y_idxs] = xy[0]

    else:
        result = np.empty((2,), np.float64)

        fort_out = _lltoxy(map_proj, truelat1, truelat2, stdlon,
                           ref_lat, ref_lon, pole_lat, pole_lon,
                           known_x, known_y, dx, dy, latinc, loninc,
                           latitude, longitude)

        # Note, comes back from fortran as y,x.  So, need to swap them.
        result[0] = fort_out[1]
        result[1] = fort_out[0]

    # Make indexes 0-based
    result = result - 1

    if as_int:
        result = np.rint(result).astype(int)

    return result


# X and Y should be 0-based
def _xy_to_ll(x, y, wrfin=None, timeidx=0, stagger=None,
              method="cat", squeeze=True, cache=None, _key=None,
              **projparams):

    """Return the latitude and longitude for specified x,y coordinates.

    The *x* and *y* arguments can be a single value or a sequence of values.

    The leftmost dimension of the returned array represents two different
    quantities:

        - return_val[0,...] will contain the latitude values.
        - return_val[1,...] will contain the longitude values.

    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.

        x (:obj:`float` or sequence): A single x-coordinate or a sequence
            of x-coordinate values to be converted.

        y (:obj:`float` or sequence): A single y-coordinate or a sequence
            of y-coordinate values to be converted.

        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.

        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.

        stagger (:obj:`str`): By default, the latitude and longitude are
            returned on the mass grid, but a staggered grid can be chosen
            with the following options:

                - 'm': Use the mass grid (default).
                - 'u': Use the same staggered grid as the u wind component,
                  which has a staggered west_east (x) dimension.
                - 'v': Use the same staggered grid as the v wind component,
                  which has a staggered south_north (y) dimension.

        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.

        _key (:obj:`int`, optional): A caching key. This is used for internal
            purposes only.  Default is None.

        **projparams: Map projection keyword arguments to set manually.

    Returns:
        :class:`xarray.DataArray` or :class:`numpy.ndarray`: The
        latitude and longitude values whose leftmost dimension is 2
        (0=latitude, 1=longitude).
        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.

    """

    if wrfin is not None:
        (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
         pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
         loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze,
                                    cache, _key)
    else:
        (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
         pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
         loninc) = _kwarg_proj_params(**projparams)

    if isinstance(x, Iterable):
        x_arr = np.asarray(x)
        y_arr = np.asarray(y)

        # Convert 0-based x, y to 1-based
        x_arr = x_arr + 1
        y_arr = y_arr + 1

        x_arr = x_arr.ravel()

        y_arr = y_arr.ravel()

        if (x_arr.size != y_arr.size):
            raise ValueError("'x' and 'y' must be the same length")

        if ref_lat.size == 1:
            outdim = [2, x_arr.size]
            extra_dims = [outdim[1]]
        else:
            # Moving domain will have moving ref_lats/ref_lons
            outdim = [2, ref_lat.size, x_arr.size]
            extra_dims = outdim[1:]

        result = np.empty(outdim, np.float64)

        for left_idxs in iter_left_indexes(extra_dims):
            lat_idxs = (0,) + left_idxs
            lon_idxs = (1,) + left_idxs

            if ref_lat.size == 1:
                ref_lat_val = ref_lat[0]
                ref_lon_val = ref_lon[0]
            else:
                ref_lat_val = ref_lat[left_idxs[-2]]
                ref_lon_val = ref_lon[left_idxs[-2]]

            x_val = x_arr[left_idxs[-1]]
            y_val = y_arr[left_idxs[-1]]

            ll = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat_val,
                         ref_lon_val, pole_lat, pole_lon, known_x, known_y,
                         dx, dy, latinc, loninc, x_val, y_val)

            result[lat_idxs] = ll[0]
            result[lon_idxs] = ll[1]

    else:
        # Convert 0-based to 1-based for Fortran
        x_val = x + 1
        y_val = y + 1

        result = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat,
                         ref_lon, pole_lat, pole_lon, known_x, known_y,
                         dx, dy, latinc, loninc, x_val, y_val)

    return result
