# 
# 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 time
import datasets
import diffusers
import bitsandbytes as bnb
from Utils import hsi, aviris, rosis, vnir


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


class hsi_dataset(torch.utils.data.Dataset):
    def __init__(self, X, H, Y):
        self.X = X
        self.H = H
        self.Y = Y
    def __len__(self):
        return self.X.shape[0]
    def __getitem__(self, i):
        return {"X": self.X[i],
                "H": self.H[i],
                "Y": self.Y[i]}

    
def get_patch(X, S, patch_size, band_block):
    num_bands = X.shape[0]
    X = tv.transforms.RandomCrop(patch_size)(X)
    H = hsi.SpectralProjection(X, S)
    start_band = np.random.randint(num_bands // band_block) * band_block
    return X[start_band:(start_band + band_block), :, :], H


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


def CCDDTrainLoop(dataloader,
                  model_unet,
                  optimizer,
                  noise_scheduler,
                  start_epoch = 0,
                  end_epoch = 300,
                  save_epochs = 10,
                  model_dir = None,
                  optimizer_file = None,
                  history_file = None,
                  unconditioned = False,
                  device = None,
                  quiet = False):

    losses = []
    num_train_timesteps = noise_scheduler.config.num_train_timesteps

    for epoch in range(start_epoch, end_epoch):
        epoch_tic = time.time()
        num_batches = len(dataloader)
        for batch_cnt, batch in enumerate(dataloader, start = 1):
            X = batch["X"].to(device)
            H = batch["H"].to(device)
            Y = batch["Y"].to(device)
            noise = torch.randn(X.shape, device = device)
            timesteps = torch.randint(0,
                                      num_train_timesteps,
                                      (X.shape[0],), device = device).long()
            X_hat = noise_scheduler.add_noise(X, noise, timesteps)
            if unconditioned:
                input = X_hat
            else:
                input = torch.cat((X_hat, Y, H), dim = 1)
            noise_hat = model_unet(input, timesteps, return_dict = False)[0]
            loss = torch.nn.functional.mse_loss(noise_hat, noise)
            losses.append(loss.item())
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
        epoch_toc = time.time()
        avg_loss = sum(losses[-num_batches:]) / num_batches
        epoch_time = epoch_toc - epoch_tic
        if not quiet:
            print("Epoch: %d -- loss: %f -- time %.4f (sec)" % (epoch,
                                                                avg_loss,
                                                                epoch_time),
              flush = True)
        if history_file is not None:
            with open(history_file, "a") as outfile:
                outfile.write("%f\n" % avg_loss)
        if (epoch + 1) % save_epochs == 0 and model_dir is not None:
            if not quiet:
                print("Saving model to %s..." % model_dir, flush = True)
            model_unet.save_pretrained(model_dir, safe_serialization = False)
            if optimizer_file is not None:
                torch.save(optimizer.state_dict(), optimizer_file)

    return model_unet, optimizer, losses


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


def CCDDTrain(training_dir,
              model_dir,
              sensor,
              spatial_ratio = 8,
              band_block = 3,
              num_train_steps = 1000,
              beta_max = 0.01,
              beta_min = 0.0001,
              learning_rate = 0.0001,
              batch_size = 64,
              patch_size = 32,
              num_patches = 8192,
              unet_channels = 128,
              unet_layers = 3,
              start_epoch = 0,
              end_epoch = 300,
              save_epochs = 10,
              history_file = None,
              unconditioned = False,
              device = None,
              quiet = False):

    if device is None:
        device_name = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device_name)

    if not os.path.isdir(training_dir):
        raise RuntimeError("Cannot find training direcotry %s" % training_dir)
    os.makedirs(model_dir, exist_ok = True)
    
    
    S, remove_absorption_bands, _ = hsi.ProcessSensor(sensor)

    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)
    
    resume_training = start_epoch != 0
    
    if not quiet:
        print('CCDD training (%s) - parameters:' % device, flush = True)
        print('  Num training time steps (T): %d' % num_train_steps,
              flush = True)
        print('  Beta max: %f' % beta_max, flush = True)
        print('  Beta min: %f' % beta_min, flush = True)
        print('  Learning rate: %f' % learning_rate, flush = True)
        print('  Batch size: %d' % batch_size, flush = True)
        print('  Patch size: %d' % patch_size, flush = True)
        print('  Num patches: %d' % num_patches, flush = True)
        print('  Band block: %d' % band_block, flush = True)
        print('  Initial epoch: %d' % start_epoch, flush = True)
        print('  End epoch: %d' % end_epoch, flush = True)
        print('  Num UNet input channels: %d' % unet_channels, flush = True)
        print('  Num UNet layers: %d' % unet_layers, flush = True)
        if unconditioned:
            print('  Conditioning: unconditioned', flush = True)
        print('', flush = True)
        print('Data parameters:', flush = True)
        print('  Sensor: %s' % sensor, flush = True)
        print('  Spatial ratio: %d' % spatial_ratio, flush = True)
        print('', flush = True)

    # X: HSI treated as HS2I
    if not quiet: print("Loading data...", flush = True)
    training_files = sorted(glob.glob(os.path.join(training_dir, '*.mat')))
    
    X = []
    H = []
    Y = []
    for training_file in training_files:
        X_current = hdf.loadmat(training_file)['X'].astype(np.float32)
        if remove_absorption_bands is not None:
            X_current = remove_absorption_bands(X_current)
        X_current = hsi.ToTensor(X_current, device = device)
        X_current = hsi.Normalize(X_current)[0]
        if X_current.shape[-3] != S.shape[1]:
            raise RuntimeError('Dataset %s (%d bands) does not match "%s" sensor (%d bands)' % (training_file, X_current.shape[-3], sensor, S.shape[1]))
        for patch in range(num_patches):
            current_patch, H_current = get_patch(X_current, S,
                                                 patch_size, band_block)
            L_current = hsi.SpatialProjection(current_patch, P)
            Y_current  = hsi.BicubicInterpolation(L_current, spatial_ratio)
            X.append(current_patch)
            H.append(H_current)
            Y.append(Y_current)
            
    dataset = hsi_dataset(torch.stack(X),
                          torch.stack(H),
                          torch.stack(Y))
    
    msi_num_bands = S.shape[0]
    
    dataloader = torch.utils.data.DataLoader(dataset,
                                             batch_size = batch_size,
                                             shuffle = True)
    
    noise_scheduler = diffusers.DDPMScheduler(num_train_timesteps =
                                              num_train_steps,
                                              beta_start = beta_min,
                                              beta_end = beta_max)
    
    if not resume_training:
        channels = unet_channels
        block_out_channels = (channels,) + tuple((channels := 2 * channels
                                                  for layer in
                                                  range(unet_layers - 1)))
        down_block_types = ("DownBlock2D",) * unet_layers
        up_block_types = ("UpBlock2D",) * unet_layers
        if unconditioned:
            in_channels = band_block
        else:
            in_channels = 2 * band_block + msi_num_bands
        model_unet = diffusers.UNet2DModel(
            in_channels = in_channels,
            out_channels = band_block,
            block_out_channels = block_out_channels,
            down_block_types = down_block_types,
            up_block_types = up_block_types)
        model_unet = model_unet.to(device)
    else:
        model_unet = diffusers.UNet2DModel.from_pretrained(
            model_dir,
            use_safetensors = False,
            device_map = device)

    model_unet.enable_gradient_checkpointing()
    
    optimizer = bnb.optim.Adam8bit(model_unet.parameters(), lr = learning_rate)
    optimizer_file = os.path.join(model_dir, "optimizer.pt")
    if resume_training:
        optimizer.load_state_dict(torch.load(optimizer_file))
    
    if history_file is not None:
        if os.path.exists(history_file) and not resume_training:
            os.remove(history_file)

    if not quiet: print("Training model...", flush = True)
    training_tic = time.time()
    model_unet, optimizer, losses = \
        CCDDTrainLoop(dataloader = dataloader,
                      model_unet = model_unet,
                      optimizer = optimizer,
                      noise_scheduler = noise_scheduler,
                      start_epoch = start_epoch,
                      end_epoch = end_epoch,
                      save_epochs = save_epochs,
                      model_dir = model_dir,
                      optimizer_file = optimizer_file,
                      history_file = history_file,
                      unconditioned = unconditioned,
                      device = device,
                      quiet = quiet)
    training_toc = time.time()
    training_time = training_toc - training_tic
    if not quiet:
        print('Training time = %.2f (sec)' % training_time, flush = True)
    
    if not quiet:
        print("Saving model to %s..." % model_dir, flush = True)
    model_unet.save_pretrained(model_dir, safe_serialization = False)
    torch.save(optimizer.state_dict(), optimizer_file)

    return model_unet, optimizer, losses
