# 
# 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 tempfile
import tarfile
import shutil
import subprocess
import numpy as np
import hdf5storage as hdf
import imageio
import xlrd
from Utils import aviris

def MakeData():
    num_rows_extracted = 256
    num_cols_extracted = 256
    
    dirname = os.path.dirname(os.path.realpath(__file__))

    xls_file = os.path.join(dirname, "losangeles.xls")
    training_dir = os.path.join(dirname, "Training")
    testing_dir = os.path.join(dirname, "Testing")
    os.makedirs(training_dir, exist_ok = True)
    os.makedirs(testing_dir, exist_ok = True)
    
    workbook = xlrd.open_workbook(xls_file)
    worksheet = workbook.sheet_by_index(0)

    with tempfile.TemporaryDirectory() as download_dir:
        for row in range(1, worksheet.nrows):
            datasetname = worksheet.cell_value(row, 1)
            fileurl = worksheet.cell_value(row, 38)
            num_rows = int(worksheet.cell_value(row, 22))
            num_cols = int(worksheet.cell_value(row, 23))
            num_bands = 224
        
            row_origin = int(worksheet.cell_value(row, 51)) - 1
            col_origin = int(worksheet.cell_value(row, 52)) - 1
        
            training = bool(worksheet.cell_value(row, 53))
        
            filename = worksheet.cell_value(row, 54)
            
            dir = training_dir if training else testing_dir

            ext = '.training.mat' if training else '.testing.mat'
            matfile = os.path.join(dir, filename + ext)

            if os.path.exists(matfile):
                print("  Nothing to do - %s exists already" % matfile)
                continue
            
            radiance_version = worksheet.cell_value(row, 7)
        
            basename = os.path.basename(fileurl)
            downloadfilename = os.path.join(download_dir,
                                            os.path.basename(fileurl))
            dirname1  = os.path.splitext(os.path.splitext(basename)[0])[0]
            dirname2  = os.path.splitext(os.path.splitext(basename)[0])[0] + \
                'rdn_' + radiance_version
            
            print('Dataset %s, rows = %d, cols = %d' % (datasetname,
                                                        num_rows, num_cols),
                  flush = True)
            
            print('  Downloading %s as %s...' % (datasetname,
                                                 downloadfilename),
                  flush = True)
            if subprocess.run(['wget', '-q', '-P', download_dir, fileurl]):
                print("Error downloading %s" % fileurl)
                continue

            with tarfile.open(downloadfilename, mode = 'r:gz') as tf:
                # Find directory, using only first scene
                names = tf.getnames()
                filename1 = "%s/%s_sc01_ort_img" % (dirname1, dirname1)
                filename2 = "%s/%s_sc01_ort_img" % (dirname2, dirname2)
                if filename1 in names:
                    print("  Found %s" % filename1)
                    scenefile = filename1
                elif filename2 in names:
                    print("  Found %s" % filename2)
                    scenefile = filename2
                else:
                    print('  Cannot find directory and/or ort_img file')
                    continue
                print('  Processing %s as %s : %dx%d at (%d, %d)...' % 
                      (scenefile, filename,
                       num_rows_extracted, num_cols_extracted,
                       row_origin, col_origin),
                      flush = True)
                tf.extract(scenefile, path = download_dir)

            scenefile = os.path.join(download_dir, scenefile)
            with open(scenefile, 'r') as infile:
                X = np.fromfile(infile, dtype='>i2').reshape(num_rows,
                                                             num_cols,
                                                             num_bands,
                                                             order='C')
                
            X = X[row_origin:(row_origin + num_rows_extracted),
                  col_origin:(col_origin + num_cols_extracted), :]
        
            print('  Writing %s...' % matfile, flush = True)
            hdf.savemat(matfile, {'X': X})
        
            ext = '.training.png' if training else '.testing.png'
            imagefile = os.path.join(dir, filename + ext)
            image = aviris.GetPseudocolorImage(X)
            imageio.imwrite(imagefile, image)

            os.remove(downloadfilename)
            os.remove(scenefile)
