#!/usr/bin/env python3

# 
# 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 sys
import torch
import argparse
from CCDD import ccdd_train as ccdd


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


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--sensor', default = 'aviris',
                        help='(string) sensor type (default = "aviris")')
    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('-lr', '--learning_rate', type=float, default=0.0001,
                        help = '(float) learning rate of ADAM optimizer (default = 0.0001')
    parser.add_argument('-bs', '--batch_size', type=int, default=64,
                        help = '(int) training batch size (default = 64)')
    parser.add_argument('-ps', '--patch_size', type=int, default=32,
                        help = '(int) patch size (default = 32)')
    parser.add_argument('-np', '--num_patches', type=int, default=8192,
                        help = '(int) number of patches per image (default = 8192)')
    parser.add_argument('-bb', '--band_block', type=int, default=3,
                        help = '(int) number of bands to block together (default = 3)')
    parser.add_argument('-uc', '--unet_channels', type=int, default=128,
                        help = '(int) UNet input channels (default = 128)')
    parser.add_argument('-ul', '--unet_layers', type=int, default=3,
                        help = '(int) num UNet layers (default = 3)')
    parser.add_argument('-e', '--epochs', type=int, default=300,
                        help = '(int) number of training epochs (default = 300)')
    parser.add_argument('--unconditioned', action = 'store_true',
                        default = False,
                        help = 'no conditioning')
    parser.add_argument('-sr', '--spatial_ratio', type=int, default=8,
                        help = '(int) spatial subsampling ratio (default = 8)')
    parser.add_argument('-r', '--resume', type=int, default=0,
                        help = '(int) resume training at given epoch')
    parser.add_argument('-se', '--save_epochs', type=int, default=10,
                        help = '(int) epoch interval for saving model')
    parser.add_argument('-hs', '--history_file', type=str, default=None,
                        help = '(string) filename of output training history (default is do not save history')
    parser.add_argument('training_dir',
                        help = '(string) directory of input training files')
    parser.add_argument('model_dir',
                        help='(string) directory of output model')
    args = parser.parse_args()
    return args


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


def main():

    args = parse_arguments()

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

    model_unet, optimizer, losses = \
        ccdd.CCDDTrain(training_dir = args.training_dir,
                       model_dir = args.model_dir,
                       sensor = args.sensor,
                       spatial_ratio = args.spatial_ratio,
                       band_block = args.band_block,
                       num_train_steps = args.num_train_steps,
                       beta_max = args.beta_max,
                       beta_min = args.beta_min,
                       learning_rate = args.learning_rate,
                       batch_size = args.batch_size,
                       patch_size = args.patch_size,
                       num_patches = args.num_patches,
                       unet_channels = args.unet_channels,
                       unet_layers = args.unet_layers,
                       start_epoch = args.resume,
                       end_epoch = args.epochs,
                       save_epochs = args.save_epochs,
                       history_file = args.history_file,
                       unconditioned = args.unconditioned,
                       device = device,
                       quiet = False)

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


if __name__ == "__main__":
    dirname = os.path.dirname(os.path.realpath(__file__))
    sys.path.insert(0, dirname)
    main()
