# 
# 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.
# 

# For data from the Headwall Hyperspec-VNIR-C imaging sensor as described by:
#   N. Yokoya and A. Iwasaki, "Airborne hyperspectral data over Chikusei,"
#   Space Appl. Lab., Univ. Tokyo, Japan, Tech. Rep. SAL-2016-05-27, May 2016.

import numpy as np
from PIL import Image
from PIL import ImageEnhance

# VNIR 224 bands total
num_bands = 128
# VNIR radiometric resolution = 14 bits
bit_depth = 14


# VNIR wavelengths: 363 to 1018 nm
wavelength = np.linspace(363, 1018, num_bands)


# VNIR water-absorption bands: none, but get rid of last two bands for
# divisibilty by 3 needed by CCDD
def RemoveAbsorptionBands(X):
    if X.shape[-1] == num_bands:
        return X[..., :-2]
    else:
        return X


# Four multispectral bands of Landsat TM (skips Band 5, 6, 7)
# Bands                                       Wavelength (micrometers)
# ------------------------------------------------------------------
# Band 1 - Blue                                0.45 -  0.52
# Band 2 - Green                               0.52 -  0.60
# Band 3 - Red                                 0.63 -  0.69
# Band 4 - NIR                                 0.76 -  0.90
# Band 5 - Shortwave Infrared (SWIR) 1         1.55 -  1.75 
# Band 6 - Thermal                            10.40 - 12.50
# Band 7 - SWIR 2                              2.08 -  2.35
def GetSpectralProjection():
    band1 = ((450 <= wavelength) & (wavelength <= 520)).reshape(1, num_bands)
    band2 = ((520 <= wavelength) & (wavelength <= 600)).reshape(1, num_bands)
    band3 = ((630 <= wavelength) & (wavelength <= 690)).reshape(1, num_bands)
    band4 = ((760 <= wavelength) & (wavelength <= 900)).reshape(1, num_bands)
    
    S = np.concatenate((band1, band2, band3, band4), axis=0).astype(np.float32)
    return (S / np.sum(S, axis=1).reshape(4, 1))[:, :-2].astype(np.float32)


def PseudocolorBands():
    return (52, 33, 20)


def GetPseudocolorImage(X):
    if X.shape[-1] == num_bands:
        X = RemoveAbsorptionBands(X)
    X = X[..., PseudocolorBands()]
    X = X / (2. ** bit_depth) * 255.
    X = np.clip(X, 0, 255)
    X = X.astype(np.uint8)
    X = Image.fromarray(X)
    return ImageEnhance.Brightness(X).enhance(15)
