import colour
import colour_datasets
import matplotlib.pyplot as plt

SHAPE = colour.SpectralShape(400, 700, 10)

DATABASE_CAMERA_SENSITIVITIES = colour_datasets.load(
    "Camera Spectral Sensitivity Database - Jiang et al. (2013)"
)
DATABASE_AMPAS_CAMERA_SENSITIVITIES = colour_datasets.load(
    "RAW to ACES Utility Data - Dyer et al. (2017)"
)

DATABASE = {
    camera: msds.copy().align(SHAPE)
    for camera, msds in DATABASE_AMPAS_CAMERA_SENSITIVITIES["camera"].items()
}

(RED_W, GREEN_W, BLUE_W), (RED_V, GREEN_V, BLUE_V) = colour.recovery.PCA_Jiang2013(
    DATABASE, 6, True
)

plt.plot(np.array(RED_W) * np.array(RED_V))
plt.plot(np.array(GREEN_W) * np.array(GREEN_V))
plt.plot(np.array(BLUE_W) * np.array(BLUE_V))
plt.show()

K = 0.01
ILLUMINANT = colour.SDS_ILLUMINANTS["D50"].copy().align(SHAPE).normalise()
CAMERA_SENSITIVITIES = (
    DATABASE_CAMERA_SENSITIVITIES["Canon 1DMarkIII"].copy().align(SHAPE)
)

MSDS_REFLECTANCES = colour.characterisation.read_training_data_rawtoaces_v1()
REFLECTANCES = np.transpose(MSDS_REFLECTANCES.copy().align(SHAPE).values)[
    np.newaxis, ...
]

RGB_REFLECTANCES_REFERENCE = np.squeeze(
    colour.msds_to_XYZ(
        REFLECTANCES,
        method="Integration",
        cmfs=CAMERA_SENSITIVITIES,
        illuminant=ILLUMINANT,
        k=K,
        shape=SHAPE,
    )
)

colour.plotting.plot_multi_colour_swatches(
    colour.cctf_encoding(RGB_REFLECTANCES_REFERENCE),
    columns=REFLECTANCES.shape[1] // 10,
)

colour.plotting.plot_multi_sds(MSDS_REFLECTANCES, legend=False)

BASIS_FUNCTIONS = np.transpose(np.array([RED_W, GREEN_W, BLUE_W]))

MSDS_CAMERA = colour.recovery.RGB_to_msds_camera_sensitivities_Jiang2013(
    RGB_REFLECTANCES_REFERENCE,
    ILLUMINANT,
    MSDS_REFLECTANCES,
    basis_functions=BASIS_FUNCTIONS,
)

plt.plot(
    CAMERA_SENSITIVITIES.signals[CAMERA_SENSITIVITIES.labels[0]].values,
    label="Red Reference",
)
plt.plot(MSDS_CAMERA.signals["red"].values, "--", label="Red Reconstructed")
plt.plot(
    CAMERA_SENSITIVITIES.signals[CAMERA_SENSITIVITIES.labels[1]].values,
    label="Green Reference",
)
plt.plot(MSDS_CAMERA.signals["green"].values, "--", label="Green Reconstructed")
plt.plot(
    CAMERA_SENSITIVITIES.signals[CAMERA_SENSITIVITIES.labels[2]].values,
    label="Blue Reference",
)
plt.plot(MSDS_CAMERA.signals["blue"].values, "--", label="Blue Reconstructed")
plt.legend()
plt.show()

RGB_REFLECTANCES_TEST = np.squeeze(
    colour.msds_to_XYZ(
        REFLECTANCES,
        method="Integration",
        cmfs=MSDS_CAMERA,
        illuminant=ILLUMINANT,
        k=K,
        shape=SHAPE,
    )
)

colour.plotting.plot_multi_colour_swatches(
    colour.cctf_encoding(RGB_REFLECTANCES_TEST),
    columns=REFLECTANCES.shape[1] // 10,
)

print(
    np.median(
        colour.delta_E(
            colour.convert(RGB_REFLECTANCES_TEST, "CIE XYZ", "CIE Lab"),
            colour.convert(RGB_REFLECTANCES_REFERENCE, "CIE XYZ", "CIE Lab"),
        )
    )
)

print(
    np.max(
        colour.delta_E(
            colour.convert(RGB_REFLECTANCES_TEST, "CIE XYZ", "CIE Lab"),
            colour.convert(RGB_REFLECTANCES_REFERENCE, "CIE XYZ", "CIE Lab"),
        )
    )
)