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

def MakeData():
    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, "pavia.mat")
    original_image = os.path.join(original_dir, "pavia.original.png")

    if os.path.exists(original_file):
        print("  Nothing to do - %s exists already" % original_file)
        return
    
    fileurl = "https://www.ehu.eus/ccwintco/uploads/e/ee/PaviaU.mat"
    print("Downloading %s..." % fileurl)
    if subprocess.run(['wget', '-q', '-O', original_file, fileurl]):
        print("Error downloading %s" % fileurl)
        return

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

    testing_dir = os.path.join(dirname, "Testing")
    os.makedirs(testing_dir, exist_ok = True)
    testing_file = os.path.join(testing_dir, "pavia.testing.mat")
    testing_image = os.path.join(testing_dir, "pavia.testing.png")
    
    X_original = hdf.loadmat(original_file)['X']
    outimage = rosis.GetPseudocolorImage(X_original)
    imageio.imwrite(original_image, outimage)
    
    # Training: lower-left 256x256
    X = X_original[-256:, :256, :]
    if os.path.isfile(training_file):
        os.remove(training_file)
    hdf.savemat(training_file, {'X': X})
    outimage = rosis.GetPseudocolorImage(X)
    imageio.imwrite(training_image, outimage)
    
    # Testing: upper-left 256x256
    X = X_original[:256, :256, :]
    if os.path.isfile(testing_file):
        os.remove(testing_file)
    hdf.savemat(testing_file, {'X': X})
    outimage = rosis.GetPseudocolorImage(X)
    imageio.imwrite(testing_image, outimage)
