"""
Calculate reflection curves.

Approaches comes from Scott Burns.
http://scottburns.us/matlab-octave-and-python-source-code-for-refl-recon-chrom-adapt/
"""
import sys
import os
import math
import plotly.graph_objects as go

sys.path.insert(0, os.getcwd())

from coloraide import algebra as alg
from coloraide.everything import ColorAll as Color
from coloraide.cmfs import CIE_1931_2DEG as CMF
from coloraide.illuminants import D65

targets = ['#ff0000', '#00ff00', '#0000ff', '#00ffff', '#ff00ff', '#ffff00', '#ffffff']
space = 'srgb'
# With the current settings, method2 requires an epsilon of 1e-14, while method 3 can use 1e-15
# More iterations may be required for for higher step resolutions.
# For our purposes, method 3 delivers colors with estimated reflectance below 1.
# Sum of RGB is close to 1.
START = 380
STEP = 10
END = 750 + STEP
MAX_ITERS = 50
METHOD = 3
EPSILON = 1e-15 if METHOD == 3 else 1e-14

if space == 'srgb':
    from coloraide.spaces.srgb_linear import XYZ_TO_RGB
elif space == 'display-p3':
    from coloraide.spaces.display_p3_linear import XYZ_TO_RGB
elif space == 'a98-rgb':
    from coloraide.spaces.a98_rgb_linear import XYZ_TO_RGB
elif space == 'rec2020':
    XYZ_TO_RGB = [
        [1.7249510969855462, -0.36225403975096726, -0.2545649705534804],
        [-0.6692669194491035, 1.6225096012804534, 0.012487041279557484],
        [0.018196517838035248, -0.04443708810957042, 0.9431475044937067]
    ]
else:
    raise ValueError('Cannot perfrom operation on this space')


A = [[], [], []]
for r in range(START, END, STEP):
    A[0].append(CMF[r][0])
    A[1].append(CMF[r][1])
    A[2].append(CMF[r][2])

W = [D65[r] for r in range(START, END, STEP)]

value = alg.matmul(A, alg.diag(W))
w = alg.vdot(A[1], W)

# Illuminant referenced CMFs
T = alg.divide(value, w)

# Make in relation to the RGB color space
TRGB = alg.matmul(XYZ_TO_RGB, T)

print('TRGB')
alg.pprint(TRGB)

# Array of finite differencing constants
size = len(W)
fill = {-1: -2, 0: 4, 1: -2}
D = [[fill.get(i - j, 0) for j in range(size)] for i in range(size)]
D[0][0] = 2
D[size - 1][size - 1] = 2

# Vectorize some functions we will use.
# If any of the functions overflow, return zero.
# `exp` is most likely to do this, the others may not unless pushed too hard.
def _tanh(x):
    """Clamp the results to zero if there is overflow."""

    try:
        return math.tanh(x)
    except OverflowError:
        return 0


def _exp(x):
    """Clamp the results to zero if there is overflow."""

    try:
        return math.exp(x)
    except OverflowError:
        return 0


def _sech2(x):
    """Clamp the results to zero if there is overflow."""

    try:
        return 1 / (math.cosh(x) ** 2)
    except OverflowError:
        return 0


tanh = alg.vectorize2(_tanh, params=1)
sech2 = alg.vectorize2(_sech2, params=1)
exp = alg.vectorize2(_exp, params=1)


def method_2(d, cmfs_w, XYZ_w, max_iterations=MAX_ITERS, tolerance=EPSILON):
    """
    Estimate reflectance curve.

    Computes a reconstructed spectral reflectance curve from
    an illuminant-referenced triplet of tristimulus values
    representing a real color (within the spectral locus).

    Parameters
    ----------
    `d`: nxn array of finite differencing constants, obtained
           from function `method_2_prep.py`. (n is the number
           of wavelength bins (rows) used in the CMFs.)
    `cmfs_w`: nx3 array of illuminant-w-referenced CMFs,
           obtained from function `method_2_prep.py`.
    `XYZ_w`: a three element vector of illum-w-referenced
           tristimulus values in the 0 <= Y <= 1 convention range.
    max_iterations: upper limit on iterations (defaults to 20).
    tolerance: determines when the linear equations for a
           stationary point are satisfied (default is 1e-8).

    Returns
    -------
    nx1 vector of spectral reflectance values (or zeros if failure).

    Notes
    -----
    The input parameters d and `cmfs_w` can be generated by function
    `method_2_prep.py`. For a given illuminant, `method_2_prep`
    needs to be called only once. Multiple calls to `method_2`
    may then follow for various tristimulus values. This is a
    typical use case, and avoids repeated computation of d and
    `cmfs_w`.

    The XYZ_w supplied to `method_2` must fall strictly within the
    spectral locus (i.e., be a real color), or the method will fail
    and return a vector of zeros.

    The reflectance curve generated will always have positive
    values, but may have some values >1 (which can be interpreted
    as an emissive source divided by the illuminant). To obtain
    reflectance curves strictly in the range 0 to 1, representing
    valid object reflectances, use `method_3.py` instead.

    `method_2` is based on "method 2" of Reference [1]. It also
    corresponds to method "LLSS" (Least Log Slope Squared) of
    Reference [2].

    References
    ----------
    [1] Burns SA. Numerical methods for smoothest reflectance
        reconstruction. Color Research & Application, Vol 45,
        No 1, 2020, pp 8-21.
    [2] Generating Reflectance Curves from sRGB Triplets, 2015
        http://scottburns.us/reflectance-curves-from-srgb/#LLSS

    Acknowledgements
    ----------------
    Thanks to Adam J. Burns and Mark (`kram1032`) for assistance
    with translating from MATLAB to Python.

    """

    n = len(cmfs_w)
    z = alg.zeros(n)
    lamb = alg.zeros(3)
    for _ in range(max_iterations):
        r = exp(z)
        dr = alg.diag(r)
        dra = alg.matmul(dr, cmfs_w)
        v = alg.matmul(dra, lamb)
        p1 = alg.add(alg.matmul(d, z), v)
        p2 = alg.subtract(alg.matmul(alg.transpose(cmfs_w), r), XYZ_w)
        F = p1 + p2
        p1 = alg.add(d, alg.diag(v))
        p2 = dra
        p3 = alg.transpose(dra)
        p4 = alg.zeros((3, 3))
        for x1, y1 in zip(p1, p2):
            x1.extend(y1)
        for x1, y1 in zip(p3, p4):
            x1.extend(y1)
        p1.extend(p3)
        J = p1
        try:
            try:
                delta = alg.solve(J, alg.multiply(F, -1))
            except Exception:
                pi = alg.pinv(J)
                delta = alg.dot(pi, alg.multiply(F, -1))
        except Exception:
            print('Ill-conditioned or singular linear system detected.')
            print(f'Check to make sure XYZ_w {XYZ_w} is within the spectral locus.')
            return alg.zeros(n)
        z = alg.add(z, delta[:n])
        lamb = alg.add(lamb, delta[n:])
        if alg.all([abs(f) < tolerance for f in F]):
            return exp(z)
    print('Maximum number of iterations reached.')
    print('Check to make sure XYZ_w is within the spectral locus.')
    return alg.zeros(n)


def method_3(d, cmfs_w, XYZ_w, max_iterations=MAX_ITERS, tolerance=EPSILON):
    """
    Estimate reflectance curve.

    Computes a reconstructed spectral reflectance curve from
    an illuminant-referenced triplet of tristimulus values
    representing an object color (within the object color solid).

    Parameters
    ----------
    `d`: nxn array of finite differencing constants, obtained
           from function `method_3_prep.py`. (n is the number
           of wavelength bins (rows) used in the CMFs.)
    `cmfs_w`: nx3 array of illuminant-w-referenced CMFs,
           obtained from function `method_3_prep.py`.
    `XYZ_w`: a three element vector of illum-w-referenced
           tristimulus values in the 0 <= Y <= 1 convention range.
    max_iterations: upper limit on iterations (defaults to 20).
    tolerance: determines when the linear equations for a
           stationary point are satisfied (default is 1e-8).

    Returns
    -------
    nx1 vector of spectral reflectance values (or zeros if failure).

    Notes
    -----
    The input parameters d and `cmfs_w` can be generated by function
    `method_3_prep.py`. For a given illuminant, `method_3_prep`
    needs to be called only once. Multiple calls to `method_3`
    may then follow for various tristimulus values. This is a
    typical use case, and avoids repeated computation of d and
    `cmfs_w`.

    The XYZ_w supplied to `method_3` must fall strictly within the
    object color solid (i.e., be a valid object color that can be
    represented by a reflectance curve strictly between 0 and 1),
    or the method will fail and return a vector of zeros.

    The reflectance curve generated will always have values in
    the range 0 to 1, corresponding to valid object colors. If
    it is desired to expand the applicability of the method to
    any real color within the spectral locus (including emissive
    colors), use `method_2.py` instead.

    `method_3` is based on "method 3" of Reference [1]. It also
    corresponds to method "LHTSS" (Least Hyperbolic Tangent Slope
    Squared) of Reference [2].

    References
    ----------
    [1] Burns SA. Numerical methods for smoothest reflectance
        reconstruction. Color Research & Application, Vol 45,
        No 1, 2020, pp 8-21.
    [2] Generating Reflectance Curves from sRGB Triplets, 2015
        http://scottburns.us/reflectance-curves-from-srgb/#LHTSS

    Acknowledgements
    ----------------
    Thanks to Adam J. Burns and Mark (`kram1032`) for assistance
    with translating from MATLAB to Python.

    """

    n = len(cmfs_w)  # number of rows
    z = alg.zeros(n)
    lam = alg.zeros(3)
    for _ in range(max_iterations):
        r = alg.divide(alg.add(tanh(z), 1), 2)
        s2 = sech2(z)
        d1 = alg.diag(alg.divide(s2, 2))
        d1a = alg.dot(d1, cmfs_w)
        d2 = alg.diag(alg.multiply(s2, tanh(z)))

        p1 = alg.add(alg.dot(d, z), alg.dot(d1a, lam))
        p2 = alg.subtract(alg.dot(alg.transpose(cmfs_w), r), XYZ_w)
        f = p1 + p2

        p1 = alg.subtract(d, alg.diag(alg.dot(alg.dot(d2, cmfs_w), lam)))
        p2 = d1a
        p3 = alg.transpose(d1a)
        p4 = alg.zeros((3, 3))
        for x1, y1 in zip(p1, p2):
            x1.extend(y1)
        for x1, y1 in zip(p3, p4):
            x1.extend(y1)
        p1.extend(p3)
        j = p1

        try:
            try:
                delta = alg.solve(j, alg.multiply(-1, f))
            except Exception:
                pi = alg.pinv(j)
                delta = alg.dot(pi, alg.multiply(f, -1))
        except Exception:
            print('Ill-conditioned or singular linear system detected.')
            print(f'Check to make sure XYZ_w {XYZ_w} is within the spectral locus.')
            return alg.zeros(n)
        z = alg.add(z, delta[0:n])
        lam = alg.add(lam, delta[n:n+3])
        if alg.all([abs(_f) < tolerance for _f in f]):
            return alg.divide(alg.add(tanh(z), 1), 2)
    print('Maximum number of iterations reached.')
    print('Check to make sure XYZ_w is within the object color solid.')
    return alg.zeros(n)


# Estimate curves for the specified colors red, green, blue, cyan, and magenta.
rho = []
for color in targets:
    c = Color(color).convert('srgb-linear').coords()
    METHOD = 2 if c == [1, 1, 1] else 3
    EPSILON = 1e-15 if METHOD == 3 else 1e-14
    if METHOD == 3:
        r = method_3(D, alg.transpose(TRGB), c)
    elif METHOD == 2:
        r = method_2(D, alg.transpose(TRGB), c)
    else:
        raise ValueError('Unrecognized method')
    rho.append(r)

print('=== CMFs D65 ===')
print('==== X ====')
alg.pprint(T[0])
print('==== Y ====')
alg.pprint(T[1])
print('==== Z ====')
alg.pprint(T[2])

# Setup plot for results
fig = go.Figure(
    layout={
        'title': 'Primary Color Reflectance Curves',
        'xaxis_title': 'Wavelength',
        'yaxis_title': 'Reflection',
        'width': 800,
        'height': 600
    }
)

print('=== Reflectance Curves ===')
for target, r in zip(targets, rho):
    print(f'==== Curve for {target} ====')
    alg.pprint(r)

    fig.add_traces(data=go.Scatter(
        x=list(range(START, END, STEP)),
        y=r,
        mode="lines",
        line={'color': target, 'width': 4, 'shape': 'spline'},
        showlegend=False
    ))

fig.show('browser')
