# 
# 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 numpy as np
import torch
from numpy import linalg as npla
import warnings
warnings.filterwarnings('ignore')
from oct2py import octave


def svt(Z, tau):
    m = Z.shape[-2]
    n = Z.shape[-1]
    if m < n:
        ZZH = Z @ Z.adjoint()
        S2, U = torch.linalg.eigh(ZZH)
        I = S2 > torch.square(tau)
        S = torch.where(I, torch.sqrt(S2), 0.)
        S[I] = (S[I] - tau) / S[I]
        S = torch.diag_embed(S)
        return U @ S @ U.adjoint() @ Z
    else:
        ZZH = Z.adjoint() @ Z
        S2, U = torch.linalg.eigh(ZZH)
        I = S2 > torch.square(tau)
        S = torch.where(I, torch.sqrt(S2), 0.)
        S[I] = (S[I] - tau) / S[I]
        S = torch.diag_embed(S)
        return Z @ U @ S @ U.adjoint()
 

def svt_fixedbasis(Z, tau, U):
    m = Z.shape[-2]
    n = Z.shape[-1]
    if m < n:
        S = torch.linalg.vector_norm(U.adjoint() @ Z, dim = -1)
        I = S > tau
        S = torch.where(I, S, 0.)
        S[I] = (S[I] - tau) / S[I]
        S = torch.diag_embed(S)
        return U @ S @ U.adjoint() @ Z
    else:
        S = torch.linalg.vector_norm(Z @ U, dim = -2)
        I = S > tau
        S = torch.where(I, S, 0.)
        S[I] = (S[I] - tau) / S[I]
        S = torch.diag_embed(S)
        return Z @ U @ S @ U.adjoint()
        
