import numpy as np

# UTILITY {{{1

F = lambda a: np.array(a, np.float64)
A = lambda *args: F(args)
and1 = lambda a: A(*a.T, np.ones_like(a.T[0])).T  # appends a column of 1s
# (works for both vectors and matrices thanks to some transposition magic)


def either1d2d(fun):
    from functools import wraps

    @wraps(fun)
    def wrap(a, *args, **kwargs):
        a = np.asfarray(a)
        assert a.ndim in (1, 2), f"expected a 1D or 2D array, got {a.ndim}"
        return fun(a, *args, **kwargs)

    return wrap


# DATA {{{1


d65 = A(95.0489, 100.0000, 108.8840)
d50 = A(96.4212, 100.0000, +82.5188)

mat_alt_xyz = A(  # for SRLAB2
    (0.320530, 0.636920, 0.042560),
    (0.161987, 0.756636, 0.081376),
    (0.017228, 0.108660, 0.874112),
)

mat_alt_lab = A(  # for SRLAB2
    (37.0950, +62.9054, -0.0008),
    (663.4684, -750.5078, +87.0328),
    (63.9569, +108.4576, -172.4152),
)

lab_magic = A(
    [0, 116, 0],
    [500, -500, 0],
    [0, 200, -200],
)

lab_offset = A(-16, 0, 0)

bradford = A(
    [+0.8951, +0.2664, -0.1614],
    [-0.7502, +1.7135, +0.0367],
    [+0.0389, -0.0685, +1.0296],
)

mat_BFD = bradford

mat_97 = A(
    [+0.8562, +0.3372, -0.1934],
    [-0.8360, +1.8327, +0.0033],
    [+0.0357, -0.0469, +1.0112],
)

mat_02 = A(
    [+0.7328, +0.4296, -0.1624],
    [-0.7036, +1.6975, +0.0061],
    [+0.0030, +0.0136, +0.9834],
)

mat_16 = A(
    [+0.401288, +0.650173, -0.051461],
    [-0.250268, +1.204414, +0.045854],
    [-0.002079, +0.048952, +0.953127],
)

# REC709_PRI
sRGB_primaries = A(  # strangely, these seem to be missing a digit of significance.
    # actually, this might be ITU-R BT.709-5 which is naturally less accurate?
    [+0.64000, +0.33000],  # +0.03000
    [+0.30000, +0.60000],  # +0.10000
    [+0.15000, +0.06000],  # +0.79000
    [+0.31270, +0.32900],  # +0.35830
)

# AP0
AP0_primaries = A(
    [+0.73470, +0.26530],  # +0.00000
    [+0.00000, +1.00000],  # +0.00000
    [+0.00010, -0.07700],  # +1.07690
    [+0.32168, +0.33767],  # +0.34065
)

# AP1
AP1_primaries = A(
    [+0.71300, +0.29300],  # -0.00600
    [+0.16500, +0.83000],  # +0.00500
    [+0.12800, +0.04400],  # +0.82800
    [+0.32168, +0.33767],  # +0.34065
)

ref_XYZ = A(  # sRGB -> XYZ (still needs to be converted from D65)
    [506752 / 1228815, 87881 / 245763, 12673 / 70218],  # sum: 3127/3290
    [87098 / 409605, 175762 / 245763, 12673 / 175545],  # sum: 1/1
    [7918 / 409605, 87881 / 737289, 1001167 / 1053270],  # sum: 3583/3290
)

_Zify = A(  # goes *after* input vector, not before.
    [1, 0, -1],
    [0, 1, -1],
    [0, 0, +1],
)

# DATA (INVERSES) {{{1

inv_alt_xyz = np.linalg.inv(mat_alt_xyz)
inv_alt_lab = np.linalg.inv(mat_alt_lab)
inv_lab_magic = np.linalg.inv(lab_magic)

inv_BFD = np.linalg.inv(mat_BFD)  # TODO: use this!
inv_97 = np.linalg.inv(mat_97)  # TODO: rename to _CAT97?
inv_02 = np.linalg.inv(mat_02)  # TODO: rename to _CAT02?
inv_16 = np.linalg.inv(mat_16)


# FUNCTIONS {{{1


@either1d2d
def cartesian_to_polar(a):
    return A(a.T[0], np.hypot(a.T[1], a.T[2]), np.arctan2(a.T[2], a.T[1])).T


@either1d2d
def polar_to_cartesian(a):
    return A(a.T[0], a.T[1] * np.cos(a.T[2]), a.T[1] * np.sin(a.T[2])).T


def rgb2lin(a):
    a = np.asfarray(a)
    return np.where(a > 0.04045, ((a + 0.055) / 1.055) ** 2.4, a / 12.92)


def lin2rgb(a):
    a = np.asfarray(a)
    return np.where(a > 0.04045 / 12.92, (a ** (1 / 2.4) * 1.055) - 0.055, a * 12.92)


@either1d2d
def xyY_2_XYZ(xyY):
    # transposition allows this function to be vectorized (over vectors).
    x, y, Y = xyY.T
    y = np.maximum(y, 1e-10)
    return (Y / y * A(x, y, 1.0 - x - y)).T


@either1d2d
def lab_to_xyz(lab, *, ref):
    ijk = (lab - lab_offset) @ inv_lab_magic.T
    xyz = np.where(ijk > np.cbrt(0.008856), ijk**3, (ijk - 16 / 116) / 7.787)
    xyz = xyz * ref
    return xyz


@either1d2d
def xyz_to_lab(xyz, *, ref):
    xyz = xyz / ref
    # absolute value just to shut up numpy:
    ijk = np.where(xyz > 0.008856, np.cbrt(np.abs(xyz)), 7.787 * xyz + 16 / 116)
    lab = ijk @ lab_magic.T + lab_offset
    return lab


# print(xyz_to_lab((30, 50, 20), ref=d50))
# print("[ 76.06926101 -58.04284385  34.04319485]")
# print(lab_to_xyz(( 76.06926101, -58.04284385, 34.04319485), ref=d50))
# stop


@either1d2d
def xyz_to_srlab2(xyz):
    rgb = xyz @ inv_XYZ.T
    xyz = rgb @ mat_alt_xyz.T
    ijk = np.where(
        xyz <= 216.0 / 24389.0, 24389.0 / 2700.0 * xyz, 1.16 * np.cbrt(xyz) - 0.16
    )
    lab = ijk @ mat_alt_lab.T
    return lab


@either1d2d
def srlab2_to_xyz(lab):
    ijk = lab @ inv_alt_lab.T
    xyz = np.where(ijk <= 0.08, 2700.0 / 24389.0 * ijk, ((ijk + 0.16) / 1.16) ** 3)
    rgb = xyz @ inv_alt_xyz.T
    xyz = rgb @ mat_XYZ.T
    return xyz


# print(srlab2_to_xyz(xyz_to_srlab2((0.3, 0.4, 0.5))))
# print(srlab2_to_xyz(xyz_to_srlab2((0.3, 0.4, 0.5))))
# stop


def ref_to_uv(ref):
    # NOTE: this assumes that ref is valid i.e. won't cause a division by zero.
    xt, yt, zt = ref  # tristimulus values (CIE 1931 XYZ)
    den = xt + yt + zt  # denominator
    xc, yc = xt / den, yt / den  # chromaticity values (still CIE 1931)
    den = xc + 15 * yc + 3 * (1 - xc - yc)
    u, v = 4 * xc / den, 6 * yc / den  # u, v coordinates (CIE 1960 UCS)
    return u, v


@either1d2d
def xyz_to_uvw(xyz, *, ref):
    u0, v0 = ref_to_uv(ref)
    x, y, z = xyz.T
    den = np.maximum(x + 15 * y + 3 * z, 1e-12)
    u = 4 * x / den
    v = 6 * y / den
    W = 25 * np.cbrt(y) - 17
    U = 13 * W * (u - u0)
    V = 13 * W * (v - v0)
    return A(U, V, W).T


@either1d2d
def uvw_to_xyz(uvw, *, ref):
    u0, v0 = ref_to_uv(ref)
    U, V, W = uvw.T
    Y = ((W + 17) / 25) ** 3
    u = U / (13 * W) + u0
    v = V / (13 * W) + v0
    den = Y * 6 / np.maximum(v, 1e-12)
    X = u / 4 * den
    Z = (den - X - 15 * Y) / 3
    return A(X, Y, Z).T


def makemat(prim, normalize=True):
    primz, white = and1(prim[:3]) @ _Zify, xyY_2_XYZ(and1(prim[3]))
    rescalant = np.cross(primz[((1, 2, 0),)], primz[((2, 0, 1),)]) @ white
    mat = (primz * rescalant[:, None]).T
    # normalize such that max inputs (1.0, 1.0, 1.0) yields (any, 1.0, any)
    return mat / np.sum(mat[1, :]) if normalize else mat


# GENERATED DATA {{{1

mat_XYZ = makemat(sRGB_primaries)
inv_XYZ = np.linalg.inv(mat_XYZ)