#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Date          : 2025-08-30
# Author        : Lancelot PINCET
# GitHub        : https://github.com/LancelotPincet
# Library       : arrLP
# Module        : coordinates

"""
This function will return a ndarray corresponding to the coordinates array of the input.
"""



# %% Libraries
import numpy as np
from arrlp import get_xp



# %% Function
def coordinates(shape, pixel=1., *, ndims=1, center=True, grid=False, origin=0., cuda=False) :
    '''
    This function will return a ndarray corresponding to the coordinates array of the input.
    
    Parameters
    ----------
    shape : int or tuple or np.ndarray
        Describes the shape of the coordinates.
        If int, all dimensions will have this value.
        If tuple, corresponds to the shape.
    pixel : float, tuple(float)
        Says what is the size of one bin, default is 1.
        If tuple, corresponds to the pixel for each dimension.
    ndims : int
        Number of dimensions if input shape is int.
        Will be overriden by any other input dimension.
    center : bool or tuple(bool)
        True to center origin, False to start by origin.
        If tuple, defines how to place origin for each dimension.
    grid : bool
        True to return full meshgrid for each dimensions, False for broadcastable arrays.
    origin : float or tuple(float)
        Value of origin (default 0).
        If tuple, value for each dimension.
    cuda : bool
        True to apply cuda.

    Returns
    -------
    coords : tuple(np.ndarray)
        return a tuple of arrays for each dimensions.

    Raises
    ------
    ValueError
        if input or passed as tuples that do not have the good number of dimensions.

    Examples
    --------
    >>> from arrlp import coordinates
    ...
    >>> array = np.ones((5,4))
    >>> coordinates = coordinates(array, pixel=10, center=(True, False), grid=False, origin=(0,3))
    ... array([[-20.], [-10.], [0.], [10.], [20.]]), array([[3., 13., 23., 33.]])
    '''
    xp = get_xp(cuda)

    # Manage shape argument
    if isinstance(shape, int) :
        shape = (shape,) * ndims
    elif isinstance(shape, tuple) :
        pass
    else :
        shape = shape.shape

    # Correct ndims
    ndims = len(shape)

    # Correct input dimensions
    pixel = iterable(pixel, ndims)
    center = iterable(center, ndims)
    origin = iterable(origin, ndims)

    #looping on shape
    coords = []
    for dim, (n, pix, cent, orig) in enumerate(zip(shape, pixel, center, origin)) :
        coord = n2coord(n, pixel=pix, center=cent, origin=orig, xp=xp)
        reshape = np.ones(ndims, dtype=int)
        reshape[dim] = len(coord)
        coord = coord.reshape(tuple(reshape))
        coords.append(coord)

    #Meshgrid
    if grid and ndims != 1 :
        return xp.meshgrid(*[coord.squeeze() for coord in coords], indexing='ij')
    else :
        return coords



def iterable(arg, ndims) :
    try :
        if len(arg) != ndims :
            raise ValueError('Argument does not have good number of dimensions')
        else :
            return arg
    except TypeError :
        return (arg,) * ndims



def n2coord(n, pixel=1., center=False, origin=0., xp=np) :
    if center :
        start, stop = -(n - 1) / 2,(n - 1) / 2
    else :
        start, stop = 0, n - 1
    return xp.linspace(start, stop, n) * pixel + origin



# %% Test function run
if __name__ == "__main__":
    from corelp import test
    test(__file__)