# 
# 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 zipfile
import shutil
import subprocess
import glob
import numpy as np
import hdf5storage as hdf
import imageio
from Utils import vnir

def MakeData():
    num_rows_extracted = 256
    num_cols_extracted = 256
    origin_row = 300
    origin_col = 1700

    dirname = os.path.dirname(os.path.realpath(__file__))

    original_dir = os.path.join(dirname, "Original")
    os.makedirs(original_dir, exist_ok = True)

    original_file = os.path.join(original_dir, "chikusei.mat")
    original_image = os.path.join(original_dir, "chikusei.original.png")

    fileurl = "http://park.itc.u-tokyo.ac.jp/sal/hyperdata/Hyperspec_Chikusei_MATLAB.zip"
    zipdir = "Chikusei_MATLAB"
    matfile = "HyperspecVNIR_Chikusei_20140729.mat"

    if os.path.exists(original_file):
        print("  Nothing to do - %s exists already" % original_file)
        return

    with tempfile.TemporaryDirectory() as tmp_dir:
        zfn = os.path.join(tmp_dir, os.path.basename(fileurl))
        print("Downloading %s as %s..." % (fileurl, zfn))
        if subprocess.run(['wget', '-q', '-P', tmp_dir, fileurl]):
            print("Error downloading %s" % fileurl)
            return
        zf = zipfile.ZipFile(zfn)
        print("Extracting %s..." % matfile)
        zf.extract("%s/%s" % (zipdir, matfile), path = tmp_dir)
        X_original = hdf.loadmat(os.path.join(tmp_dir, zipdir, matfile))["chikusei"].astype(np.float32)
        print("Saving %s..." % original_file)
        hdf.savemat(original_file, {'X': X_original.astype(np.uint16)})

    training_dir = os.path.join(dirname, "Training")
    os.makedirs(training_dir, exist_ok = True)
    training_file = os.path.join(training_dir, "chikusei.training.mat")
    training_image = os.path.join(training_dir, "chikusei.training.png")

    testing_dir = os.path.join(dirname, "Testing")
    os.makedirs(testing_dir, exist_ok = True)
    testing_file = os.path.join(testing_dir, "chikusei.testing.mat")
    testing_image = os.path.join(testing_dir, "chikusei.testing.png")
    
    outimage = vnir.GetPseudocolorImage(X_original)
    imageio.imwrite(original_image, outimage)
    
    # Training: upper-right 256x256
    X = X_original[origin_row:(origin_row + num_rows_extracted),
                               origin_col:(origin_col + num_cols_extracted), :]
    if os.path.isfile(training_file):
        os.remove(training_file)
    hdf.savemat(training_file, {'X': X})
    outimage = vnir.GetPseudocolorImage(X)
    imageio.imwrite(training_image, outimage)
    
    # Testing: lower-right 256x256
    X = X_original[-origin_row:(-origin_row + num_rows_extracted),
                   origin_col:(origin_col + num_cols_extracted), :]
    if os.path.isfile(testing_file):
        os.remove(testing_file)
    hdf.savemat(testing_file, {'X': X})
    outimage = vnir.GetPseudocolorImage(X)
    imageio.imwrite(testing_image, outimage)
