from collections.abc import Iterable
from abc import abstractmethod
from matplotlib.pyplot import hist
from pandas import lreshape
from tqdm.auto import tqdm

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
import scipy
import wandb
from fastcore.parallel import parallel
from .color_matrices import *

__all__ = ['ColorMatching',
           'ColorMatchingInstance']

def interpolate_Ti_tf(lambdas):
    T_1 = T_lambda[:,1]
    T_1 = tfp.math.interp_regular_1d_grid(lambdas, T_lambda[0,0], T_lambda[-1,0], 
                                            T_1, fill_value='extrapolate')
    T_2 = T_lambda[:,2]
    T_2 = tfp.math.interp_regular_1d_grid(lambdas, T_lambda[0,0], T_lambda[-1,0], 
                                            T_2, fill_value='extrapolate')
    T_3 = T_lambda[:,3]
    T_3 = tfp.math.interp_regular_1d_grid(lambdas, T_lambda[0,0], T_lambda[-1,0], 
                                            T_3, fill_value='extrapolate')
    return T_1, T_2, T_3

def monochomatic_stimulus_2_tf(central_lambda, lambdas, width, max_radiance, background):
    """
    Generates a quasi-monochromatic spectrum. Gaussian centered in central lambda with desired width and peak on max_radiance over a background.
    """
    spectrum = max_radiance*tf.cast(tf.exp(-(lambdas-central_lambda)**2/width**2), tf.float32) + background_stimulus_tf(lambdas, background)
    return spectrum

def background_stimulus_tf(lambdas, background):
    """
    Generates a quasi-monochromatic spectrum. Gaussian centered in central lambda with desired width and peak on max_radiance over a background.
    """
    spectrum = tf.cast(tf.ones_like(lambdas, dtype=tf.float32)*background, tf.float32)
    return spectrum

def xyz2nrgb_tf(xyz, gamma, clip=False):
    """
    Changes from xyz to nrgb values (ready to be displayed in the screen). Expects and xyz color column vector or images.
    This function is TensorFlow's gradient friendly.

    Parameters
    ----------
    xyz: np.array
        Three-element array corresponding to a color in XYZ space.
    gamma: int or List
        Value(s) to be used in the transformation to digital values.
        The final result will be rgb = ng**(1/gamma).
    clip: bool
        Boolean determining if clipping will be used or not when the
        resultant values exceed the range [0,1].
    """
    rgb = Mxyz2ng @ xyz

    if clip and (tf.reduce_min(rgb) < 0 or tf.reduce_max(rgb) > 1):
        rgb = tf.clip_by_value(rgb, 0, 1)

    if type(gamma) == list:
        rgb = tf.pow(rgb, 1/gamma[:,None])
    elif type(gamma) == float:
        rgb = rgb**(1/gamma)
    else:
        raise TypeError('Gamma should be either a single value (float) or a list with 3.')
    
    return rgb

def get_recover_idxs(central_wavelengths, central_wavelengths_sorted):
    return np.array([np.where(original == central_wavelengths_sorted)[0][0] for original in central_wavelengths])

def recover_histories(histories, idx_recover):
    ## Must copy each list or it freaks out
    histories_recovered = [history.copy() for history in histories]
    for history in histories_recovered:
        weights_recovered = []
        for weights in history["Weights"]:
            weights_r = weights[idx_recover]
            weights_recovered.append(weights_r)
        history["Weights"] = weights_recovered
    return histories_recovered    

class ColorMatchingInstance():
    """
    A color matching experiment instance, where instance means
    only a single wavelength. A full color matching experiment
    is compromised of a full run over a collection of wavelengths.
    """
    def __init__(self,
                 wavelength,
                 central_wavelengths,
                 lambdas,
                 max_radiance,
                 background_radiance,
                 img_size,
                 space_transform_fn,
                 width_lambda=10,
                 width_central_lambdas=5,
                 initial_weights=[0.001, 0.001, 0.001, 0.001],
                 norm_grads=False,
                 norm_imgs=False):
        self.wavelength = wavelength
        self.central_wavelengths = central_wavelengths
        self.lambdas = lambdas
        self.max_radiance = max_radiance
        self.background_radiance = background_radiance
        self.width_lambda = width_lambda
        self.width_central_lambdas = width_central_lambdas
        self.img_size = img_size
        self.space_transform_fn = space_transform_fn
        self.norm_grads = norm_grads
        self.norm_imgs = (lambda x: x) if not norm_imgs else norm_imgs
        self.weights = tf.Variable(initial_weights,
                                   trainable=True,
                                   dtype=tf.float32)
        self._Ts = None
        self.loss = None
        self.optimizer = None
        self._history = None

    @property
    def Ts(self):
        if self._Ts is None:
            self._Ts = interpolate_Ti_tf(self.lambdas)
        return self._Ts

    def compile(self, loss, optimizer):
        self.loss = loss
        self.optimizer = optimizer

    def generate_images(self):
        spectrum_lambda = monochomatic_stimulus_2_tf(self.wavelength, self.lambdas, width=self.width_lambda, max_radiance=0.75*2*self.max_radiance, background=5*self.background_radiance)
        Y_l = km*tf.reduce_sum(self.Ts[1]*spectrum_lambda)
        spectrum_white = tf.cast(tf.ones_like(self.lambdas, dtype=tf.float32)*self.background_radiance*10, tf.float32)
        Y_w = km*tf.reduce_sum(self.Ts[1]*spectrum_white)
        spectrum_white = spectrum_white*Y_l/Y_w

        for i in range(len(self.central_wavelengths)):
            spectrum = monochomatic_stimulus_2_tf(self.central_wavelengths[i], self.lambdas, width=self.width_central_lambdas, max_radiance=0.75*self.max_radiance*tf.abs(self.weights[i]), background=self.background_radiance)
            if self.weights[i] >= 0:
                spectrum_lambda = spectrum_lambda + spectrum
            else:
                spectrum_white = spectrum_white + spectrum

        t1_l = km*tf.reduce_sum(self.Ts[0]*spectrum_lambda)
        t2_l = km*tf.reduce_sum(self.Ts[1]*spectrum_lambda)
        t3_l = km*tf.reduce_sum(self.Ts[2]*spectrum_lambda)
        t1_w = km*tf.reduce_sum(self.Ts[0]*spectrum_white)
        t2_w = km*tf.reduce_sum(self.Ts[1]*spectrum_white)
        t3_w = km*tf.reduce_sum(self.Ts[2]*spectrum_white)

        t_l = tf.convert_to_tensor([t1_l, t2_l, t3_l], dtype = tf.float32)[:,None]
        t_w = tf.convert_to_tensor([t1_w, t2_w, t3_w], dtype = tf.float32)[:,None]

        rgb_lambda = self.space_transform_fn(t_l)
        rgb_white = self.space_transform_fn(t_w)

        img_lambda = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_lambda, perm=[1,0]), tf.float32)
        img_white = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_white, perm=[1,0]), tf.float32)

        return self.norm_imgs(img_lambda), self.norm_imgs(img_white)#, rgb_2_l, rgb_2_w

    @property
    def history(self):
        if self._history is None:
            print('The experiment has not been fitted yet.')
        return self._history

    def fit(self, model, epochs, verbose=True, use_tqdm=True, lr_callback=None, early_stopping=None):
        if self._history is None:
            self._history = {'Loss':[], 'GradsL2':[], 'Weights':[]}
        pbar = tqdm(range(epochs)) if use_tqdm else range(epochs)
        for epoch in pbar:
            with tf.GradientTape() as tape:
                img_l, img_w = self.generate_images()
                imgs = tf.concat([img_l, img_w], axis=0)
                response_l, response_w = model.predict(imgs)
                response_l, response_w = response_l[None,:,:,:], response_w[None,:,:,:]
                loss = self.loss(response_l, response_w)
            
            grads = tape.gradient(loss, self.weights)
            self.optimizer.apply_gradients(zip([grads], [self.weights]))
            self.history['Loss'].append(loss.numpy().item())
            self.history['GradsL2'].append(tf.reduce_sum(grads**2).numpy().item())
            self.history['Weights'].append(self.weights.numpy())

            if lr_callback is not None and epoch >= 1:
                lr_callback.check_improvement(self.history)
                lr_callback.reduce_lr()

            if early_stopping is not None and epoch >= 1:
                early_stopping.check_improvement(self.history)
                if early_stopping.stop_training():
                    break

            if verbose and not use_tqdm:
                if epoch % verbose == 0:
                    print(f'Epoch {epoch+1} -> Loss: {self.history["Loss"][-1]} | GradsL2: {self.history["GradsL2"][-1]}')
        return self.history

class ColorMatchingInstanceWandb():
    """
    A `ColorMatchingInstace` with `wandb` logging.
    """
    import wandb

    def __init__(self,
                 wavelength,
                 central_wavelengths,
                 lambdas,
                 max_radiance,
                 background_radiance,
                 img_size,
                 space_transform_fn,
                 to_rgb=None, # Function to project the images to RGB to log them properly.
                 width_lambda=10,
                 width_central_lambdas=5,
                 initial_weights=[0.001, 0.001, 0.001, 0.001],
                 norm_grads=False):
        self.wavelength = wavelength
        self.central_wavelengths = central_wavelengths
        self.lambdas = lambdas
        self.max_radiance = max_radiance
        self.background_radiance = background_radiance
        self.width_lambda = width_lambda
        self.width_central_lambdas = width_central_lambdas
        self.img_size = img_size
        self.space_transform_fn = space_transform_fn
        self.to_rgb = to_rgb
        self.norm_grads = norm_grads
        self.weights = tf.Variable(initial_weights,
                                   trainable=True,
                                   dtype=tf.float32)
        self._Ts = None
        self.loss = None
        self.optimizer = None
        self._history = None

    @property
    def Ts(self):
        if self._Ts is None:
            self._Ts = interpolate_Ti_tf(self.lambdas)
        return self._Ts

    def compile(self, loss, optimizer):
        self.loss = loss
        self.optimizer = optimizer

    def generate_images(self):
        spectrum_lambda = monochomatic_stimulus_2_tf(self.wavelength, self.lambdas, width=self.width_lambda, max_radiance=0.75*2*self.max_radiance, background=5*self.background_radiance)
        Y_l = km*tf.reduce_sum(self.Ts[1]*spectrum_lambda)
        spectrum_white = tf.cast(tf.ones_like(self.lambdas, dtype=tf.float32)*self.background_radiance*10, tf.float32)
        Y_w = km*tf.reduce_sum(self.Ts[1]*spectrum_white)
        spectrum_white = spectrum_white*Y_l/Y_w

        for i in range(len(self.central_wavelengths)):
            spectrum = monochomatic_stimulus_2_tf(self.central_wavelengths[i], self.lambdas, width=self.width_central_lambdas, max_radiance=0.75*self.max_radiance*tf.abs(self.weights[i]), background=self.background_radiance)
            if self.weights[i] >= 0:
                spectrum_lambda = spectrum_lambda + spectrum
            else:
                spectrum_white = spectrum_white + spectrum

        t1_l = km*tf.reduce_sum(self.Ts[0]*spectrum_lambda)
        t2_l = km*tf.reduce_sum(self.Ts[1]*spectrum_lambda)
        t3_l = km*tf.reduce_sum(self.Ts[2]*spectrum_lambda)
        t1_w = km*tf.reduce_sum(self.Ts[0]*spectrum_white)
        t2_w = km*tf.reduce_sum(self.Ts[1]*spectrum_white)
        t3_w = km*tf.reduce_sum(self.Ts[2]*spectrum_white)

        t_l = tf.convert_to_tensor([t1_l, t2_l, t3_l], dtype = tf.float32)[:,None]
        t_w = tf.convert_to_tensor([t1_w, t2_w, t3_w], dtype = tf.float32)[:,None]

        rgb_lambda = self.space_transform_fn(t_l)
        rgb_white = self.space_transform_fn(t_w)

        img_lambda = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_lambda, perm=[1,0]), tf.float32)
        img_white = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_white, perm=[1,0]), tf.float32)

        return img_lambda, img_white#, rgb_2_l, rgb_2_w

    @property
    def history(self):
        if self._history is None:
            print('The experiment has not been fitted yet.')
        return self._history

    def fit(self, model, epochs, verbose=True, use_tqdm=True, lr_callback=None, early_stopping=None):
        if self._history is None:
            self._history = {'Loss':[], 'GradsL2':[], 'Weights':[]}
        pbar = tqdm(range(epochs)) if use_tqdm else range(epochs)
        for epoch in pbar:
            with tf.GradientTape() as tape:
                img_l, img_w = self.generate_images()
                imgs = tf.concat([img_l, img_w], axis=0)
                response_l, response_w = model.predict(imgs)
                response_l, response_w = response_l[None,:,:,:], response_w[None,:,:,:]
                loss = self.loss(response_l, response_w)
            
            grads = tape.gradient(loss, self.weights)
            self.optimizer.apply_gradients(zip([grads], [self.weights]))
            self.history['Loss'].append(loss.numpy().item())
            self.history['GradsL2'].append(tf.reduce_sum(grads**2).numpy().item())
            self.history['Weights'].append(self.weights.numpy())

            if lr_callback is not None and epoch >= 1:
                lr_callback.check_improvement(self.history)
                lr_callback.reduce_lr()

            if early_stopping is not None and epoch >= 1:
                early_stopping.check_improvement(self.history)
                if early_stopping.stop_training():
                    break

            if verbose and not use_tqdm:
                if epoch % verbose == 0:
                    print(f'Epoch {epoch+1} -> Loss: {self.history["Loss"][-1]} | GradsL2: {self.history["GradsL2"][-1]}')
            
            ## WandB logging
            wandb.log({k:v[-1] for k, v in self.history.items()})
            for i, w in enumerate(self.history["Weights"][-1]):
                wandb.log({f"Weight_{i}":w})
            if epoch % verbose == 0:
                img_l, img_w = img_l.numpy().squeeze(), img_w.numpy().squeeze()
                if self.to_rgb is not None: img_l, img_w = self.to_rgb(img_l), self.to_rgb(img_w)
                fig, axes = plt.subplots(1,2)
                axes[0].imshow(img_l)
                axes[0].set_title("Stimuli")
                axes[1].imshow(img_w)
                axes[1].set_title("White")
                wandb.log({"Images":wandb.Image(fig)})
        return self.history

class ReduceLROnPlateau():
    def __init__(self,
                 factor,
                 patience,
                 cooldown,
                 min_lr,
                 delta,
                 optimizer,
                 verbose=False):
        self.factor = factor
        self.patience = patience
        self.cooldown = cooldown
        self.min_lr = min_lr
        self.delta = delta
        self.optimizer = optimizer
        self.verbose = verbose
        self.steps_no_improve = 0
        self.steps_cooldown = cooldown
        self.min_loss = 99999999

    def check_improvement(self, history):
        # if np.abs((history['Loss'][-1]-self.min_loss)/self.min_loss) <= self.delta:
        if history['Loss'][-1] >= self.min_loss:
            self.steps_no_improve += 1
        else:
            self.min_loss = history['Loss'][-1]
            self.steps_no_improve = 0
    
    def reduce_lr(self):
        lr = float(tf.keras.backend.get_value(self.optimizer.learning_rate))
        if lr > self.min_lr:
            if self.steps_no_improve == self.patience:
                if self.steps_cooldown < self.cooldown:
                    self.steps_cooldown += 1
                else:    
                    new_lr = lr*self.factor
                    tf.keras.backend.set_value(self.optimizer.lr, new_lr)
                    self.steps_cooldown = 0
                    if self.verbose:
                        print(f'Learning Rate changed: {lr} -> {new_lr}')
    def reset(self):
        self.steps_no_improve = 0
        self.steps_cooldown = self.cooldown
        self.min_loss = 99999999

class EarlyStopping():
    def __init__(self,
                 patience,
                 delta,
                 verbose=False):
        self.patience = patience
        self.delta = delta
        self.verbose = verbose
        self.steps_no_improve = 0
        self.min_loss = 99999999

    def check_improvement(self, history):
        # if np.abs((history['Loss'][-1]-self.min_loss)/self.min_loss) <= self.delta:
        if history['Loss'][-1] >= self.min_loss:
            self.steps_no_improve += 1
        else:
            self.min_loss = history['Loss'][-1]
            self.steps_no_improve = 0
    
    def stop_training(self):
        if self.steps_no_improve == self.patience:
            if self.verbose:
                print('Stopping training.')
            return True
        else:
            return False

    def reset(self):
        self.steps_no_improve = 0
        self.min_loss = 99999999

class EarlyStoppingRelative():
    def __init__(self,
                 patience,
                 delta,
                 verbose=False):
        self.patience = patience
        self.delta = delta
        self.verbose = verbose
        self.steps_no_improve = 0
        self.min_loss = 99999999

    def check_improvement(self, history):
        # if np.abs((history['Loss'][-1]-self.min_loss)/self.min_loss) <= self.delta:
        prev_loss = history['Loss'][-2*self.patience//3:-1]
        if not prev_loss: 
            prev_loss = self.min_loss
        else:
            prev_loss = np.min(prev_loss)

        if history['Loss'][-1] >= prev_loss:
            self.steps_no_improve += 1
        else:
            self.min_loss = history['Loss'][-1]
            self.steps_no_improve = 0
    
    def stop_training(self):
        if self.steps_no_improve == self.patience:
            if self.verbose:
                print('Stopping training.')
            return True
        else:
            return False

    def reset(self):
        self.steps_no_improve = 0
        self.min_loss = 99999999

class ColorMatchingExperiment(ColorMatchingInstance):
    """
    Color matching experiment where different monochromatic lights
    are given and the user has to minimize the distance with respect
    to a white image.
    This process is performed in an iterative way by optimizing the
    ammount of each color that is put into the images from a set
    of four available wavelenghts.
    """
    
    def __init__(self,
                 wavelengths=np.linspace(380, 770, 50),
                 initial_weights=[0.001, 0.001, 0.001, 0.001],
                 **kwargs):
        # super(ColorMatchingExperiment, self).__init__(wavelength=None,
        #                                               **kwargs)
        self.wavelengths = wavelengths
        if not any([isinstance(el, Iterable) for el in initial_weights]):
            self.initial_weights = [initial_weights]*len(self.wavelengths)
        else: self.initial_weights = initial_weights
        self.instances = [ColorMatchingInstance(wavelength=wavelength, initial_weights=initial_weights, **kwargs) for wavelength, initial_weights in zip(wavelengths, self.initial_weights)]

    def compile(self, loss, optimizer):
        for instance in self.instances:
            instance.loss = loss
            instance.optimizer = optimizer

    def fit(self, model, epochs, verbose=True, use_tqdm=True, lr_cb_fn=None, es_cb_fn=None):
        histories = []
        pbar = tqdm(self.instances) if use_tqdm else self.instances
        for instance in pbar:
            if verbose and not use_tqdm:
                print(f'Wavelength: {instance.wavelength}')
            if lr_cb_fn is not None:
                lr_callback = lr_cb_fn()
            else:
                lr_callback = None

            if es_cb_fn is not None:
                early_stopping = es_cb_fn()
            else:
                early_stopping = None

            history = instance.fit(model,
                                   epochs=epochs,
                                   verbose=verbose,
                                   use_tqdm=use_tqdm,
                                   lr_callback=lr_callback,
                                   early_stopping=early_stopping)
            histories.append(history)
        return histories

class ColorMatchingExperimentWandb(ColorMatchingInstance):
    """
    Color matching experiment where different monochromatic lights
    are given and the user has to minimize the distance with respect
    to a white image.
    This process is performed in an iterative way by optimizing the
    ammount of each color that is put into the images from a set
    of four available wavelenghts.
    """
    
    def __init__(self,
                 project, # project name in WandB.
                #  to_rgb=None,
                 wavelengths=np.linspace(380, 770, 50),
                 **kwargs):
        # super(ColorMatchingExperiment, self).__init__(wavelength=None,
        #                                               **kwargs)
        self.project = project
        # self.to_rgb = to_rgb
        self.wavelengths = wavelengths
        self.instances = [ColorMatchingInstanceWandb(wavelength=wavelength, **kwargs) for wavelength in wavelengths]

    def compile(self, loss, optimizer):
        for instance in self.instances:
            instance.loss = loss
            instance.optimizer = optimizer

    def fit(self, model, epochs, verbose=True, use_tqdm=True, lr_cb_fn=None, es_cb_fn=None):
        histories = []
        pbar = tqdm(self.instances) if use_tqdm else self.instances
        for instance in pbar:
            wandb.init(project=self.project,
                       name=f"Wavelength={instance.wavelength}")
            if verbose and not use_tqdm:
                print(f'Wavelength: {instance.wavelength}')
            if lr_cb_fn is not None:
                lr_callback = lr_cb_fn()
            else:
                lr_callback = None

            if es_cb_fn is not None:
                early_stopping = es_cb_fn()
            else:
                early_stopping = None

            history = instance.fit(model,
                                   epochs=epochs,
                                   verbose=verbose,
                                   use_tqdm=use_tqdm,
                                   lr_callback=lr_callback,
                                   early_stopping=early_stopping)
            histories.append(history)
            wandb.finish()
        return histories


class ColorMatchingExperimentWandbParallel(ColorMatchingInstance):
    """
    Color matching experiment where different monochromatic lights
    are given and the user has to minimize the distance with respect
    to a white image.
    This process is performed in an iterative way by optimizing the
    ammount of each color that is put into the images from a set
    of four available wavelenghts.
    """
    
    def __init__(self,
                 project, # project name in WandB.
                #  to_rgb=None,
                 wavelengths=np.linspace(380, 770, 50),
                 n_workers=2, # Number of cores to use for parallel computing. False means no parallel.
                 **kwargs):
        # super(ColorMatchingExperiment, self).__init__(wavelength=None,
        #                                               **kwargs)
        self.project = project
        # self.to_rgb = to_rgb
        self.wavelengths = wavelengths
        self.n_workers = n_workers
        self.instances = [ColorMatchingInstanceWandb(wavelength=wavelength, **kwargs) for wavelength in wavelengths]

    def compile(self, loss, optimizer):
        for instance in self.instances:
            instance.loss = loss
            instance.optimizer = optimizer

    def _fit_instance(self, instance, model, epochs, verbose=True, use_tqdm=True, lr_cb_fn=None, es_cb_fn=None):
        wandb.init(project=self.project,
                    name=f"Wavelength={instance.wavelength}")
        if verbose and not use_tqdm:
            print(f'Wavelength: {instance.wavelength}')
        if lr_cb_fn is not None:
            lr_callback = lr_cb_fn()
        else:
            lr_callback = None

        if es_cb_fn is not None:
            early_stopping = es_cb_fn()
        else:
            early_stopping = None

        history = instance.fit(model,
                                epochs=epochs,
                                verbose=verbose,
                                use_tqdm=use_tqdm,
                                lr_callback=lr_callback,
                                early_stopping=early_stopping)
        wandb.finish()
        return history
    
    def fit(self, model, epochs, verbose=True, use_tqdm=True, lr_cb_fn=None, es_cb_fn=None):
        histories = parallel(self._fit_instance, items=self.instances, model=model, epochs=epochs, verbose=verbose, use_tqdm=False, lr_cb_fn=lr_cb_fn, es_cb_fn=es_cb_fn, n_workers=self.n_workers, progress=use_tqdm)
        return histories


class LuminanceMatchingInstance(ColorMatchingInstance):
    """
    A luminance matching experiment instance, where instance means
    only a single wavelength. A full luminance matching experiment
    is compromised of a full run over a collection of wavelengths.
    """

    def __init__(self, 
                 initial_weights, 
                 wavelength_ref,
                 **kwargs):
        super(LuminanceMatchingInstance, self).__init__(**kwargs)
        self.weights = tf.Variable(initial_weights,
                                   trainable=True,
                                   dtype=tf.float32,
                                   constraint=tf.keras.constraints.NonNeg())
        self.wavelength_ref = wavelength_ref

    def generate_images(self):
        spectrum_lambda = monochomatic_stimulus_2_tf(self.wavelength, self.lambdas, width=10, max_radiance=tf.abs(self.weights)*self.max_radiance, background=self.background_radiance)
        spectrum_white = monochomatic_stimulus_2_tf(self.wavelength_ref, self.lambdas, width=10, max_radiance=self.max_radiance, background=self.background_radiance)

        t1_l = km*tf.reduce_sum(self.Ts[0]*spectrum_lambda)
        t2_l = km*tf.reduce_sum(self.Ts[1]*spectrum_lambda)
        t3_l = km*tf.reduce_sum(self.Ts[2]*spectrum_lambda)
        t1_w = km*tf.reduce_sum(self.Ts[0]*spectrum_white)
        t2_w = km*tf.reduce_sum(self.Ts[1]*spectrum_white)
        t3_w = km*tf.reduce_sum(self.Ts[2]*spectrum_white)

        t_l = tf.convert_to_tensor([t1_l, t2_l, t3_l], dtype = tf.float32)[:,None]
        t_w = tf.convert_to_tensor([t1_w, t2_w, t3_w], dtype = tf.float32)[:,None]

        rgb_lambda = self.space_transform_fn(t_l)
        rgb_white = self.space_transform_fn(t_w)

        img_lambda = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_lambda, perm=[1,0]), tf.float32)
        img_white = tf.ones((1,*self.img_size,3), dtype=tf.float32)*tf.cast(tf.transpose(rgb_white, perm=[1,0]), tf.float32)

        return img_lambda, img_white#, rgb_2_l, rgb_2_w

class LuminanceMatchingExperiment(ColorMatchingExperiment):
    """
    Luminance matching experiment where different monochromatic lights
    are given and the user has to minimize the distance with respect
    to a white image.
    This process is performed in an iterative way by optimizing the
    ammount 'brightness' (luminance) into the images.
    """
    def __init__(self,
                 wavelengths=np.linspace(380, 770, 50),
                 wavelength_ref=550.0,
                 **kwargs):
        self.wavelengths = wavelengths
        self.wavelength_ref=wavelength_ref
        self.instances = [LuminanceMatchingInstance(wavelength=wavelength, wavelength_ref=wavelength_ref, **kwargs) for wavelength in wavelengths]