# 
# 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
from PIL import Image
from PIL import ImageEnhance

# AVIRIS 224 bands total
num_bands = 224
# AVIRIS radiometric resolution = 14 bits
bit_depth = 14


# AVIRIS wavelengths: 400 to 2500 nm
wavelength = np.linspace(400, 2500, num_bands)


# AVIRIS water-absorption bands 108-112, 154-167, and 224
water_absorption_bands = np.array([band for band in range(108, 113)] +
                                  [band for band in range(154, 168)] +
                                  [224]) - 1


def AbsorptionBoolean():
    I = np.full((num_bands), True)
    I[water_absorption_bands] = False
    return I


def RemoveAbsorptionBands(X):
    if X.shape[-1] == num_bands:
        return X[..., AbsorptionBoolean()]
    else:
        return X


# Six multispectral bands of Landsat TM (skips Band 6)
# 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)
    band5 = ((1550 <= wavelength) & (wavelength <= 1750)).reshape(1, num_bands)
    band7 = ((2080 <= wavelength) & (wavelength <= 2350)).reshape(1, num_bands)
    
    S = np.concatenate((band1, band2, band3, band4, band5, band7), axis=0).astype(np.float32)
    return (S / np.sum(S, axis=1).reshape(6, 1))[:, AbsorptionBoolean()].astype(np.float32)


def PseudocolorBands():
    return (28, 16, 8)


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(2.75)
