import numpy
import pytest

from ..math.regular_grid import ScatterDataInterpolator


@pytest.mark.parametrize(
    "ndim, ndatadim, units, unit_conv",
    [
        pytest.param(
            2,
            0,
            ["um", "um"],
            [1.0, 1.0],
            id="2D_homogeneous_0Ddata",
        ),
        pytest.param(
            2,
            1,
            ["um", "um"],
            [1.0, 1.0],
            id="2D_homogeneous_1Ddata",
        ),
        pytest.param(
            2,
            2,
            ["um", "um"],
            [1.0, 1.0],
            id="2D_homogeneous_2Ddata",
        ),
        pytest.param(
            2,
            1,
            ["um", "um"],
            [1.0, 1.0],
            id="2D_homogeneous_3Ddata",
        ),
        pytest.param(
            3,
            0,
            ["um", "um", "um"],
            [1.0, 1.0, 1.0],
            id="3D_homogeneous_0Ddata",
        ),
        pytest.param(
            2,
            0,
            ["um", "mm"],
            [1.0, 1e3],
            id="2D_heterogeneous_0Ddata",
        ),
        pytest.param(
            3,
            0,
            ["um", "um", "mm"],
            [1.0, 1.0, 1e3],
            id="3D_heterogeneous_0Ddata",
        ),
        pytest.param(
            3,
            1,
            ["um", "um", "mm"],
            [1.0, 1.0, 1e3],
            id="3D_heterogeneous_1Ddata",
        ),
        pytest.param(
            4,
            0,
            ["um", "mm", "deg", "rad"],
            [1.0, 1e3, 1.0, 180 / numpy.pi],
            id="4D_mixed_0Ddata",
        ),
    ],
)
def test_regular_regridding_mixed_scaled(ndim, ndatadim, units, unit_conv):
    # Generate raw scatter coordinates
    state = numpy.random.RandomState(42)
    N_points = 100
    scatter_coords = [state.uniform(-10, 10, N_points) for _ in units]

    # Arbitrary linear relationship between values and coordinates
    coeffs = [i + 1 for i in range(ndim)]
    offset = 5
    values = (
        sum(c * x * uc for c, x, uc in zip(coeffs, scatter_coords, unit_conv)) + offset
    )

    # Add extra data dimensions for each scatter point
    data_shape = tuple(range(10, 10 + ndatadim))
    values = _add_data_dimensions(values, data_shape)

    # Create interpolator with raw coordinates
    names = [f"dim{i}" for i in range(ndim)]
    interpolator = ScatterDataInterpolator(
        scatter_coords,
        names,
        units,
        method="linear",
        reference_units="first",
        scale_method="range",
    )
    interpolated = interpolator.regrid(values)

    # Interpolated values as expected from the
    # - linear relationship between values and coordinates
    # - and linear interpolation in scaled space.
    grid_coords = interpolator.expanded_grid_coordinates.T
    expected = sum(c * x for c, x in zip(coeffs, grid_coords)) + offset
    expected = _add_data_dimensions(expected, data_shape)
    expected = expected.reshape(interpolated.shape)

    # Validate interpolation
    assert numpy.nanmax(numpy.abs(interpolated - expected)) < 1e-9


def _add_data_dimensions(array: numpy.ndarray, data_shape: tuple) -> numpy.ndarray:
    """
    Expand a 1D array to include additional detector/data dimensions.

    :param array: Input 1D array of shape (N,).
    :param data_shape: Tuple representing the shape of additional data dimensions to add.
    :returns: Array of shape (N, *data_shape), where the original array is broadcasted
              along the new dimensions.
    """
    reshaped_array = array[(...,) + (numpy.newaxis,) * len(data_shape)]
    return numpy.tile(reshaped_array, data_shape)
