# 
# 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 os
import glob
import numpy as np
import torch
import torchvision as tv
import hdf5storage as hdf
import imageio
import matplotlib.pyplot as plt
from torch.utils.flop_counter import FlopCounterMode
from contextlib import nullcontext
import argparse
import time
import diffusers
import sys
from Utils import hsi


##########################################################


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--sensor',
                        default = 'aviris',
                        help='(string) sensor type (default = "aviris")')
    parser.add_argument('-sr', '--spatial_ratio',
                        type = int,
                        default = 8,
                        help = '(int) spatial subsampling ratio (default = 8)')
    parser.add_argument('-bb', '--band_block',
                        type = int,
                        default = 3,
                        help = '(int) number of bands to block together (default = 3)')
    parser.add_argument('-T', '--num_train_steps',
                        type = int,
                        default = 1000,
                        help = '(int) number of diffusion time steps (default = 1000)')
    parser.add_argument('-bmx', '--beta_max',
                        type = float,
                        default = 0.01,
                        help = '(float) maximum beta value (default = 0.01)')
    parser.add_argument('-bmn', '--beta_min',
                        type = float,
                        default = 0.0001,
                        help = '(float) minimum beta value (default = 0.0001)')
    parser.add_argument('-ns', '--num_inference_steps',
                        type = int,
                        default = 10,
                        help = '(int) number of inference time steps (default = 10)')
    parser.add_argument('-li', '--num_landweber_iterations',
                        type = int,
                        default = 2,
                        help = '(int) number of Landweber iterations (default = 5)')
    parser.add_argument('--ddim', action = 'store_true',
                        default = False,
                        help = 'use DDIM scheduler (default = DDPM)')
    parser.add_argument('--unconstrained', action = 'store_true',
                        default = False,
                        help = 'no enforcement of S or P constraints (implies Landweber iterations = 1)')
    parser.add_argument('--unconditioned', action = 'store_true',
                        default = False,
                        help = 'no conditioning')
    parser.add_argument('--transpose_adjoint', action = 'store_true',
                        default = False,
                        help = 'use transpose as spatial adjoint operator (default = bicubic interpolation)')
    parser.add_argument('--tau',
                        type = float,
                        default = None,
                        help = '(float) spectral regularizer constant (default = 3/2/sigma_1^2, where sigma_1 is the largest singular value of the S matrix)')
    parser.add_argument('--eta',
                        type = float,
                        default = 1.0,
                        help = '(float) spatial regularizer constant (default = 1.0')
    parser.add_argument('--flops', action = 'store_true',
                        default = False,
                        help = 'report number of FLOPS')
    parser.add_argument('testing_dir',
                        help = '(string) directory of testing images (MAT format)')
    parser.add_argument('model_dir',
                        help='(string) directory of trained model')
    parser.add_argument('results_dir',
                        help='(string) directory of output results')
    args = parser.parse_args()
    return args


##########################################################


def diffusion_update(sample, t, Lprime, H, model, scheduler,
                     num_bands, band_block, unconditioned):
    for band in range(0, num_bands, band_block):
        band_range = range(band, band + band_block)
        sample_block = sample[:, band_range, ...]
        if unconditioned:
            input = sample_block
        else:
            input = torch.cat((sample_block,
                               Lprime[:, band_range, ...],
                               H), dim = 1)
        with torch.inference_mode():
            noise_pred = model(input, t)["sample"]
        sample[:, band_range, ...] = scheduler.step(noise_pred, t,
                                                    sample_block).prev_sample
    return sample


def landweber_spectral_update(X, H, S, tau):
    residual = H - hsi.SpectralProjection(X, S)
    residual = hsi.SpectralProjection(residual, torch.t(S))
    return X + tau * residual

        
def landweber_spatial_update(X, L, P, eta, transpose_adjoint = False):
    spatial_ratio = P.shape[0]
    residual = L - hsi.SpatialProjection(X, P)
    if transpose_adjoint:
        residual = hsi.SpatialProjectionTranspose(residual, P)
    else:
        residual = hsi.BicubicInterpolation(residual, spatial_ratio)
    return X + eta * residual


def landweber(num_iterations, tau, eta, sample, t, Lprime, H, L, S, P,
              model, scheduler, num_bands, band_block,
              unconstrained = False,
              unconditioned = False,
              transpose_adjoint = False):
    if unconstrained:
        sample = diffusion_update(sample, t, Lprime, H,
                                  model, scheduler, num_bands, band_block,
                                  unconditioned)
    else:
        for i in range(num_iterations):
            sample = diffusion_update(sample, t, Lprime, H,
                                      model, scheduler, num_bands, band_block,
                                      unconditioned)
            sample = landweber_spectral_update(sample, H, S, tau)
            sample = landweber_spatial_update(sample, L, P, eta,
                                              transpose_adjoint)
    return sample


##########################################################


def CCDDReconstructImage(Lprime, H, L, S, P,
                         model_unet,
                         noise_scheduler,
                         tau,
                         eta,
                         band_block = 3,
                         num_landweber_iterations = 2,
                         unconstrained = False,
                         unconditioned = False,
                         transpose_adjoint = False,
                         device = None):
    
    if device is None:
        device_name = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device_name)

    X_hat = torch.randn_like(Lprime)

    msi_num_bands, hsi_num_bands = S.shape

    for t in noise_scheduler.timesteps:
        X_hat = landweber(num_landweber_iterations, tau, eta,
                          X_hat, t, Lprime, H, L, S, P,
                          model_unet, noise_scheduler,
                          hsi_num_bands, band_block,
                          unconstrained,
                          unconditioned,
                          transpose_adjoint)

    return X_hat


##########################################################


def CCDDReconstructFile(filename,
                        sensor,
                        model_unet,
                        noise_scheduler,
                        results_dir = None,
                        output_error_image = False,
                        spatial_ratio = 8,
                        band_block = 3,
                        num_landweber_iterations = 2,
                        tau = None,
                        eta = 1.0,
                        unconstrained = False,
                        unconditioned = False,
                        transpose_adjoint = False,
                        flops = False,
                        device = None,
                        quiet = False):

    if device is None:
        device_name = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device_name)
        
    S, remove_absorption_bands, get_pseudocolor_image = \
        hsi.ProcessSensor(sensor)

    if tau is None:
        sigma_1 = np.linalg.svd(S, full_matrices = False, compute_uv = False)[0]
        tau = 3/ 2 / sigma_1 ** 2
        
    S = torch.tensor(S, dtype = torch.float32, device = device)
    
    P = hsi.GaussianKernel(shape = (spatial_ratio, spatial_ratio))
    P = torch.tensor(P, dtype = torch.float32, device = device)

    if not os.path.isfile(filename):
        raise RuntimeError("Cannot find dataset %s" % filename)
    if not quiet: print('Processing image %s' % filename, flush = True)
        
    X = hdf.loadmat(filename)['X'].astype(np.float32)
    if remove_absorption_bands is not None:
        X = remove_absorption_bands(X)
    X = hsi.ToTensor(X, device = device)
    if X.shape[-3] != S.shape[1]:
        raise RuntimeError('Dataset %s (%d bands) does not match "%s" sensor (%d bands)' % (filename, X.shape[-3], sensor, S.shape[1]))
    num_rows = (X.shape[1] // spatial_ratio) * spatial_ratio
    num_cols = (X.shape[2] // spatial_ratio) * spatial_ratio
    if (num_rows != X.shape[1]) or (num_cols != X.shape[2]):
        if not quiet:
            print("Cropping %s to spatial size %d x %d" % (filename,
                                                           num_rows, num_cols),
                  flush = True)
        X = X[:, :num_rows, :num_cols]

    X, Xmax, Xmin = hsi.Normalize(X)
    X = X.expand(1, -1, -1, -1)

    # MSI
    H = hsi.SpectralProjection(X, S)
    
    # HSI
    L = hsi.SpatialProjection(X, P)
    Lprime = hsi.BicubicInterpolation(L, spatial_ratio)

    if flops:
        flop_counter = FlopCounterMode(display = False)
        
    tic = time.time()
    with flop_counter if flops else nullcontext():
        X_hat = CCDDReconstructImage(
            Lprime, H, L, S, P,
            model_unet = model_unet,
            noise_scheduler = noise_scheduler,
            tau = tau,
            eta = eta,
            band_block = band_block,
            num_landweber_iterations = num_landweber_iterations,
            unconstrained = unconstrained,
            unconditioned = unconditioned,
            transpose_adjoint = transpose_adjoint,
            device = device)
    toc = time.time()
    if not quiet: print('  Reconstruction time = %.2f (sec)' % (toc - tic),
                        flush = True)
    if flops and not quiet:
        print('  TFLOPS = %.2f' % 
              (flop_counter.get_total_flops() / 1e12), flush = True)
    
    X = torch.squeeze(X, 0)
    X_hat = torch.squeeze(X_hat, 0)

    # Renormalize to [0, 1] for quality measures, for consistency with LRTA
    X = X / 2 + 0.5
    X_hat = X_hat / 2 + 0.5

    psnr = hsi.PSNR(X, X_hat, max_val = 1.0)
    rmse = hsi.RMSE(X, X_hat)
    sam = hsi.SAM(X, X_hat)
    if not quiet:
        print('    PSNR = %.2f dB' % psnr, flush = True)
        print('    RMSE = %e' % rmse, flush = True)
        print('     SAM = %.4f deg' % sam, flush = True)

    if results_dir is not None:
        outfile = os.path.splitext(os.path.basename(filename))[0]
        outimagefile = os.path.join(results_dir,
                                    outfile + ".reconstructed.mat")
        if not quiet:
            print('  Saving reconstructed image as %s...' % outimagefile,
                  flush = True)
        
        if os.path.isfile(outimagefile):
            os.remove(outimagefile)
        hdf.savemat(outimagefile, {'X_hat': hsi.ToNumpy(X_hat)})
    
        if output_error_image:
            errplotfile = os.path.join(results_dir,
                                       outfile + ".error.png")
            if not quiet:
                print('  Saving error image as %s...' % errplotfile,
                      flush = True)
            hsi.SaveErrorPlot(X, X_hat, errplotfile)
    
        if get_pseudocolor_image is not None:
            outimagefile = os.path.join(results_dir,
                                        outfile + ".reconstructed.png")
            X_hat = hsi.Denormalize(X_hat, Xmax, Xmin, type = "minmax")
            if not quiet:
                print('  Saving pseudocolor image as %s...' % outimagefile,
                      flush = True)
            outimage = get_pseudocolor_image(hsi.ToNumpy(X_hat))
            imageio.imwrite(outimagefile, outimage)
            
    return psnr, rmse, sam


##########################################################


def CCDDReconstruct(testing_dir,
                    model_dir,
                    results_dir,
                    sensor,
                    spatial_ratio = 8,
                    band_block = 3,
                    num_train_steps = 1000,
                    beta_max = 0.01,
                    beta_min = 0.0001,
                    num_inference_steps = 10,
                    num_landweber_iterations = 2,
                    tau = None,
                    eta = 1.0,
                    ddim = False,
                    unconstrained = False,
                    unconditioned = False,
                    transpose_adjoint= False,
                    flops = False,
                    device = None,
                    quiet = False):

    if not os.path.isdir(testing_dir):
        raise RuntimeError("Cannot find testing direcoty %s" % testing_dir)
    if not os.path.isdir(model_dir):
        raise RuntimeError("Cannot find model direcoty %s" % model_dir)
    os.makedirs(results_dir, exist_ok = True)
    
    if unconditioned and unconstrained:
        raise RuntimeError("Cannot have unconditioned and unconstrained reconstruction simultaneously")
        
    if device is None:
        device_name = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device_name)
        
    if unconstrained:
        num_landweber_iterations = 1

    if not quiet:
        print('CCDD testing (%s) - parameters:' % device, flush = True)
        if ddim:
            print('  Scheduler: DDIM', flush = True)
        else:
            print('  Scheduler: DDPM', flush = True)
        print('  Num inference steps: %d' % num_inference_steps, flush = True)
        if unconditioned:
            print('  Conditioning: unconditioned', flush = True)
        else:
            print('  Conditioning: conditioned', flush = True)
        if unconstrained:
            print('  Constraints: unconstrained', flush = True)
            print('  Num Landweber iterations: not using S or P constraints', flush = True)
        else:
            print('  Constraints: constrained', flush = True)
            print('  Num Landweber iterations: %d' % num_landweber_iterations, flush = True)
        if transpose_adjoint:
            print('  Spatial adjoint: transpose', flush = True)
        else:
            print('  Adjoint: bicubic interpolation', flush = True)
        if tau is None:
            print('  tau: tau_star', flush = True)
        else:
            print('  tau: %f' % tau, flush = True)
        print('  eta: %f\n' % eta, flush = True)
        
    # X: HSI treated as HS2I
    testing_files = sorted(glob.glob(os.path.join(testing_dir, '*.mat')))
    
    if not quiet: print('Loading model from %s...' % model_dir, flush = True)
    model_unet = diffusers.UNet2DModel.from_pretrained(model_dir,
                                                       use_safetensors = False,
                                                       device_map = device)
    num_parameters = sum(p.numel() for p in model_unet.parameters())
    if not quiet: print('  Number of UNet parameters: %.2f (million)\n' % \
                        (num_parameters / 1e6),
          flush = True)
    
    if ddim:
        noise_scheduler = diffusers.DDIMScheduler(num_train_timesteps =
                                                  num_train_steps,
                                                  beta_start = beta_min,
                                                  beta_end = beta_max)
    else:
        noise_scheduler = diffusers.DDPMScheduler(num_train_timesteps =
                                              num_train_steps,
                                              beta_start = beta_min,
                                              beta_end = beta_max)
    
    noise_scheduler.set_timesteps(num_inference_steps = num_inference_steps)

    psnrs = []
    rmses = []
    sams = []
    
    for filename in testing_files:
        psnr, rmse, sam = \
            CCDDReconstructFile(
                filename = filename,
                sensor = sensor,
                model_unet = model_unet,
                noise_scheduler = noise_scheduler,
                results_dir = results_dir,
                output_error_image = True,
                spatial_ratio = spatial_ratio,
                band_block = band_block,
                num_landweber_iterations = num_landweber_iterations,
                tau = tau,
                eta = eta,
                unconstrained = unconstrained,
                unconditioned = unconditioned,
                transpose_adjoint = transpose_adjoint,
                flops = flops,
                device = device,
                quiet = quiet)
        psnrs.append(psnr)
        rmses.append(rmse)
        sams.append(sam)

    return psnrs, rmses, sams
