"""
Color utility functions for accessing and interpolating color data.

This module provides functions for loading and interpolating color matching functions
and illuminant data to arbitrary wavelength ranges. It serves as an interface between
the color transformation functions and the underlying color data.
"""

import jax.numpy as jnp

from jaxlayerlumos.colors import color_matching_functions
from jaxlayerlumos.colors import illuminants


def get_cmfs(bx, str_color_space="cie1931"):
    """
    Get color matching functions interpolated to target wavelengths.
    
    This function loads color matching functions for the specified color space
    and interpolates them to the target wavelength range.
    
    Args:
        bx (jnp.ndarray): Target wavelengths in meters.
        str_color_space (str, optional): Color space name. Currently only supports
                                        "cie1931". Defaults to "cie1931".
    
    Returns:
        jnp.ndarray: Color matching functions with shape (n_wavelengths, 4) containing
                    [wavelength, x_bar, y_bar, z_bar] columns.
    
    Raises:
        ValueError: If the color space is not supported.
    
    Note:
        The target wavelengths must be within the range of the available color
        matching function data.
    """
    assert isinstance(bx, jnp.ndarray)

    if str_color_space == "cie1931":
        cmfs = color_matching_functions.cmfs_cie1931
    else:
        raise ValueError

    cmfs = jnp.array(cmfs)
    assert cmfs.shape[1] == 4
    assert jnp.min(cmfs[:, 0]) <= jnp.min(bx)
    assert jnp.max(bx) <= jnp.max(cmfs[:, 0])

    cmfs_interpolated = jnp.concatenate(
        [
            bx[..., jnp.newaxis],
            jnp.interp(bx, cmfs[:, 0], cmfs[:, 1])[..., jnp.newaxis],
            jnp.interp(bx, cmfs[:, 0], cmfs[:, 2])[..., jnp.newaxis],
            jnp.interp(bx, cmfs[:, 0], cmfs[:, 3])[..., jnp.newaxis],
        ],
        axis=1,
    )

    return cmfs_interpolated


def get_illuminant(bx, str_illuminant="d65"):
    """
    Get illuminant data interpolated to target wavelengths.
    
    This function loads illuminant data for the specified illuminant and interpolates
    it to the target wavelength range.
    
    Args:
        bx (jnp.ndarray): Target wavelengths in meters.
        str_illuminant (str, optional): Illuminant name. Currently only supports
                                       "d65". Defaults to "d65".
    
    Returns:
        jnp.ndarray: Illuminant data with shape (n_wavelengths, 2) containing
                    [wavelength, relative_spectral_power] columns.
    
    Raises:
        ValueError: If the illuminant is not supported.
    
    Note:
        The target wavelengths must be within the range of the available illuminant data.
    """
    assert isinstance(bx, jnp.ndarray)

    if str_illuminant == "d65":
        illuminant = illuminants.illuminant_d65
    else:
        raise ValueError

    illuminant = jnp.array(illuminant)
    assert illuminant.shape[1] == 2
    assert jnp.min(illuminant[:, 0]) <= jnp.min(bx)
    assert jnp.max(bx) <= jnp.max(illuminant[:, 0])

    illuminant_interpolated = jnp.concatenate(
        [
            bx[..., jnp.newaxis],
            jnp.interp(bx, illuminant[:, 0], illuminant[:, 1])[..., jnp.newaxis],
        ],
        axis=1,
    )

    return illuminant_interpolated
