"""
Coordinate and grid helpers for dLux Fresnel modules.

This module normalises “grid-like” inputs into explicit 1D coordinate
vectors (x, y). All consumers assume *separable*, *rectilinear* grids:

  - Separable: coordinates factor as x(x_index), y(y_index).
  - Rectilinear: each axis has a single spacing (equally spaced samples).
  - Anisotropic: dx and dy may differ between axes.

Supported specifications
------------------------
We accept three basic forms:

1) Uniform centred grid:

       spec = (N, d)
       spec = ((Nx, Ny), (dx, dy))

   where N is int or (int, int) and d is float/int or (float/int, float/int).
   Coordinates are generated with dLux.utils.nd_coords for each axis.

2) Single 1D coordinate vector:

       spec = x    # Array, shape (N,)

   interpreted as symmetric 2D grid with x and y identical.

3) Explicit per-axis coordinates:

       spec = (x, y)  # tuple of Arrays, each shape (N,)

All arrays are expected to be 1D, monotonically increasing, and
approximately equally spaced along each axis. No shape or monotonicity
checks are enforced here.
"""

from __future__ import annotations
from jax import Array
import jax.numpy as np
import jax.tree as jtu

__all__ = [
    "unpack_size",
    "unpack_scale",
    "unpack_coords",
    "unpack_coord_spec",
    "unpack_coord_specs",
    "pad_coords",
    "pad_to",
    "crop_to",
    "nd_coords",
]


# ---------------------------------------------------------------------
# Basic unpack helpers (scalar → pair)
# ---------------------------------------------------------------------


def unpack_size(N: int | tuple[int, int]) -> tuple[int, int]:
    """
    Normalise a size specification to (Nx, Ny).

    Parameters
    ----------
    N : int or (int, int)
        Grid size. If scalar, applies to both axes.

    Returns
    -------
    Nx, Ny : int
    """
    if isinstance(N, tuple):
        Nx, Ny = N
        return int(Nx), int(Ny)
    return int(N), int(N)


def unpack_scale(d: float | tuple) -> tuple[float, float]:
    """
    Normalise a sampling specification to (dx, dy).

    Parameters
    ----------
    d : float or (float, float)
        Grid spacing. If scalar, applies to both axes. Integers are also
        accepted and cast to float.

    Returns
    -------
    dx, dy : float
    """
    if isinstance(d, tuple):
        dx, dy = d
        return dx, dy
    val = d
    return val, val


def unpack_coords(coords: Array | tuple) -> tuple[Array, Array]:
    """
    Normalise a coordinate specification to (x, y) 1D arrays.

    Parameters
    ----------
    coords : Array or (Array, Array)
        - 1D Array: symmetric grid, x = y = coords
        - (x, y)   : explicit per-axis coordinates

    Returns
    -------
    x, y : Array
        1D coordinate vectors along x and y.

    Notes
    -----
    We do *not* accept arbitrary 2D coordinate arrays here. All downstream
    code assumes separability (rectilinear grids).
    """
    if isinstance(coords, Array):
        if coords.ndim != 1:
            raise ValueError(
                "unpack_coords: Array coords must be 1D. "
                f"Got ndim={coords.ndim} and shape={coords.shape}."
            )
        return coords, coords

    x, y = coords
    return x, y


# ---------------------------------------------------------------------
# CoordSpec -> (x, y) on a centred grid
# ---------------------------------------------------------------------


def _is_size_like(obj) -> bool:
    if isinstance(obj, int):
        return True
    if (
        isinstance(obj, tuple)
        and len(obj) == 2
        and all(isinstance(v, int) for v in obj)
    ):
        return True
    if isinstance(obj, Array) and obj.shape == ():
        return True
    if (
        isinstance(obj, tuple)
        and len(obj) == 2
        and all(isinstance(v, Array) and v.shape == () for v in obj)
    ):
        return True
    return False


def _is_scale_like(obj) -> bool:
    if isinstance(obj, (int, float)):
        return True
    if (
        isinstance(obj, tuple)
        and len(obj) == 2
        and all(isinstance(v, (int, float)) for v in obj)
    ):
        return True
    if isinstance(obj, Array) and obj.shape == ():
        return True
    if (
        isinstance(obj, tuple)
        and len(obj) == 2
        and all(isinstance(v, Array) and v.shape == () for v in obj)
    ):
        return True
    return False


def unpack_coord_spec(spec: Array | tuple) -> tuple[Array, Array]:
    """
    Normalise a coordinate specification into explicit (x, y) vectors.

    Supported forms
    ---------------
    1) Uniform centred/offset grid:

         spec = (N, d)
         spec = ((Nx, Ny), (dx, dy))
         spec = (N, d, x0)
         spec = ((Nx, Ny), (dx, dy), (x0, y0))

       where:
       - N is int or (int, int)
       - d is float/int or (float/int, float/int)
       - x0/y0 give the coordinate of the grid centre

       Coordinates are generated via `dLux.utils.nd_coords` and shifted by
       the specified centre offset.

    2) Single 1D coordinate vector:

         spec = x    # Array, shape (N,)

       Interpreted as symmetric grid: x = y = spec.

    3) Explicit per-axis coordinates:

         spec = (x, y)    # tuple of 1D Arrays

    Returns
    -------
    x, y : Array
        1D coordinate vectors along x and y.
    """
    # Case 2: single explicit 1D coordinate vector
    if isinstance(spec, Array):
        return unpack_coords(spec)

    if isinstance(spec, tuple):
        # ------------------------------------------------------------------
        # Uniform grid specs
        # ------------------------------------------------------------------
        if len(spec) == 2:
            a, b = spec

            if _is_size_like(a) and _is_scale_like(b):
                Nx, Ny = unpack_size(a)
                dx, dy = unpack_scale(b)

                x = nd_coords(Nx, dx)
                y = nd_coords(Ny, dy)
                return x, y

            # Otherwise treat as explicit (x, y) coords
            elif not _is_size_like(a) and not _is_scale_like(b):
                return unpack_coords(spec)

        elif len(spec) == 3:
            a, b, c = spec

            if _is_size_like(a) and _is_scale_like(b) and _is_scale_like(c):
                Nx, Ny = unpack_size(a)
                dx, dy = unpack_scale(b)
                x0, y0 = unpack_scale(c)

                x = nd_coords(Nx, dx, x0)
                y = nd_coords(Ny, dy, y0)
                return x, y

    raise TypeError(
        "unpack_coord_spec: unsupported spec type. Expected Array, (x, y) "
        "tuple, or scalar tuple (N, d), ((Nx, Ny), (dx, dy)), "
        "(N, d, x0), or ((Nx, Ny), (dx, dy), (x0, y0)). "
        f"Got type {type(spec)} with value {spec}."
    )


def unpack_coord_specs(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
) -> tuple:
    """
    Convenience wrapper to normalise both input and output specs.

    Parameters
    ----------
    spec_in, spec_out
        Coordinate specifications for input and output planes.

    Returns
    -------
    (x_in, y_in), (x_out, y_out) : tuple of Arrays
    """
    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)
    return (x_in, y_in), (x_out, y_out)


def pad_coords(
    x: Array,
    y: Array,
    npad: int | tuple[int, int],
) -> tuple[Array, Array]:
    """
    Symmetrically pad 1D coordinate vectors to larger sizes.

    The padded coordinates preserve:
    - the original pixel spacing
    - the original grid centre

    Parameters
    ----------
    x, y : Array
        1D coordinate vectors.
    npad : int or tuple[int, int]
        Target padded size. If scalar, the same size is used for both axes.
    check_uniform : bool, optional
        If True, verify that x and y are uniformly spaced.

    Returns
    -------
    x_pad, y_pad : Array
        Padded 1D coordinate vectors.
    """
    Nx_pad, Ny_pad = unpack_size(npad)

    Nx_in = x.shape[0]
    Ny_in = y.shape[0]

    if Nx_pad < Nx_in or Ny_pad < Ny_in:
        raise ValueError(
            f"Target padded size {(Nx_pad, Ny_pad)} must be >= input size {(Nx_in, Ny_in)}."
        )

    if (Nx_pad % 2) != (Nx_in % 2):
        raise ValueError(
            f"x-axis parity mismatch: input Nx={Nx_in}, padded Nx={Nx_pad}. "
            "Symmetric padding requires even->even or odd->odd."
        )

    if (Ny_pad % 2) != (Ny_in % 2):
        raise ValueError(
            f"y-axis parity mismatch: input Ny={Ny_in}, padded Ny={Ny_pad}. "
            "Symmetric padding requires even->even or odd->odd."
        )

    dx = x[1] - x[0]
    dy = y[1] - y[0]

    xc = 0.5 * (x[0] + x[-1])
    yc = 0.5 * (y[0] + y[-1])

    x_pad = nd_coords(Nx_pad, dx, xc)
    y_pad = nd_coords(Ny_pad, dy, yc)

    return x_pad, y_pad


def pad_to(array: Array, npixels: int | tuple[int, int]) -> Array:
    """
    Zero-pad a 2D array of shape (Ny, Nx) to (Ny_pad, Nx_pad).

    Padding is symmetric about the grid centre, so input and output sizes
    must have matching parity along each axis (even→even or odd→odd).

    Parameters
    ----------
    array : Array
        Input array of shape (Ny, Nx).
    npixels : int or tuple[int, int]
        Target size (Nx_pad, Ny_pad). If scalar, both axes are padded
        to the same size.

    Returns
    -------
    Array
        Padded array of shape (Nx_pad, Ny_pad).
    """
    Nx_pad, Ny_pad = unpack_size(npixels)

    Ny_in, Nx_in = array.shape

    if Nx_pad < Nx_in or Ny_pad < Ny_in:
        raise ValueError(
            f"Target size {(Nx_pad, Ny_pad)} must be >= input size {(Nx_in, Ny_in)}."
        )

    if (Nx_in % 2) != (Nx_pad % 2):
        raise ValueError(
            f"x parity mismatch: Nx_in={Nx_in}, Nx_pad={Nx_pad}. "
            "Only even→even or odd→odd padding allowed."
        )

    if (Ny_in % 2) != (Ny_pad % 2):
        raise ValueError(
            f"y parity mismatch: Ny_in={Ny_in}, Ny_pad={Ny_pad}. "
            "Only even→even or odd→odd padding allowed."
        )

    px = (Nx_pad - Nx_in) // 2
    py = (Ny_pad - Ny_in) // 2

    return np.pad(array, ((py, py), (px, px)))


def crop_to(array: Array, npixels: int | tuple[int, int]) -> Array:
    """
    Centrally crop a 2D array of shape (Ny, Nx) to (Ny_crop, Nx_crop).

    Cropping is symmetric about the grid centre, so input and output sizes
    must have matching parity along each axis (even→even or odd→odd).

    Parameters
    ----------
    array : Array
        Input array of shape (Ny, Nx).
    npixels : int or tuple[int, int]
        Target size (Nx_crop, Ny_crop). If scalar, both axes are cropped
        to the same size.

    Returns
    -------
    Array
        Cropped array of shape (Ny_crop, Nx_crop).
    """
    Nx_crop, Ny_crop = unpack_size(npixels)

    Ny_in, Nx_in = array.shape

    if Nx_crop > Nx_in or Ny_crop > Ny_in:
        raise ValueError(
            f"Target size {(Nx_crop, Ny_crop)} must be <= input size {(Nx_in, Ny_in)}."
        )

    if (Nx_in % 2) != (Nx_crop % 2):
        raise ValueError(
            f"x parity mismatch: Nx_in={Nx_in}, Nx_crop={Nx_crop}. "
            "Only even→even or odd→odd cropping allowed."
        )

    if (Ny_in % 2) != (Ny_crop % 2):
        raise ValueError(
            f"y parity mismatch: Ny_in={Ny_in}, Ny_crop={Ny_crop}. "
            "Only even→even or odd→odd cropping allowed."
        )

    cx = (Nx_in - Nx_crop) // 2
    cy = (Ny_in - Ny_crop) // 2

    return array[cy : cy + Ny_crop, cx : cx + Nx_crop]


### Raised dLux coords utilities ###
def _cast_tuple(x, name):

    # Validate npixels and ensure tuple
    if isinstance(x, int):
        x = (x,)
    elif isinstance(x, tuple):
        for n in x:
            if not isinstance(n, int):
                raise ValueError(f"All {name} must be integers.")
    else:
        raise ValueError(f"{name} must be an int or a tuple of ints.")

    return x


def _cast_scalar(x, ndim, name):

    _is_numeric = lambda x: isinstance(x, (int, float))
    _is_scalar_array = lambda x: isinstance(x, Array) and x.ndim == 0
    _is_scalar = lambda x: _is_numeric(x) or _is_scalar_array(x)

    if _is_scalar(x):
        x = (x,) * ndim
    elif isinstance(x, Array):
        if x.ndim != 1 or x.shape[0] != ndim:
            raise ValueError(f"Length of {name} array must match number of dimensions.")
        x = tuple(x)
    elif isinstance(x, tuple):
        if len(x) != ndim:
            raise ValueError(f"Length of {name} must match number of dimensions.")
        for z in x:
            if not _is_scalar(z):
                raise ValueError(
                    f"All {name} must be scalars (int, float, or scalar Array)."
                )
    else:
        raise ValueError(
            f"{name} must be a scalar (int, float, or scalar Array) or "
            f"a tuple of scalars."
        )

    return x


def _input_len(x, name):
    if isinstance(x, tuple):
        return len(x)
    if isinstance(x, Array):
        if x.ndim == 0:
            return 1
        if x.ndim == 1:
            return x.shape[0]
        raise ValueError(
            f"{name} must be a scalar, tuple, or 1D array; got ndim={x.ndim}."
        )
    return 1


def nd_coords(
    npixels: int | tuple[int, ...],
    pixel_scales: float | tuple[float, ...] = 1.0,
    offsets: float | tuple[float, ...] = 0.0,
    indexing: str = "xy",
) -> Array:
    """
    Returns a set of nd pixel center coordinates, with an optional offset. Each
    dimension can have a different number of pixels, pixel scale and offset by passing
    in tuples of values: `nd_coords((10, 10), (1, 2), (0, 1))`. pixel scale and offset
    can also be passed in as floats to apply those values to all dimensions, i.e.:
    `nd_coords((10, 10), 1, 0)`.

    The indexing argument is the same as in numpy.meshgrid, i.e.: giving the
    string ‘ij’ returns a meshgrid with matrix indexing, while ‘xy’ returns a
    meshgrid with Cartesian indexing. In the 2-D case with inputs of length M
    and N, the outputs are of shape (N, M) for ‘xy’ indexing and (M, N) for
    ‘ij’ indexing. In the 3-D case with inputs of length M, N and P, outputs
    are of shape (N, M, P) for ‘xy’ indexing and (M, N, P) for ‘ij’ indexing.

    Parameters
    ----------
    npixels : int | tuple[int, ...]
        The number of pixels in each dimension.
    pixel_scales : float | tuple[float, ...] = 1.0
        The pixel_scales of each dimension. If a tuple, the length
        of the tuple must match the number of dimensions. If a float, the same
        scale is applied to all dimensions.
    offsets : float | tuple[float, ...] = 0.0
        The offset of the pixel centers in each dimension. If a tuple, the
        length of the tuple must match the number of dimensions. If a float,
        the same offset is applied to all dimensions.
        set to 0.
    indexing : str = 'xy'
        The indexing of the output. Default is 'xy'. See numpy.meshgrid for
        more details.

    Returns
    -------
    coordinates : Array
        The positions of the pixel centers in the given dimensions.
    """
    if indexing not in ["xy", "ij"]:
        raise ValueError("indexing must be either 'xy' or 'ij'.")

    # Validate npixels and ensure tuple
    npixels = _cast_tuple(npixels, "npixels")

    # Get the number of dimensions
    ndim = max(
        len(npixels), _input_len(pixel_scales, "mean"), _input_len(offsets, "std")
    )

    # Validate and ensure tuples
    pixel_scales = _cast_scalar(pixel_scales, ndim, "pixel_scales")
    offsets = _cast_scalar(offsets, ndim, "offsets")

    if len(npixels) != ndim:
        npixels *= ndim

    def pixel_fn(n, offset, scale):
        start = -(n - 1) / 2 * scale - offset
        end = (n - 1) / 2 * scale - offset
        return np.linspace(start, end, n)

    # Generate the linear edges of each axes
    lin_pixels = jtu.map(pixel_fn, npixels, offsets, pixel_scales)

    # Output (x, y) for 2d, else in order.
    positions = np.array(np.meshgrid(*lin_pixels, indexing=indexing))

    # Squeeze the output in case of 1d input
    return np.squeeze(positions)
