import numpy as np
import multiprocessing
from colour import *
from colour.models import eotf_inverse_sRGB
from colour.recovery import coefficients_Jakob2019
from gsoc_common import model_sd, D65, D65_xy, plot_comparison

# Solve for a specific RGB color
def optimize_RGB(linear_RGB, coeffs_0):
    coeffs, error = coefficients_Jakob2019(
        linear_RGB,
        RGB_COLOURSPACES["sRGB"],
        dimensionalise=False,
        coefficients_0=coeffs_0
    )

    if error > 1e-6:
        with open("out/bad.txt", "a") as fd:
            fd.write("\n")
            fd.write("linear_RGB=%s\n" % linear_RGB)
            fd.write("XYZ=%s\n" % XYZ)
            fd.write("coeffs_0=%s\n" % coeffs_0)
            fd.write("coeffs=%s, error=%g\n" % (coeffs, error))

    # A low budget graph to quickly see what's going on
    log_error = np.log10(error)
    bars = "|" * max(0, int(5 * (log_error + 9)))

    print("%12.5g %12.5g %12.5g %5.3f %s" % (*linear_RGB, log_error, bars))
    return coeffs

# Solve for all lightness values of a fully saturated RGB color
def optimize_chromaticity(linear_RGB):
    # alpha's aren't spaced equally (see the article)
    def smoothstep(x):
        return x**2 * (3 - 2 * x)

    steps = np.arange(0, LIGHTNESS_STEPS)
    alphas = smoothstep(smoothstep(steps / LIGHTNESS_STEPS))

    i_mid = LIGHTNESS_STEPS // 5
    coeffs_mid = optimize_RGB(linear_RGB * alphas[i_mid], (0, 0, 0))

    coeffs_0 = coeffs_mid
    for i in range(i_mid + 1, LIGHTNESS_STEPS):
        coeffs_0 = optimize_RGB(linear_RGB * alphas[i], coeffs_0)

    coeffs_0 = coeffs_mid
    for i in reversed(range(0, i_mid)):
        coeffs_0 = optimize_RGB(linear_RGB * alphas[i], coeffs_0)



# This program runs XYZ_to_sd_Jakob2019 converges on a large number of inputs,
# covering an entire RGB gamut, to see if it'll diverge somewhere.
if __name__ == "__main__":
    CHROMA_STEPS = 8
    LIGHTNESS_STEPS = 64

    with open("out/bad.txt", "w") as fd:
        fd.write("Going through %dx%dx%d cubes\n"
                 % (LIGHTNESS_STEPS, CHROMA_STEPS, CHROMA_STEPS))

    args = []
    for A in np.linspace(0, 1, CHROMA_STEPS):
        for B in np.linspace(0, 1, CHROMA_STEPS):
            for RGB in [np.array([1, A, B]),
                        np.array([A, 1, B]),
                        np.array([A, B, 1])]:
                args.append(RGB)

    pool = multiprocessing.Pool()
    pool.map(optimize_chromaticity, args)