#!/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_reconstruct as ccdd


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


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 main():

    args = parse_arguments()

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

    ccdd.CCDDReconstruct(
        testing_dir = args.testing_dir,
        model_dir = args.model_dir,
        results_dir = args.results_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,
        num_inference_steps = args.num_inference_steps,
        num_landweber_iterations = args.num_landweber_iterations,
        tau = args.tau,
        eta = args.eta,
        ddim = args.ddim,
        unconstrained = args.unconstrained,
        unconditioned = args.unconditioned,
        transpose_adjoint = args.transpose_adjoint,
        flops = args.flops,
        device = device,
        quiet = False)

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


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