#!/usr/bin/env python3

# 
# FB-LRTA: Fixed-Basis Low-Rank Tensor Approximation
# Copyright (C) 2023-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 numpy as np
import torch
import hdf5storage as hdf
import argparse
import time
import imageio
import sys
sys.path.insert(0, '.')
from Common import hsi
import fblrta


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

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('-e', '--errplot', default = None,
                        help='(string) filename of error plot (default = None)')
    parser.add_argument('-cb', '--colorbar', default = None,
                        help='(string) filename of colorbar for error plot (default = None)')
    parser.add_argument('--crop', type = int, default = None,
                        help='(int) size of upper-left crop of input image (default = None')
    parser.add_argument('--no_fixed_basis', action = 'store_true',
                        default = False,
                        help = 'use SVT rather than fixed-basis SVT')
    parser.add_argument('inimagefile',
                        help = '(string) input testing image (MAT format)')
    parser.add_argument('resultsfile',
                        help='(string) filename of output reconstructed image (MAT format)')
    parser.add_argument('outimagefile',
                        help='(string) filename of output pseudocolor image (PNG format)')
    args = parser.parse_args()
    return args

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

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

args = parse_arguments()

fixed_basis = not args.no_fixed_basis

# X: HSI treated as HS2I
print('Reading testing image file %s...' % args.inimagefile, flush = True)
testing_data = hdf.loadmat(args.inimagefile)
X = testing_data['X'].astype(np.float32)

S, remove_absorption_bands, get_pseudocolor_image = \
    hsi.ProcessSensor(args.sensor)

# Remove water-absorption bands
if remove_absorption_bands is not None:
    X = remove_absorption_bands(X)

# Crop to upper-left
if args.crop is not None:
    print('Cropping image to %dx%d' % (args.crop, args.crop), flush = True)
    X = X[:args.crop, :args.crop , :]
    
msi_num_rows, msi_num_cols, hsi_num_bands = X.shape
hsi_num_rows = msi_num_rows // args.spatial_ratio
hsi_num_cols = msi_num_cols // args.spatial_ratio
msi_num_bands = S.shape[0]

if hsi_num_rows * args.spatial_ratio != msi_num_rows or \
   hsi_num_cols * args.spatial_ratio != msi_num_cols:
    msi_num_rows = hsi_num_rows * args.spatial_ratio
    msi_num_cols = hsi_num_cols * args.spatial_ratio
    print('Image size (%d x %d) not divisible by %d -- cropping image to %dx%d' % \
          (X.shape[0], X.shape[1], args.spatial_ratio,
           msi_num_rows, msi_num_cols),
          flush = True)
    X = X[:msi_num_rows, :msi_num_cols , :]
    
# Normalize to [0, 1]
X, Xmax, Xmin = hsi.Normalize(X, type = "minmax")

P = hsi.GetSpatialProjectionAsMatrix(args.spatial_ratio,
                                     msi_num_rows,
                                     msi_num_cols).astype(np.float32)

model = fblrta.fblrta(S, P, device = device, fixed_basis = fixed_basis)

# HSI
psf = hsi.GaussianKernel(shape = (args.spatial_ratio,
                                  args.spatial_ratio)).astype(np.float32)
L = hsi.SpatialProjection(X, psf)

# MSI
H = hsi.SpectralProjection(X, S)

# Initial bicubic interpolation
X0 = hsi.BicubicInterpolation(L, args.spatial_ratio).astype(np.float32)

# Get fixed PCA-transform basis from HSI
if fixed_basis:
    print('Calculating U...', flush = True)
    L3 = hsi.Unfold3(L)
    U = np.linalg.eigh(L3 @ L3.T)[1]
    U = torch.tensor(U, dtype = torch.float32, device = device)
else:
    U = None
    
if fixed_basis:
    print('FB-LRTA reconstruction...', flush =True)
else:
    print('FB-LRTA reconstruction (no fixed basis)...', flush =True)
tic = time.time()
X_hat = model(hsi.ToTensor(X0, dtype = torch.float32, device = device),
              hsi.ToTensor(H, dtype = torch.float32, device = device),
              hsi.ToTensor(L, dtype = torch.float32, device = device),
              U)
toc = time.time()
print('    Reconstruction time = %.2f (sec)' % (toc - tic), flush = True)

X_hat = hsi.ToNumpy(X_hat)

print('    PSNR = %.2f dB' % hsi.PSNR(X, X_hat, max_val = 1.0), flush = True)
print('    RMSE = %e' % hsi.RMSE(X, X_hat), flush = True)
print('     SAM = %.4f deg' % hsi.SAM(X, X_hat), flush = True)

print('Saving reconstructed image as %s...' % args.resultsfile, flush = True)

if os.path.isfile(args.resultsfile):
    os.remove(args.resultsfile)
hdf.savemat(args.resultsfile, {'X_hat': X_hat})

if args.errplot is not None:
    hsi.SaveErrorPlot(X, X_hat, args.errplot, args.colorbar)

X_hat = hsi.Denormalize(X_hat, Xmax, Xmin, type = "minmax")
outimage = get_pseudocolor_image(X_hat)
imageio.imwrite(args.outimagefile, outimage)
