# 
# CCDD: Constrained Conditional Denoising Diffusion
# Copyright (C) 2025  James E. Fowler
# 
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Library General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
# 
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Library General Public License for more details.
# 
# You should have received a copy of the GNU Library General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., 675 Mass Ave, Cambridge,
# MA 02139, USA.
# 

import numpy as np
import numpy.linalg as npla
import numpy.matlib as ml
import scipy
import matplotlib.pyplot as plt
import time
from PIL import Image
import torch
import torch.nn.functional as F
import imageio
import tiler
# from torcheval.metrics import StructuralSimilarity as SSIM
# import torcheval.metrics.image.ssim as tm
import skimage.metrics as skim
from Utils import aviris, aviris220, rosis, vnir, hydice


###########################################################################
#
# HSI images in ndarray form have format (channels-last):
#    (batch_size, num_rows, num_cols, num_bands)
#
# HSI images in torch.tensor form have format (channels-first):
#    (batch_size, num_bands, num_rows, num_cols)
#
# hsi.ToTensor() and hsi.ToNumpy() automatically convert between
# channels-first and channels-last formats
#
###########################################################################


def ToTensor(X, dtype = None, device = None):
    if torch.is_tensor(X):
        return X
    X = torch.tensor(X, dtype = dtype, device = device)
    if X.ndim == 4:
        X = torch.permute(X, (0, 3, 1, 2))
    else:
        X = torch.permute(X, (2, 0, 1))
    return X


def ToNumpy(X):
    if not torch.is_tensor(X):
        return X
    if X.ndim == 4:
        X = torch.permute(X, (0, 2, 3, 1))
    else:
        X = torch.permute(X, (1, 2, 0))
    return np.array(X.cpu())


def Normalize(X, type = None):
    if torch.is_tensor(X):
        Xmax = torch.max(X)
        Xmin = torch.min(X)
    else:
        Xmax = np.max(X)
        Xmin = np.min(X)
    X = (X - Xmin) / (Xmax - Xmin)
    if type == "minmax":
        return X, Xmax, Xmin
    else:
        return (X - 0.5) * 2, Xmax, Xmin


def Denormalize(X, Xmax, Xmin, type = None):
    if type != "minmax":
        X = X / 2 + 0.5
    X = X * (Xmax - Xmin) + Xmin
    return X

    
def GaussianKernel(shape = (8, 8), sigma = 2.0):
    m, n = [(ss - 1.)/2. for ss in shape]
    y,x = np.ogrid[-m:m+1,-n:n+1]
    h = np.exp( -(x*x + y*y) / (2. * sigma * sigma) )
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    return h


# Use 'C' or 'torch' for row-major order, 'F' or 'matlab' for column-major order
def GetSpatialProjectionAsMatrix(spatial_ratio, num_rows, num_cols,
                                 psf = None,
                                 sparse = False,
                                 tensor = False,
                                 order = 'C'):
    if psf is None:
        psf = GaussianKernel(shape = (spatial_ratio, spatial_ratio))
    m, n = psf.shape
    M = num_rows // spatial_ratio
    N = num_cols // spatial_ratio
    num = m * n
    NUM = M * N

    if order == 'F' or order == 'matlab':
        I = np.kron(np.arange(1, NUM + 1).reshape(NUM, 1),
                    np.ones((num, 1))).astype(int).flatten() - 1
        Temp1 = np.kron(num_rows * np.arange(0, n),
                        np.ones((1, m))); 
        Temp2 = ml.repmat(np.arange(1, m + 1), 1, n)
        Temp3 = Temp1 + Temp2
        
        X, Y = np.meshgrid(np.arange(0, M), np.arange(0, N))
        Temp4 = (X * m + Y * num_rows * n).reshape(M * N, 1)
        Temp5 = ml.repmat(Temp3, 1, NUM)
        Temp6 = np.kron(Temp4.T, np.ones((1, num)))
        J = (Temp5 + Temp6).astype(int).flatten() - 1
    elif order == 'C' or order == 'torch':
        I = np.kron(np.arange(1, NUM + 1).reshape(NUM, 1),
                    np.ones((num, 1))).astype(int).flatten() - 1
        Temp1 = np.kron(num_cols * np.arange(0, m),
                        np.ones((1, n))); 
        Temp2 = ml.repmat(np.arange(1, n + 1), 1, m)
        Temp3 = Temp1 + Temp2
        
        Y, X = np.meshgrid(np.arange(0, N), np.arange(0, M))
        Temp4 = (Y * n + X * num_cols * m).reshape(NUM, 1)
        Temp5 = ml.repmat(Temp3, 1, NUM)
        Temp6 = np.kron(Temp4.T, np.ones((1, num)))
        J = (Temp5 + Temp6).astype(int).flatten() - 1
    else:
        print("Invalid order (%s)" % order)
        return None
        
    P = ml.repmat(psf.reshape(num, 1), NUM, 1).flatten()

    P = scipy.sparse.coo_matrix((P, (I, J)), dtype=np.float32).T
    
    if not sparse:
        if tensor:
            P = torch.tensor(P)
        else:
            P = P.toarray()
    else:
        if tensor:
            values = torch.FloatTensor(P.data)
            indices = torch.LongTensor(np.vstack((P.row, P.col)))
            P = torch.sparse.FloatTensor(indices, values,
                                         torch.Size(P.shape))
    return P


# def GetSpatialProjectionAsMatrix(spatial_ratio, num_rows, num_cols,
#                                  psf = None,
#                                  sparse = False,
#                                  tensor = False):
#     if psf is None:
#         psf = GaussianKernel(shape = (spatial_ratio, spatial_ratio))
#     m, n = psf.shape
#     M = num_rows // spatial_ratio
#     N = num_cols // spatial_ratio
#     num = m * n
#     NUM = M * N
#     I = np.kron(np.arange(1, NUM + 1).reshape(NUM, 1),
#                 np.ones((num, 1))).astype(int).flatten() - 1
#     Temp1 = np.kron(num_cols * np.arange(0, m),
#                     np.ones((1, n))); 
#     Temp2 = ml.repmat(np.arange(1, n + 1), 1, m)
#     Temp3 = Temp1 + Temp2
# 
#     Y, X = np.meshgrid(np.arange(0, N), np.arange(0, M))
#     Temp4 = (Y * n + X * num_cols * m).reshape(NUM, 1)
#     Temp5 = ml.repmat(Temp3, 1, NUM)
#     Temp6 = np.kron(Temp4.T, np.ones((1, num)))
#     J = (Temp5 + Temp6).astype(int).flatten() - 1
#     
#     P = ml.repmat(psf.reshape(num, 1), NUM, 1).flatten()
# 
#     P = scipy.sparse.coo_matrix((P, (I, J)), dtype=np.float32).T
#     
#     if not sparse:
#         if tensor:
#             P = torch.tensor(P)
#         else:
#             P = P.toarray()
#     else:
#         if tensor:
#             values = torch.FloatTensor(P.data)
#             indices = torch.LongTensor(np.vstack((P.row, P.col)))
#             P = torch.sparse.FloatTensor(indices, values,
#                                          torch.Size(P.shape))
#     return P


def SpectralProjectionUnfolded(X3, S):
    return S @ X3


def SpectralProjection(X, S):
    if torch.is_tensor(X):
        num_rows = X.shape[-2]
        num_cols = X.shape[-1]
    else:
        num_rows = X.shape[-3]
        num_cols = X.shape[-2]
    X3 = Unfold3(X)
    return Fold3(SpectralProjectionUnfolded(X3, S), num_rows, num_cols)

    
def SpatialProjectionUnfolded(X3, P):
    return X3 @ P


def SpatialProjection(X, P, strides = None, padding = 0):
    return_numpy = False
    if not torch.is_tensor(X):
        X = ToTensor(X)
        return_numpy = True
    if not torch.is_tensor(P):
        P = torch.tensor(P)
        
    return_squeeze = False
    if X.dim() == 3:
        X = torch.unsqueeze(X, 0)
        return_squeeze = True

    if strides == None:
        strides = (P.shape[0], P.shape[1])
    num_bands = X.shape[1]
    P = P.expand([num_bands, 1, -1, -1])

    X_tilde = F.conv2d(X, P, stride = strides, groups = num_bands,
                       padding = padding)
    if return_squeeze:
        X_tilde = torch.squeeze(X_tilde, 0)

    if return_numpy:
        return ToNumpy(X_tilde)
    else:
        return X_tilde
            
#     else:
#         return_squeeze = False
#         if (len(X.shape)) == 3:
#             X = X[np.newaxis, :, :, :]
#             return_squeeze = True
# 
#         num_batches = X.shape[0]
#         num_bands = X.shape[3]
#         P = np.tile(P[:, :, np.newaxis, np.newaxis], (1, 1, hsi_num_bands, 1))
#         num_out_channels = P.shape[3]
# 
#         X_tilde = np.zeros((num_batches, X.shape[1], X.shape[2],
#                             num_bands * num_out_channels), dtype = np.float32)
#         for batch in range(num_batches):
#             for band in range(num_bands):
#                 for channel in range(num_out_channels):
#                     out_band = band * num_out_channels + channel
#                     X_tilde[batch, :, :, out_band] = \
#                         scipy.ndimage.correlate(X[batch, :, :, band].squeeze(),
#                                                 P[:, :, band, channel].squeeze(),
#                                                 mode = 'constant')
# 
#         X_tilde = X_tilde[:,
#                           (P.shape[0] // 2)::P.shape[0],
#                           (P.shape[1] // 2)::P.shape[1], :]
#                                             
#         if return_squeeze:
#             X_tilde = np.squeeze(X_tilde, 0)


    
def SpatialProjectionTranspose(X_tilde, P, strides = None, padding = 0):
    return_numpy = False
    if not torch.is_tensor(X_tilde):
        X_tilde = ToTensor(X_tilde)
        return_numpy = True
    if not torch.is_tensor(P):
        P = torch.tensor(P)

    return_squeeze = False
    if X_tilde.dim() == 3:
        X_tilde = torch.unsqueeze(X_tilde, 0)
        return_squeeze = True
        
    if strides == None:
        strides = (P.shape[0], P.shape[1])
    num_bands = X_tilde.shape[1]
    P = P.expand([num_bands, 1, -1, -1])
    
    X = F.conv_transpose2d(X_tilde, P, stride = strides, groups = num_bands,
                           padding = 0)
    if return_squeeze:
        X = torch.squeeze(X, 0)

    if return_numpy:
        return ToNumpy(X)
    else:
        return X

    
# def SpatialProjectionTranspose(X_tilde, P, num_rows, num_cols):
#     return_numpy = False
#     if not torch.is_tensor(X_tilde):
#         X_tilde = tf.convert_to_tensor(X_tilde)
#         return_numpy = True
#         
#     return_squeeze = False
#     if (len(X_tilde.shape)) == 3:
#         X_tilde = tf.expand_dims(X_tilde, 0)
#         return_squeeze = True
# 
#     strides = [1, P.shape[0], P.shape[1], 1]
#     shape = [tf.shape(X_tilde)[0], num_rows, num_cols, tf.shape(X_tilde)[3]]
#     shape = tf.stack(shape)
#     X = tf.nn.depthwise_conv2d_backprop_input(shape,
#                                               P, X_tilde,
#                                               strides = strides,
#                                               padding = 'SAME')
# 
#     if return_squeeze:
#         X = tf.squeeze(X, 0)
#         
#     if return_numpy:
#         return X.numpy()
#     else:
#         return X


def Unfold1(X):
    if torch.is_tensor(X):
        if X.ndim == 4:
            batch_size, num_bands, num_rows, num_cols = X.shape
            X = torch.permute(X, (0, 2, 3, 1))
#             X = torch.permute(X, (0, 2, 1, 3))
            return X.reshape(-1, num_rows, num_cols * num_bands)
        else:
            num_bands, num_rows, num_cols = X.shape
            X = torch.permute(X, (1, 2, 0))
#             X = torch.permute(X, (1, 0, 2))
            return X.reshape(num_rows, num_cols * num_bands)
    else:
        if X.ndim == 4:
            batch_size, num_rows, num_cols, num_bands = X.shape
            return X.reshape(-1, num_rows, num_cols * num_bands)
        else:
            num_rows, num_cols, num_bands = X.shape
            return X.reshape(num_rows, num_cols * num_bands)
        
    
def Fold1(X1, num_cols, num_bands):
    if X1.ndim == 3:
        num_rows = X1.shape[1]
    else:
        num_rows = X1.shape[0]
        
    if torch.is_tensor(X1):
        if X1.ndim == 3:
            X = X1.reshape(-1, num_rows, num_cols, num_bands)
            X = torch.permute(X, (0, 3, 1, 2))
            return X
#             X = X1.reshape(-1, num_rows, num_bands, num_cols)
#             return torch.permute(X, (0, 2, 1, 3))
        else:
            X = X1.reshape(num_rows, num_cols, num_bands)
            X = torch.permute(X, (2, 0, 1))
            return X
#             X = X1.reshape(num_rows, num_bands, num_cols)
#             return torch.permute(X, (1, 0, 2))
    else:
        if X1.ndim == 3:
            X = X1.reshape(-1, num_rows, num_cols, num_bands)
            return X
        else:
            X = X1.reshape(num_rows, num_cols, num_bands)
            return X

        
def Unfold2(X):
    if torch.is_tensor(X):
        if X.ndim == 4:
            batch_size, num_bands, num_rows, num_cols = X.shape
            X = torch.permute(X, (0, 3, 2, 1))
            return X.reshape(-1, num_cols, num_rows * num_bands)
        else:
            num_bands, num_rows, num_cols = X.shape
            X = torch.permute(X, (2, 1, 0))
            return X.reshape(num_cols, num_rows * num_bands)
    else:
        if X.ndim == 4:
            batch_size, num_rows, num_cols, num_bands = X.shape
            X = np.transpose(X, (0, 2, 1, 3))
            return X.reshape(-1, num_cols, num_rows * num_bands)
        else:
            num_rows, num_cols, num_bands = X.shape
            X = np.transpose(X, (1, 0, 2))
            return X.reshape(num_cols, num_rows * num_bands)
        
            
def Fold2(X1, num_rows, num_bands):
    if X1.ndim == 3:
        num_cols = X1.shape[1]
    else:
        num_cols = X1.shape[0]

    if torch.is_tensor(X1):
        if X1.ndim == 3:
            X = X1.reshape(-1, num_cols, num_rows, num_bands)
            X = torch.permute(X, (0, 3, 2, 1))
            return X
        else:
            X = X1.reshape(num_cols, num_rows, num_bands)
            X = torch.permute(X, (2, 1, 0))
            return X
    else:
        if X1.ndim == 3:
            X = X1.reshape(-1, num_cols, num_rows, num_bands)
            X = np.transpose(X, (0, 2, 1, 3))
            return X
        else:
            X = X1.reshape(num_cols, num_rows, num_bands)
            X = torch.permute(X, (1, 0, 2))
            return X
        

def Unfold3(X):
    if torch.is_tensor(X):
        if X.ndim == 4:
            batch_size, num_bands, num_rows, num_cols = X.shape
            X3 = X.reshape(batch_size, num_bands, num_rows * num_cols)
#             X3 = torch.clone(X3, memory_format = torch.contiguous_format)
            return X3
        else:
            num_bands, num_rows, num_cols = X.shape
            X3 = X.reshape(num_bands, num_rows * num_cols)
#             X3 = torch.clone(X3, memory_format = torch.contiguous_format)
            return X3
    else:
        if X.ndim == 4:
            batch_size, num_rows, num_cols, num_bands = X.shape
            X = np.transpose(X, (0, 3, 1, 2))
            X3 = X.reshape(batch_size, num_bands, num_rows * num_cols)
#             X3 = X3.copy()
            return X3
        else:
            num_rows, num_cols, num_bands = X.shape
            X = np.transpose(X, (2, 0, 1))
            X3 = X.reshape(num_bands, num_rows * num_cols)
#             X3 = X3.copy()
            return X3

        
def Fold3(X3, num_rows, num_cols):
    if X3.ndim == 3:
        num_bands = X3.shape[1]
    else:
        num_bands = X3.shape[0]
        
    if torch.is_tensor(X3):
        if X3.ndim == 3:
            X = X3.reshape(-1, num_bands, num_rows, num_cols)
#             X = torch.clone(X, memory_format = torch.contiguous_format)
            return X
        else:
            X = X3.reshape(num_bands, num_rows, num_cols)
#             X = torch.clone(X, memory_format = torch.contiguous_format)
            return X
    else:
        if X3.ndim == 3:
            X = X3.reshape(-1, num_bands, num_rows, num_cols)
            X = np.transpose(X, (0, 2, 3, 1))
#             X = X.copy()
            return X
        else:
            X = X3.reshape(num_bands, num_rows, num_cols)
            X = np.transpose(X, (1, 2, 0))
#             X = X.copy()
            return X

        
# def BicubicInterpolation(L, spatial_ratio):
#     if torch.is_tensor(L):
#         device = L.device
#         L = ToNumpy(L)
#         tensor = True
#     else:
#         tensor = False
#     
#     if L.ndim == 3:
#         L = np.expand_dims(L, 0)
#         squeeze = True
#     else:
#         squeeze = False
# 
#     num_patches = L.shape[0]
#     num_rows = L.shape[1] * spatial_ratio
#     num_cols = L.shape[2] * spatial_ratio
#     num_bands = L.shape[3]
#     
#     X0 = np.zeros((num_patches, num_rows, num_cols, num_bands))
# 
#     for patch in range(num_patches):
#         for band in range(num_bands):
#             X0[patch, :, :, band] = \
#                 np.array(Image.fromarray(L[patch, :, :, band].squeeze()).resize(size=(num_cols, num_rows), resample=Image.BICUBIC))
# 
#     if tensor:
#         X0 = ToTensor(X0).to(device)
#     if squeeze:
#         return X0.squeeze()
#     else:
#         return X0

    
def BicubicInterpolation(L, spatial_ratio):
    return_numpy = False
    if not torch.is_tensor(L):
        L = ToTensor(L)
        return_numpy = True

    return_squeeze = False
    if (len(L.shape)) == 3:
        L = torch.unsqueeze(L, 0)
        return_squeeze = True
        
    X0 = torch.nn.functional.interpolate(L,
                                         scale_factor = spatial_ratio,
                                         mode = "bicubic")
                                             
    if return_squeeze:
        X0 = torch.squeeze(X0, 0)
    if return_numpy:
        return ToNumpy(X0)
    else:
        return X0

# 
#     else:
#         if L.ndim == 3:
#             L = np.expand_dims(L, 0)
#             return_squeeze = True
#         else:
#             return_squeeze = False
# 
#         num_patches = L.shape[0]
#         num_rows = L.shape[1] * spatial_ratio
#         num_cols = L.shape[2] * spatial_ratio
#         num_bands = L.shape[3]
#     
#         X0 = np.zeros((num_patches, num_rows, num_cols, num_bands))
# 
#         for patch in range(num_patches):
#             for band in range(num_bands):
#                 X0[patch, :, :, band] = \
#                     np.array(Image.fromarray(L[patch, :, :, band].squeeze()).resize(size=(num_cols, num_rows), resample=Image.BICUBIC))
#             
#         if return_squeeze:
#             return X0.squeeze()
#         else:
#             return X0


def GetPCATransform(X, num_pcs = 'all'):
    if torch.is_tensor(X):
        dtype = X.dtype
        device = X.device
        X = ToNumpy(X)
        tensor = True
    else:
        tensor = False

    L, U = np.linalg.eigh(Unfold3(X) @ Unfold3(X).T)
    if num_pcs != 'all':
        U = U[..., :num_pcs]

    if tensor:
        U = torch.tensor(U, dtype = dtype, device = device)
    return U


def Tile(X, padding, tile_shape, overlap):
    X = np.pad(X, (padding, padding, (0, 0)), mode='reflect')
    X_tiler = tiler.Tiler(data_shape = X.shape,
                          tile_shape = tile_shape,
                          overlap = overlap,
                          channel_dimension = 2)
    X_tiled = np.zeros((len(X_tiler),) + tile_shape, dtype = np.float32)
    X_tiler = Tiler(X, padding, tile_shape, overlap)
    X_tiled = np.zeros((len(X_tiler),) + tile_shape, dtype = np.float32)
    for tile_id, tile in X_tiler.iterate(X):
        X_tiled[tile_id, :, :, :] = tile

    return X_tiled, X_tiler


def MSE(X, X_hat):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)
    return skim.mean_squared_error(X, X_hat)


def RMSE(X, X_hat):
    return np.sqrt(MSE(X, X_hat))


def PSNR(X, X_hat, max_val = 255.):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)
    return skim.peak_signal_noise_ratio(X, X_hat, data_range = max_val)


def SNR(X, X_hat):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)
    mse = MSE(X, X_hat)
    if mse > 0:
        var = np.var(X)
        snr = 10 * np.log10(var / mse)
    else:
        snr = np.nan
    return snr

    
# def SAM(X, X_hat):
#     X = ToNumpy(X)
#     X_hat = ToNumpy(X_hat)
#     num_rows, num_cols, num_bands = X.shape
#     s = np.zeros((num_rows, num_cols))
#     for row in range(num_rows):
#         for col in range(num_cols):
#             x = X[row, col, :].squeeze()
#             x_hat = X_hat[row, col, :].squeeze()
#             x_norm = npla.norm(x)
#             x_hat_norm = npla.norm(x_hat)
#             if (x_norm < 1e-7) or (x_hat_norm < 1e-7):
#                 s[row, col] = np.nan
#             else:
#                 s[row, col] = np.real(np.arccos(np.dot(x, x_hat) /
#                                                 x_norm /
#                                                 x_hat_norm)) / \
#                                                 np.pi * 180
#     return np.mean(s)


def SAM(X, X_hat):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)
    X_norm = npla.norm (X, axis = -1)
    X_hat_norm = npla.norm (X_hat, axis = -1)
    I = (X_norm >= 1e-7) & (X_hat_norm >= 1e-7)
    s = np.zeros_like(X_norm)
    dp = np.einsum("...i,...i", X, X_hat)
    s[I] = dp[I] / X_norm[I] / X_hat_norm[I]
    s = np.clip(s, -1, 1)
    s[I] = np.real(np.arccos(s[I])) / np.pi * 180
    return np.mean(s[I])


def SSIM(X, X_hat, max_val = 255.):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)
    X = np.clip(X, 0.0, 1.0)
    X_hat = np.clip(X_hat, 0.0, 1.0)
    ssim = skim.structural_similarity(X, X_hat,
                                      data_range = max_val, channel_axis = -1)
    return ssim


# def TV(X):
#     if len(X.shape) == 3:
#         X = tf.expand_dims(X, axis = 0)
#     if len(X_hat.shape) == 3:
#         X_hat = tf.expand_dims(X_hat, axis = 0)
#     return tf.image.total_variation(X)
# 
# def nuclear_norm(X, U):
#     if (len(X.shape) == 3):
#         X = tf.expand_dims(X, axis = 0)
#     if (len(U.shape) == 2):
#         U = tf.expand_dims(U, axis = 0)
#     X3 = Unfold3(X)
#     Z = Unfold3(tfla.matmul(X3, U, adjoint_a = True))
#     S2 = tfla.matmul(Z, Z, adjoint_b = True)
#     
#     tf.debugging.check_numerics(S2, 'problematic S2')
#     return tfla.trace(tf.math.sqrt(S2))


def GetErrorPlot(X, X_hat, return_colormap = True, location = "bottom"):
    X = ToNumpy(X)
    X_hat = ToNumpy(X_hat)

    diff = np.sqrt(np.mean((X - X_hat) ** 2, axis = -1))
#     j = plt.imshow(diff, vmin = 0., vmax = 0.2, cmap='viridis')
    j = plt.imshow(diff, vmin = 0., vmax = 0.2, cmap='nipy_spectral')
    plt.gca().set_axis_off()
    plot = plt.gcf()
    if not return_colormap:
        return plot
    _, ax = plt.subplots()
    plt.colorbar(j, ax = ax, location = location)
    ax.remove()
    colorbar = plt.gcf()
    return plot, colorbar


def SaveErrorPlot(X, X_hat, errfile, colorbarfile = None, location = "bottom"):
    if colorbarfile is None:
        plot = GetErrorPlot(X, X_hat, return_colormap = False)
        plot.savefig(errfile, bbox_inches = 'tight', pad_inches = 0)
    else:
        plot, colorbar = GetErrorPlot(X, X_hat, location = location)
        plot.savefig(errfile, bbox_inches = 'tight', pad_inches = 0)
        colorbar.savefig(colorbarfile, bbox_inches = 'tight', pad_inches = 0)


def ProcessSensor(sensor = None):
    remove_absorption_bands = None
    if sensor == 'aviris':
        S = aviris.GetSpectralProjection().astype(np.float32)
        remove_absorption_bands = aviris.RemoveAbsorptionBands
        get_pseudocolor_image = aviris.GetPseudocolorImage
    elif sensor == 'aviris220':
        S = aviris220.GetSpectralProjection().astype(np.float32)
        remove_absorption_bands = aviris220.RemoveAbsorptionBands
        get_pseudocolor_image = aviris220.GetPseudocolorImage
    elif sensor == 'rosis':
        S = rosis.GetSpectralProjection().astype(np.float32)
        remove_absorption_bands = rosis.RemoveAbsorptionBands
        get_pseudocolor_image = rosis.GetPseudocolorImage
    elif sensor == 'vnir':
        S = vnir.GetSpectralProjection().astype(np.float32)
        remove_absorption_bands = vnir.RemoveAbsorptionBands
        get_pseudocolor_image = vnir.GetPseudocolorImage
    elif sensor == 'hydice':
        S = hydice.GetSpectralProjection().astype(np.float32)
        remove_absorption_bands = hydice.RemoveAbsorptionBands
        get_pseudocolor_image = hydice.GetPseudocolorImage
    else:
        print('Invalid sensor type (%s)' % sensor, flush = True)
        return None
    return S, remove_absorption_bands, get_pseudocolor_image
    
