/* Copyright 2015. The Regents of the University of California.
 * Copyright 2016-2017. Martin Uecker.
 * Copyright 2024-2025. Institute of Biomedical Imaging. TU Graz.
 * All rights reserved. Use of this source code is governed by
 * a BSD-style license which can be found in the LICENSE file.
 *
 * Authors: 
 * 2014 Frank Ong
 * 2014 Martin Uecker
 */

#include <stdlib.h>
#include <complex.h>
#include <math.h>
#include <stdbool.h>
#include <assert.h>

#include "misc/misc.h"
#include "misc/mri.h"
#include "misc/debug.h"

#include "num/multind.h"
#include "num/flpmath.h"
#include "num/lapack.h"
#include "num/blas.h"
#include "num/linalg.h"
#include "num/ops.h"
#include "num/iovec.h"
#include "num/fft.h"

#include "iter/thresh.h"

#include "svthresh.h"


float svthresh_nomeanu(int M, int N, float lambda, complex float* dst, const complex float* src)
{
	int MN = M * N;

	complex float* basis = md_alloc(1, (long[1]){ M }, CFL_SIZE);
	complex float* coeff = md_alloc(1, (long[1]){ N }, CFL_SIZE);
	complex float* tmp = md_alloc(1, (long[1]){ MN }, CFL_SIZE);

	for (int i = 0; i < M; i++)
		basis[i] = 1. / sqrtf((float)M);


	md_clear(1, (long[1]){ N }, coeff, CFL_SIZE);
	md_clear(1, (long[1]){ MN }, tmp, CFL_SIZE);

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			coeff[j] += basis[i] * src[i + j * M];

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			tmp[i + j * M] = src[i + j*M] - coeff[j] * basis[i];

	svthresh(M, N, lambda , dst, tmp);

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			dst[i + j * M] += coeff[j] * basis[i];

	return 0;
}


float svthresh_nomeanv(int M, int N, float lambda, complex float* dst, const complex float* src)
{
	int MN = M * N;

	complex float* basis = md_alloc(1, (long[1]){ N }, CFL_SIZE);
	complex float* coeff = md_alloc(1, (long[1]){ M }, CFL_SIZE);
	complex float* tmp = md_alloc(1, (long[1]){ MN }, CFL_SIZE);

	for (int i = 0; i < N; i++)
		basis[i] = 1. / sqrtf((float)N);

	md_clear(1, (long[1]){ M }, coeff, CFL_SIZE);
	md_clear(1, (long[1]){ MN }, tmp, CFL_SIZE);

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			coeff[i] += basis[j] * src[i + j * M];

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			tmp[i + j * M] = src[i + j * M] - coeff[i] * basis[j];


	svthresh(M, N, lambda, dst, tmp);

	for (int j = 0; j < N; j++)
		for (int i = 0; i < M; i++)
			dst[i + j * M] += coeff[i] * basis[j];


	return 0;
}

/**
 * Singular Value Thresholding
 *
 * @param M - matrix column size
 * @param N - matrix row size
 * @param lambda - regularization parameter
 * @param A - input/output matrix
 */
float svthresh(long M, long N, float lambda, complex float* dst, complex float* src) //FIXME: destroys input
{
	long minMN = MIN(M, N);
	long dimsU[3] = { M, minMN, 1};
	long dimsVT[3] = { minMN, N, 1 };
	long dimsS[3] = { minMN, 1, 1 };
//	long dimsAA[3] = {minMN, minMN,1};

	long strsVT[3];
	long strsS[3];
	md_calc_strides(3, strsVT, dimsVT, CFL_SIZE);
	md_calc_strides(3, strsS, dimsS, FL_SIZE);


	complex float* U = md_alloc_sameplace(3, dimsU, CFL_SIZE, src);
	complex float* VT = md_alloc_sameplace(3, dimsVT, CFL_SIZE, src );
	float* S = md_alloc_sameplace(3, dimsS, FL_SIZE, src);

//	complex float* AA = md_alloc_sameplace(3, dimsAA, CFL_SIZE, src );
//	lapack_normal_multiply( M, N, (M > N), (complex float (*) [])AA, (const complex float (*) [])src );
	

	// SVD
	lapack_svd_econ(M, N, (complex float (*)[])U, (complex float (*)[])VT, S, (complex float (*)[N])src);

	// Thresh S
	md_softthresh(3, dimsS, lambda, 0, S, S);

	// VT = S * VT
	md_mul2(3, dimsVT, strsVT, (float*)VT, strsVT, (float*)VT, strsS, S);
	md_mul2(3, dimsVT, strsVT, (float*)VT + 1, strsVT, (float*)VT + 1, strsS, S);

	// dst = U * VT
	blas_matrix_multiply(M, N, minMN, (complex float (*)[N])dst, (const complex float (*)[minMN])U, (const complex float (*)[N])VT);

	md_free(U);
	md_free(VT);
	md_free(S);

	return 0;
}





float nuclearnorm(long M, long N, /* const */ complex float* d)
{ // FIXME: destroys input

	long minMN = MIN(M,N);
	long dimsU[3]	= { M, minMN, 1 };
	long dimsVT[3]	= { minMN, N, 1 };
	long dimsS[3]	= { minMN, 1, 1 };

	complex float* U = md_alloc_sameplace(3, dimsU, CFL_SIZE, d);
	complex float* VT = md_alloc_sameplace(3, dimsVT, CFL_SIZE, d );
	float* S = md_alloc_sameplace(3, dimsS, FL_SIZE, d);

	// SVD
	lapack_svd_econ(M, N, (complex float (*) []) U, (complex float (*) []) VT, S, (complex float (*) [N])d);

	float nnorm = 0.;
	for (int i = 0; i < minMN; i++)
		nnorm += S[i];

	md_free(U);
	md_free(VT);
	md_free(S);

	return nnorm;
}


float maxsingular(long M, long N, /*const*/ complex float* d)
{	// FIXME: destroys input
	long dimsU[2] = { M, N };
	long dimsV[2] = { N, N };
  
	complex float* U = md_alloc(2, dimsU, sizeof(complex float));
	complex float* VT = md_alloc(2, dimsV, sizeof(complex float));

	float* S = xmalloc(sizeof(float[MIN(M, N)]));

	// SVD
	lapack_svd_econ(M, N, (complex float (*)[])U,
		 (complex float (*)[]) VT, S,
		 (complex float (*)[N])d);

	float value = S[0];

	md_free(U);
	md_free(VT);
	xfree(S);

	return value;
}



/***********
 * Blockproc functions
 ************/

struct svthresh_blockproc_data {
	unsigned long mflags;
	float lambda;
	int remove_mean;
};

struct svthresh_blockproc_data* svthresh_blockproc_create(unsigned long mflags, float lambda, int remove_mean)
{
	PTR_ALLOC(struct svthresh_blockproc_data, data);
	data->mflags = mflags;
	data->lambda = lambda;
	data->remove_mean = remove_mean;
	return data;
}

float svthresh_blockproc(const void* _data, const long blkdims[DIMS], complex float* dst, const complex float* src0)
{
	const struct svthresh_blockproc_data* data = _data;

	assert(dst == src0);
	complex float* src = dst;

	int M = 1;
	int N = (int)md_calc_size(DIMS, blkdims);

	for (int i = 0; i < DIMS; i++) {

		if (!MD_IS_SET(data->mflags, i))
			continue;

		M *= blkdims[i];
		N /= blkdims[i];
	}

	switch (data->remove_mean) {
	case 1:
		svthresh_nomeanu(M, N, data->lambda , dst, src);
		break;
	case 2:
		svthresh_nomeanv(M, N, data->lambda , dst, src);
		break;
	case 0:
		svthresh(M, N, data->lambda , dst, src);
		break;
	default:
		assert(0);
	}

	return 0;
}



float nucnorm_blockproc(const void* _data, const long blkdims[DIMS], complex float* dst, const complex float* src0)
{
	const struct svthresh_blockproc_data* data = _data;

	assert(dst == src0);
	complex float* src = dst;

	long M = 1;
	long N = md_calc_size(DIMS, blkdims);

	for (int i = 0; i < DIMS; i++) {

		if (!MD_IS_SET(data->mflags, i))
			continue;

		M *= blkdims[i];
		N /= blkdims[i];
	}

	float G = sqrtf(M) + sqrtf(N);

	return G * nuclearnorm(M, N, src);
}


