# 
# 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 scipy.sparse.linalg as la
import numpy as np
import torch
import torch.nn as nn
import hdf5storage as hdf
import argparse
import sys
sys.path.insert(0, './Common')
import hsi
import svt


def norm2(A):
    Anorm2 = la.svds(A, 1, return_singular_vectors = False,
                     solver = 'lobpcg') ** 2
    return Anorm2.astype(np.float32)


class fblrta(nn.Module):
    @torch.inference_mode()
    def __init__(self,
                 S = None, P = None,
                 beta = 0.5,
                 gamma = 0.5,
                 num_iterations = 60,
                 device = None,
                 fixed_basis = True):
        super().__init__()

        self.beta = torch.tensor(beta, device = device)
        self.gamma = torch.tensor(gamma, device = device)
        self.num_iterations = num_iterations
        self.S = torch.tensor(S, device = device)
        self.P = torch.tensor(P, device = device)
        self.tau = torch.tensor(norm2(P) * beta + gamma * norm2(S),
                                device = device)
        self.fixed_basis = fixed_basis


    @torch.inference_mode()
    def fblrta_iteration(self, X3, H3, L3, Lambda1, Lambda2, U):

        X3 -= self.beta/self.tau * (X3 @ self.P -
                                    L3 + Lambda1/self.beta) @ self.P.adjoint()
        X3 -= self.gamma/self.tau * self.S.adjoint() @ (self.S @ X3 -
                                                        H3 + Lambda2/self.gamma)

# Use svt.svt() instead of svt.svt_fixedbasis() to use the original LRTA SVT
# rather than FB-LRTA's fixed-basis SVT
        if self.fixed_basis:
            X3 = svt.svt_fixedbasis(X3, 1/self.tau, U)
        else:
            X3 = svt.svt(X3, 1/self.tau)
        
        Lambda1 += self.beta * (X3 @ self.P - L3)
        Lambda2 += self.gamma * (self.S @ X3 - H3)
        
        return X3, Lambda1, Lambda2

    
    @torch.inference_mode()
    def forward(self, X0, H, L, U):
        
        num_bands, num_rows, num_cols = X0.shape
        
        X3 = hsi.Unfold3(X0)
        H3 = hsi.Unfold3(H)
        L3 = hsi.Unfold3(L)
        Lambda1 = torch.zeros_like(L3)
        Lambda2 = torch.zeros_like(H3)
        
        for iteration in range(self.num_iterations):
            X3, Lambda1, Lambda2 = self.fblrta_iteration(X3, H3, L3,
                                                         Lambda1, Lambda2, U)
        
        X_hat = hsi.Fold3(X3, num_rows, num_cols)
        
        return X_hat
