# coding: utf-8
import numpy as np

"""Color temperature correction script for Shade 3D.

Apply a color temperature correction to the rendering image.
"""

__author__ = "CASPAR003 <caspar003@caspar003.info>"
__date__ = "2016-08-12"

def color_temperature_correction(Tto, Tfrom = 6500.0):
    """Apply a color temperature correction to the image window.
    
    @param Tto Target color tempreture (unit: K).
    @param Tfrom Current color tempreture (unit: K, default: 6500K).
    
    Reference:
    T. Fujiwara. Color coordinate transformation. 2012-01.
    URL: http://w3.kcua.ac.jp/~fujiwara/infosci/colorspace/
    """
    # Validate arguments.
    if Tto < 1000.0 or 15000.0 < Tto:
        print "[Error] Argument Tto is invalid."
        return
    if Tfrom < 1000.0 or 15000.0 < Tfrom:
        print "[Error] Argument Tfrom is invalid."
        return
    
    # Get the rendering image and validate it.
    img = xshade.scene().rendering.image
    if img is None:
        print "[Error] The rendering image is None." 
        return
    if not img.has_image:
        print "[Error] The rendering image has no image." 
        return
    
    # sRGB vertex matrix N.
    N = np.array([
        [0.64, 0.30, 0.15],
        [0.33, 0.60, 0.06],
        [0.03, 0.10, 0.79]
    ])
    
    # Assemble RGB(Tfrom) -> XYZ color conversion matrix M0.
    u, v = blackbody_radiation_color_uv(Tfrom)
    x = 3.0 * u / (2.0 * u - 8.0 * v + 4.0)
    y = 2.0 * v / (2.0 * u - 8.0 * v + 4.0)
    XYZ_white = np.array([x/y, 1.0, (1.0 - x - y)/y])
    T_rgb = np.dot(np.linalg.inv(N), XYZ_white)
    T_rgb /= np.linalg.norm(T_rgb)
    M0 = np.dot(N, np.diag(T_rgb))
    
    # Assemble RGB(Tto) -> XYZ color conversion matrix M1.
    u, v = blackbody_radiation_color_uv(Tto)
    x = 3.0 * u / (2.0 * u - 8.0 * v + 4.0)
    y = 2.0 * v / (2.0 * u - 8.0 * v + 4.0)
    XYZ_white = np.array([x/y, 1.0, (1.0 - x - y)/y])
    T_rgb = np.dot(np.linalg.inv(N), XYZ_white)
    T_rgb /= np.linalg.norm(T_rgb)
    M1 = np.dot(N, np.diag(T_rgb))
    
    # Assemble RGB(Tfrom) -> RGB(Tto) color conversion matrix C.
    C = np.dot(np.linalg.inv(M1), M0)
    
    # Apply color conversion matrix C to the image.
    width, height = img.size
    for i in range(width):
        for j in range(height):
            rgb = np.array(img.get_pixel(i, j))
            rgb = np.dot(C, rgb)
            rgb = np.clip(rgb, 0.0, float("inf"))
            img.set_pixel(i, j, tuple(rgb))
    
    # Update image window
    xshade.scene().rendering.update_image()
    return



def blackbody_radiation_color_uv(T):
    """Blackbody radiation color in CIE 1960 uv color space.
    
    @param[in] T Target color tempreture (unit: K).
    @return u, v Blackbody radiation color in CIE 1960 uv color space.
    
    Scope of application of this approximation is from 1000K to 15000K.
    
    Reference:
    Approximation, Planckian locus, Wikipedia. Ref.: 2016-08-12.
    URL: https://en.wikipedia.org/wiki/Planckian_locus#Approximation
    """
    u = (0.860117757 + 1.54118254e-4 * T + 1.28641212e-7 * T * T) \
        / (1.0 + 8.42420235e-4 * T + 7.08145163e-7 * T * T)
    v = (0.317398726 + 4.22806245e-5 * T + 4.20481691e-8 * T * T) \
        / (1.0 - 2.89741816e-5 * T + 1.61456053e-7 * T * T)
    return u, v



if __name__ == "__main__": color_temperature_correction(3200.0)