# Copyright (c) 2018 MetPy Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
"""Test the `slices` module."""

import numpy as np
import pytest
import xarray as xr

from metpy.interpolate import cross_section, geodesic, interpolate_to_slice
from metpy.testing import assert_array_almost_equal, needs_cartopy
from metpy.units import units


@pytest.fixture()
def test_ds_lonlat():
    """Return dataset on a lon/lat grid with no time coordinate for use in tests."""
    data_temp = np.linspace(250, 300, 5 * 6 * 7).reshape((5, 6, 7)) * units.kelvin
    data_rh = np.linspace(0, 1, 5 * 6 * 7).reshape((5, 6, 7)) * units.dimensionless
    ds = xr.Dataset(
        {
            'temperature': (['isobaric', 'lat', 'lon'], data_temp),
            'relative_humidity': (['isobaric', 'lat', 'lon'], data_rh)
        },
        coords={
            'isobaric': xr.DataArray(
                np.linspace(1000, 500, 5),
                name='isobaric',
                dims=['isobaric'],
                attrs={'units': 'hPa'}
            ),
            'lat': xr.DataArray(
                np.linspace(30, 45, 6),
                name='lat',
                dims=['lat'],
                attrs={'units': 'degrees_north'}
            ),
            'lon': xr.DataArray(
                np.linspace(255, 275, 7),
                name='lon',
                dims=['lon'],
                attrs={'units': 'degrees_east'}
            )
        }
    )
    return ds.metpy.parse_cf()


@pytest.fixture()
def test_ds_xy():
    """Return dataset on a x/y grid with a time coordinate for use in tests."""
    data_temperature = np.linspace(250, 300, 5 * 6 * 7).reshape((1, 5, 6, 7)) * units.kelvin
    ds = xr.Dataset(
        {
            'temperature': (['time', 'isobaric', 'y', 'x'], data_temperature),
            'lambert_conformal': ([], '')
        },
        coords={
            'time': xr.DataArray(
                np.array(['2018-07-01T00:00'], dtype='datetime64[ns]'),
                name='time',
                dims=['time']
            ),
            'isobaric': xr.DataArray(
                np.linspace(1000, 500, 5),
                name='isobaric',
                dims=['isobaric'],
                attrs={'units': 'hPa'}
            ),
            'y': xr.DataArray(
                np.linspace(-1500, 0, 6),
                name='y',
                dims=['y'],
                attrs={'units': 'km'}
            ),
            'x': xr.DataArray(
                np.linspace(-500, 3000, 7),
                name='x',
                dims=['x'],
                attrs={'units': 'km'}
            )
        }
    )
    ds['temperature'].attrs = {
        'grid_mapping': 'lambert_conformal'
    }
    ds['lambert_conformal'].attrs = {
        'grid_mapping_name': 'lambert_conformal_conic',
        'standard_parallel': 50.0,
        'longitude_of_central_meridian': -107.0,
        'latitude_of_projection_origin': 50.0,
        'earth_shape': 'spherical',
        'earth_radius': 6367470.21484375
    }
    return ds.metpy.parse_cf()


@pytest.mark.parametrize('bad_units', [False, True])
def test_interpolate_to_slice_against_selection(test_ds_lonlat, bad_units):
    """Test interpolate_to_slice on a simple operation."""
    data = test_ds_lonlat['temperature']

    # interpolate_to_slice shouldn't care about units
    if bad_units:
        # Needed so we can go back to using attribute metadata
        data = data.metpy.dequantify()
        data.attrs['units'] = 'my_bad_units'

    path = np.array([[265.0, 30.],
                     [265.0, 36.],
                     [265.0, 42.]])
    test_slice = interpolate_to_slice(data, path)
    true_slice = data.sel({'lat': [30., 36., 42.], 'lon': 265.0})
    # Coordinates differ, so just compare the data
    assert_array_almost_equal(true_slice.data, test_slice.data, 5)


@needs_cartopy
def test_geodesic(test_ds_xy):
    """Test the geodesic construction."""
    crs = test_ds_xy['temperature'].metpy.pyproj_crs
    path = geodesic(crs, (36.46, -112.45), (42.95, -68.74), 7)
    truth = np.array([[-4.99495719e+05, -1.49986599e+06],
                      [9.84044354e+04, -1.26871737e+06],
                      [6.88589099e+05, -1.02913966e+06],
                      [1.27269045e+06, -7.82037603e+05],
                      [1.85200974e+06, -5.28093957e+05],
                      [2.42752546e+06, -2.67710326e+05],
                      [2.99993290e+06, -9.39107692e+02]])
    assert_array_almost_equal(path, truth, 0)


@needs_cartopy
def test_cross_section_dataarray_and_linear_interp(test_ds_xy):
    """Test the cross_section function with a data array and linear interpolation."""
    data = test_ds_xy['temperature']
    start, end = ((36.46, -112.45), (42.95, -68.74))
    data_cross = cross_section(data, start, end, steps=7)
    truth_values = np.array([[[250.00095489, 251.53646673, 253.11586664, 254.73477364,
                               256.38991013, 258.0794356, 259.80334269],
                              [260.04880178, 261.58431362, 263.16371353, 264.78262053,
                               266.43775702, 268.12728249, 269.85118958],
                              [270.09664867, 271.63216051, 273.21156042, 274.83046742,
                               276.48560391, 278.17512938, 279.89903647],
                              [280.14449556, 281.6800074, 283.25940731, 284.87831431,
                               286.5334508, 288.22297627, 289.94688336],
                              [290.19234245, 291.72785429, 293.3072542, 294.9261612,
                               296.58129769, 298.27082316, 299.99473025]]])
    truth_values_x = np.array([-499495.71907062, 98404.43537514, 688589.09865512,
                               1272690.44926197, 1852009.73516881, 2427525.45740665,
                               2999932.89862589])
    truth_values_y = np.array([-1499865.98780602, -1268717.36799267, -1029139.66048478,
                               -782037.60343652, -528093.95678826, -267710.32566917,
                               -939.10769171])
    index = xr.DataArray(range(7), name='index', dims=['index'])

    data_truth_x = xr.DataArray(
        truth_values_x,
        name='x',
        coords={
            'metpy_crs': data['metpy_crs'],
            'y': (['index'], truth_values_y),
            'x': (['index'], truth_values_x),
            'index': index,
        },
        dims=['index']
    )
    data_truth_y = xr.DataArray(
        truth_values_y,
        name='y',
        coords={
            'metpy_crs': data['metpy_crs'],
            'y': (['index'], truth_values_y),
            'x': (['index'], truth_values_x),
            'index': index,
        },
        dims=['index']
    )
    data_truth = xr.DataArray(
        truth_values * units.kelvin,
        name='temperature',
        coords={
            'time': data['time'],
            'isobaric': data['isobaric'],
            'index': index,
            'metpy_crs': data['metpy_crs'],
            'y': data_truth_y,
            'x': data_truth_x
        },
        dims=['time', 'isobaric', 'index']
    )

    xr.testing.assert_allclose(data_truth, data_cross)


def test_cross_section_dataarray_projection_noop(test_ds_xy):
    """Test the cross_section function with a projection dataarray."""
    data = test_ds_xy['lambert_conformal']
    start, end = ((36.46, -112.45), (42.95, -68.74))
    data_cross = cross_section(data, start, end, steps=7)
    xr.testing.assert_identical(data, data_cross)


@needs_cartopy
def test_cross_section_dataset_and_nearest_interp(test_ds_lonlat):
    """Test the cross_section function with a dataset and nearest interpolation."""
    start, end = (30.5, 255.5), (44.5, 274.5)
    data_cross = cross_section(test_ds_lonlat, start, end, steps=7, interp_type='nearest')
    nearest_values = test_ds_lonlat.isel(lat=xr.DataArray([0, 1, 2, 3, 3, 4, 5], dims='index'),
                                         lon=xr.DataArray(range(7), dims='index'))
    truth_temp = nearest_values['temperature'].metpy.unit_array
    truth_rh = nearest_values['relative_humidity'].metpy.unit_array
    truth_values_lon = np.array([255.5, 258.20305939, 261.06299342, 264.10041516,
                                 267.3372208, 270.7961498, 274.5])
    truth_values_lat = np.array([30.5, 33.02800969, 35.49306226, 37.88512911, 40.19271688,
                                 42.40267088, 44.5])
    index = xr.DataArray(range(7), name='index', dims=['index'])

    data_truth = xr.Dataset(
        {
            'temperature': (['isobaric', 'index'], truth_temp),
            'relative_humidity': (['isobaric', 'index'], truth_rh)
        },
        coords={
            'isobaric': test_ds_lonlat['isobaric'],
            'index': index,
            'metpy_crs': test_ds_lonlat['metpy_crs'],
            'lat': (['index'], truth_values_lat),
            'lon': (['index'], truth_values_lon)
        },
    )

    xr.testing.assert_allclose(data_truth, data_cross)


def test_interpolate_to_slice_error_on_missing_coordinate(test_ds_lonlat):
    """Test that the proper error is raised with missing coordinate."""
    # Use a variable with a coordinate removed
    data_bad = test_ds_lonlat['temperature'].copy()
    del data_bad['lat']
    path = np.array([[265.0, 30.],
                     [265.0, 36.],
                     [265.0, 42.]])

    with pytest.raises(ValueError):
        interpolate_to_slice(data_bad, path)


def test_cross_section_error_on_missing_coordinate(test_ds_lonlat):
    """Test that the proper error is raised with missing coordinate."""
    # Use a variable with no crs coordinate
    data_bad = test_ds_lonlat['temperature'].copy()
    del data_bad['metpy_crs']
    start, end = (30.5, 255.5), (44.5, 274.5)

    with pytest.raises(ValueError):
        cross_section(data_bad, start, end)
