# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
#


def ds_x_to_json(ds, widget):
    return ds2json(
        ds,
        widget.zonal_speed,
        widget.meridional_speed,
        widget.latitude_dimension,
        widget.longitude_dimension,
        widget.units,
    )


def ds2json(ds, u_var, v_var, lat_dim="latitude", lon_dim="longitude", units=None):
    """
    Assumes that the velocity components are given on a regular grid
    (fixed spacing in latitude and longitude).

    Parameters
    ----------
    u_var : str
        Name of the U-component (zonal) variable.
    v_var : str
        Name of the V-component (meridional) variable.
    lat_dim : str, optional
        Name of the latitude dimension/coordinate
        (default: 'latitude').
    lon_dim : str, optional
        Name of the longitude dimension/coordinate
        (default: 'longitude').
    units : str, optional
        Velocity units (default: try getting units from the
        'units' attributes of `u_var` and `v_var`).
    """
    import numpy as np

    ds = ds.copy()
    for var_name in (u_var, v_var):
        var_dims = ds[var_name].dims

        if set(var_dims) != set([lat_dim, lon_dim]):
            raise ValueError(
                "Invalid dimensions for variable '{}' in Dataset: "
                "should include only {}, found {}.".format(
                    var_name, (lat_dim, lon_dim), var_dims
                )
            )

        # If dataset contains nans replace with 0
        ds[var_name] = ds[var_name].fillna(0)

    if units is None:
        u_var_units = ds[u_var].attrs.get("units")
        v_var_units = ds[v_var].attrs.get("units")

        if u_var_units != v_var_units:
            raise ValueError(
                "Different units found for U-component '{}' and "
                "V-component '{}' variables: '{}' and '{}'".format(
                    u_var, v_var, u_var_units, v_var_units
                )
            )

        units = u_var_units

    if units is None:
        units = ""

    # Data should be in gaussian grid format (latitudes descending)
    if np.any(np.diff(ds[lat_dim].values) >= 0):
        ds = ds.sel(**{lat_dim: slice(None, None, -1)})

    # infer grid specifications (assume a rectangular grid)
    lat = ds[lat_dim].values
    lon = ds[lon_dim].values

    lon_left = float(lon.min())
    lon_right = float(lon.max())
    lat_lower = float(lat.min())
    lat_upper = float(lat.max())

    dx = float((lon_right - lon_left) / (lon.size - 1))
    dy = float((lat_upper - lat_lower) / (lat.size - 1))

    nx = lon.size
    ny = lat.size

    u_v_spec = ([2, 3], ["Eastward current", "Northward current"], [u_var, v_var])

    velocity_data = []

    for p_number, p_name, var_name in zip(*u_v_spec):
        velocity_data.append(
            {
                "header": {
                    "parameterUnit": units,
                    "parameterNumber": p_number,
                    "dx": dx,
                    "dy": dy,
                    "parameterNumberName": p_name,
                    "la1": lat_upper,
                    "la2": lat_lower,
                    "parameterCategory": 2,
                    "lo2": lon_right,
                    "nx": nx,
                    "ny": ny,
                    "refTime": "2017-02-01 23:00:00",
                    "lo1": lon_left,
                },
                "data": ds[var_name].values.flatten().tolist(),
            }
        )

    return velocity_data