#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Modified by D. Moser in 2015-09-22

"""Convolve the CIE 1931 RGB to the Black Body function."""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mplcolors
import pyhdust as hdt
import pyhdust.images as phim
import pyhdust.phc as phc

lbd = np.linspace(370, 730, 101)  # nm
Temps = np.arange(1000., 12000 + 1, 1000)  # Kelvin

color = ['red', 'green', 'blue']
# CIE 1931 RGB table == Natural
fdat = np.loadtxt(hdt.hdtpath() + 'pyhdust/refs/rgb_eff.txt',
    unpack=True)
# Modified
fdatmod = fdat.copy()
fdatmod[2] *= 1.5
# Colorblind
fdatcb = fdat.copy()
fdatcb[2, 10:] = fdat[2, 0:-10]
fdatcb[2, :10] = 0

fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(12, 6))
for i in range(3):
    ax0.plot(fdat[0], fdat[i + 1] - 0.0, color=color[i]) 
    ax0.plot(fdatmod[0], fdatmod[i + 1] - 0.5, color=color[i], ls='--') 
    ax0.plot(fdatcb[0], fdatcb[i + 1] - 1.0, color=color[i], ls='-.') 

for i in range(len(Temps)):
    BB = phc.BBlbd(Temps[i], lbd * 1e-7)
    cor = phim.doColorConv(lbd, BB, fdat)
    ax1.plot([i], [0.0], 'o', ms=10 + 1 * i,
        color=mplcolors.rgb2hex(cor / 255.))
    cor = phim.doColorConv(lbd, BB, fdatmod)
    ax1.plot([i], [-0.5], 'o', ms=10 + 1 * i,
        color=mplcolors.rgb2hex(cor / 255.))
    cor = phim.doColorConv(lbd, BB, fdatcb)
    ax1.plot([i], [-1.0], 'o', ms=10 + 1 * i,
        color=mplcolors.rgb2hex(cor / 255.))

ax1.set_xlim([-1, 12])
ax1.set_ylim([-1.2, .3])
plt.savefig('BB_rgb.png')