"""
Oklab Colourspace
=================

Define the *Oklab* colourspace transformations.

-   :func:`colour.XYZ_to_Oklab`
-   :func:`colour.Oklab_to_XYZ`

References
----------
-   :cite:`Ottosson2020` : Ottosson, B. (2020). A perceptual color space for
    image processing. Retrieved December 24, 2020, from
    https://bottosson.github.io/posts/oklab/
"""

from __future__ import annotations

from functools import partial

import numpy as np

from colour.algebra import spow
from colour.hints import (  # noqa: TC001
    Domain1,
    NDArrayFloat,
    Range1,
)
from colour.models import Iab_to_XYZ, XYZ_to_Iab

__author__ = "Colour Developers"
__copyright__ = "Copyright 2013 Colour Developers"
__license__ = "BSD-3-Clause - https://opensource.org/licenses/BSD-3-Clause"
__maintainer__ = "Colour Developers"
__email__ = "colour-developers@colour-science.org"
__status__ = "Production"

__all__ = [
    "MATRIX_1_XYZ_TO_LMS",
    "MATRIX_1_LMS_TO_XYZ",
    "MATRIX_2_LMS_TO_LAB",
    "MATRIX_2_LAB_TO_LMS",
    "XYZ_to_Oklab",
    "Oklab_to_XYZ",
]

MATRIX_1_XYZ_TO_LMS: NDArrayFloat = np.array(
    [
        [0.8189330101, 0.3618667424, -0.1288597137],
        [0.0329845436, 0.9293118715, 0.0361456387],
        [0.0482003018, 0.2643662691, 0.6338517070],
    ]
)
"""*CIE XYZ* tristimulus values to normalised cone responses matrix."""

MATRIX_1_LMS_TO_XYZ: NDArrayFloat = np.linalg.inv(MATRIX_1_XYZ_TO_LMS)
"""Normalised cone responses to *CIE XYZ* tristimulus values matrix."""

MATRIX_2_LMS_TO_LAB: NDArrayFloat = np.array(
    [
        [0.2104542553, 0.7936177850, -0.0040720468],
        [1.9779984951, -2.4285922050, 0.4505937099],
        [0.0259040371, 0.7827717662, -0.8086757660],
    ]
)
"""Normalised cone responses to *Oklab* colourspace matrix."""

MATRIX_2_LAB_TO_LMS: NDArrayFloat = np.linalg.inv(MATRIX_2_LMS_TO_LAB)
"""*Oklab* colourspace to normalised cone responses matrix."""


def XYZ_to_Oklab(XYZ: Domain1) -> Range1:
    """
    Convert from *CIE XYZ* tristimulus values to *Oklab* colourspace.

    Parameters
    ----------
    XYZ
        *CIE XYZ* tristimulus values.

    Returns
    -------
    :class:`numpy.ndarray`
        *Oklab* colourspace array.

    Notes
    -----
    +------------+-----------------------+-----------------+
    | **Domain** | **Scale - Reference** | **Scale - 1**   |
    +============+=======================+=================+
    | ``XYZ``    | 1                     | 1               |
    +------------+-----------------------+-----------------+

    +------------+-----------------------+-----------------+
    | **Range**  | **Scale - Reference** | **Scale - 1**   |
    +============+=======================+=================+
    | ``Lab``    | 1                     | 1               |
    +------------+-----------------------+-----------------+

    -   Input *CIE XYZ* tristimulus values must be adapted to
        *CIE Standard Illuminant D Series* *D65*.

    References
    ----------
    :cite:`Ottosson2020`

    Examples
    --------
    >>> XYZ = np.array([0.20654008, 0.12197225, 0.05136952])
    >>> XYZ_to_Oklab(XYZ)  # doctest: +ELLIPSIS
    array([ 0.5163401...,  0.154695 ...,  0.0628957...])
    """

    return XYZ_to_Iab(
        XYZ,
        partial(spow, p=1 / 3),
        MATRIX_1_XYZ_TO_LMS,
        MATRIX_2_LMS_TO_LAB,
    )


def Oklab_to_XYZ(Lab: Domain1) -> Range1:
    """
    Convert from *Oklab* colourspace to *CIE XYZ* tristimulus values.

    Parameters
    ----------
    Lab
        *Oklab* colourspace array.

    Returns
    -------
    :class:`numpy.ndarray`
        *CIE XYZ* tristimulus values.

    Notes
    -----
    +------------+-----------------------+-----------------+
    | **Domain** | **Scale - Reference** | **Scale - 1**   |
    +============+=======================+=================+
    | ``Lab``    | 1                     | 1               |
    +------------+-----------------------+-----------------+

    +------------+-----------------------+-----------------+
    | **Range**  | **Scale - Reference** | **Scale - 1**   |
    +============+=======================+=================+
    | ``XYZ``    | 1                     | 1               |
    +------------+-----------------------+-----------------+

    References
    ----------
    :cite:`Ottosson2020`

    Examples
    --------
    >>> Lab = np.array([0.51634019, 0.15469500, 0.06289579])
    >>> Oklab_to_XYZ(Lab)  # doctest: +ELLIPSIS
    array([ 0.2065400...,  0.1219722...,  0.0513695...])
    """

    return Iab_to_XYZ(
        Lab,
        partial(spow, p=3),
        MATRIX_2_LAB_TO_LMS,
        MATRIX_1_LMS_TO_XYZ,
    )