/* Copyright 2014-2017. The Regents of the University of California.
 * Copyright 2015-2021. Uecker Lab. University Medical Center Göttingen.
 * Copyright 2022-2025. TU Graz. Institute of Biomedical Imaging.
 * All rights reserved. Use of this source code is governed by
 * a BSD-style license which can be found in the LICENSE file.
 *
 * Authors:
 * 2013-2016 Martin Uecker
 * 2015, 2017 Jonathan Tamir
 */

#include <stdbool.h>
#include <complex.h>

#include "num/multind.h"
#include "num/flpmath.h"
#include "num/fft.h"
#include "num/init.h"

#include "misc/mri.h"
#include "misc/mmio.h"
#include "misc/misc.h"
#include "misc/opts.h"
#include "misc/nested.h"




static const char help_str[] = "Perform homodyne reconstruction along dimension dim.";





struct wdata {

	float frac;
	int pfdim;
	long wdims[DIMS];
	long wstrs[DIMS];
	complex float* weights;
};


/**
 * Applies the Homodyne filter.
 * @param N k-space dimension
 * @param p k-space position
 * @param frac is the fraction of acquired k-space
 * @param alpha is the offset of the ramp, between 0 and 1
 * @param clear clear acquired k-space
 *
 * The ramp portion is given by 2*(alpha - 1) / (end - start) * (p - end) + alpha
 * alpha = 0 is a full ramp, alpha = 1 is a horizontal line
 */
static float homodyne_filter(int N, float frac, float alpha, bool clear, long p)
{
	if (frac <= 0.5)
		return 1.;

	float start = (float)N * (1 - frac);
	float end = (float)N * frac;

	float ret = clear ? 0. : 1.;


	if (p < start)
		ret = 2.;
	else if (p >= start && p < end)
		ret = 2 * (alpha - 1) / (end - start) * ((float)p - end) + alpha;

	return ret;
}



static complex float* estimate_phase(struct wdata wdata, unsigned long flags,
		int N, const long dims[N], const complex float* idata, bool center_fft)
{
	long cdims[N];
	md_copy_dims(N, cdims, dims);
	// cdims[0] = cdims[1] = cdims[2] = 24;
	cdims[wdata.pfdim] = (wdata.frac - 0.5) * (double)dims[wdata.pfdim];

	complex float* center = md_alloc(N, cdims, CFL_SIZE);
	complex float* phase = md_alloc(N, dims, CFL_SIZE);

	md_resize_center(N, cdims, center, dims, idata, CFL_SIZE);
	md_resize_center(N, dims, phase, cdims, center, CFL_SIZE);
	md_free(center);

	(center_fft ? ifftuc : ifftu)(N, dims, flags, phase, phase);
	md_zphsr(N, dims, phase, phase);

	return phase;
}

static void homodyne(struct wdata wdata, unsigned long flags, int N, const long dims[N],
		const long strs[N], complex float* data, const complex float* idata,
		const long pstrs[N], const complex float* phase, bool center_fft)
{
	md_zmul2(N, dims, strs, data, strs, idata, wdata.wstrs, wdata.weights);
	(center_fft ? ifftuc : ifftu)(N, dims, flags, data, data);

	md_zmulc2(N, dims, strs, data, strs, data, pstrs, phase);
	md_zreal(N, dims, data, data);
}



int main_homodyne(int argc, char* argv[argc])
{
	int pfdim = -1;
	float frac = 0.f;
	const char* in_file = NULL;
	const char* out_file = NULL;

	struct arg_s args[] = {

		ARG_INT(true, &pfdim, "dim"),
		ARG_FLOAT(true, &frac, "fraction"),
		ARG_INFILE(true, &in_file, "input"),
		ARG_OUTFILE(true, &out_file, "output"),
	};

	bool clear = false;
	bool image = false;
	bool center_fft = true;
	const char* phase_ref = NULL;

	float alpha = 0.;

	num_init();

	const struct opt_s opts[] = {

		OPT_FLOAT('r', &alpha, "alpha", "Offset of ramp filter, between 0 and 1. alpha=0 is a full ramp, alpha=1 is a horizontal line"),
		OPT_SET('I', &image, "Input is in image domain"),
		OPT_SET('C', &clear, "Clear unacquired portion of kspace"),
		OPT_INFILE('P', &phase_ref, "phase_ref>", "Use <phase_ref> as phase reference"),
		OPT_CLEAR('n', &center_fft, "use uncentered ffts"),
	};

	cmdline(&argc, argv, ARRAY_SIZE(args), args, help_str, ARRAY_SIZE(opts), opts);


	const int N = DIMS;
	long dims[N];
	complex float* idata = load_cfl(in_file, N, dims);
	complex float* data = create_cfl(out_file, N, dims);

	assert((0 <= pfdim) && (pfdim < N));
	assert(frac > 0.);

	if (image) {

		(center_fft ? fftuc : fftu)(N, dims, FFT_FLAGS, idata, idata);
	}


	long strs[N];
	md_calc_strides(N, strs, dims, CFL_SIZE);

	struct wdata wdata;
	wdata.frac = frac;
	wdata.pfdim = pfdim;
	md_select_dims(N, MD_BIT(pfdim), wdata.wdims, dims);
	md_calc_strides(N, wdata.wstrs, wdata.wdims, CFL_SIZE);
	wdata.weights = md_alloc(N, wdata.wdims, CFL_SIZE);

	NESTED(void, comp_weights, (const long pos[]))
	{
		wdata.weights[md_calc_offset(DIMS, wdata.wstrs, pos) / (long)CFL_SIZE]
			= homodyne_filter((int)wdata.wdims[pfdim], frac, alpha, clear, pos[pfdim]);
	};

	md_loop(N, wdata.wdims, comp_weights);

	long pstrs[N];
	long pdims[N];
	complex float* phase = NULL;

	if (NULL == phase_ref) {

		phase = estimate_phase(wdata, FFT_FLAGS, N, dims, idata, center_fft);

		md_copy_dims(N, pdims, dims);

	} else {

		phase = load_cfl(phase_ref, N, pdims);
	}

	md_calc_strides(N, pstrs, pdims, CFL_SIZE);

	homodyne(wdata, FFT_FLAGS, N, dims, strs, data, idata, pstrs, phase, center_fft);

	md_free(wdata.weights);

	if (NULL == phase_ref) {

		md_free(phase);

	} else {

		unmap_cfl(N, pdims, phase);

		xfree(phase_ref);
	}

	unmap_cfl(N, dims, idata);
	unmap_cfl(N, dims, data);

	return 0;
}

