"""
Conversion code from wavelength to RGB. 

This code is adapted from https://noah.org/wiki/Wavelength_to_RGB_in_Python (now unavailable), which
is itself based on http://www.physics.sfasu.edu/astro/color/spectra.html by Dan Bruton (also unavailable).

Some notes taken from the original files:

Color definitions:

    Color   Wavelength(nm) Frequency(THz)
    Red     620-750        484-400
    Orange  590-620        508-484
    Yellow  570-590        526-508
    Green   495-570        606-526
    Blue    450-495        668-606
    Violet  380-450        789-668

Peak wavelength responses for each type of photoreceptor cell in the human eye:
    S cone: 437 nm
    M cone: 533 nm
    L cone: 564 nm
    rod:    550 nm in bright daylight, 498 nm when dark adapted. 
                Rods adapt to low light conditions by becoming more sensitive.
"""
import numpy as np

def wavelength_to_rgb(wavelength, gamma=0.8, attenuation=0.):
    """
    Convert a given wavelength (in nanometers) of light to an 
    approximate RGB color value. The wavelength must be given
    in nanometers in the range from 380 nm through 750 nm.

    attenuation is the intensity at both ends of the spectra (default is 0)
    """
    # Prepare output
    scalar = np.isscalar(wavelength)
    w = np.asarray(wavelength)
    sh = w.shape
    N = w.size
    w.resize((N,))
    out = np.zeros(shape=(N,3), dtype=float)

    R, G, B = 0, 1, 2
    
    # Violet
    m = (w>=380)&(w<440)
    att = attenuation + (1-attenuation) * (w[m] - 380) / (440-380)    
    out[m, R] = att * (440-w[m]) / (440-380)
    out[m, B] = att

    # Blue
    m = (w>=440)&(w<490)
    out[m, G] = (w[m] - 440)/(490-440)
    out[m, B] = 1.

    # Green
    m = (w>=490)&(w<510)
    out[m, G] = 1.
    out[m, B] = (510-w[m])/(510-490)

    # Yellow
    m = (w>=510)&(w<580)
    out[m, R] = (w[m] - 510)/(580-510)
    out[m, G] = 1.

    # Orange
    m = (w>=580)&(w<645)
    out[m, R] = 1.
    out[m, G] = (645 -w[m])/(645-580)

    # Red
    m = (w>=645)&(w<750)
    out[m, R] = attenuation + (1-attenuation) * (750 - w[m]) / (750-645)

    # gamma
    out **= gamma

    if scalar:
        return out.item()
    else:
        out.resize(sh + (3,))
        return out