  function [psf, fwhm, var, con] = qpwls_psf_slow_st(A, C, beta, mask, W, varargin)
%|function [psf, fwhm, var, con] = qpwls_psf_slow_st(A, C, beta, mask, W, [option])
%|
%| Compute exact PSF as per fessler:96:srp
%| for quadratically-penalized shift-invariant unweighted least-squares
%| uses psf = [A'WA + beta C'C]^{-1} A'W y, where usually y = A e
%|
%| PSF is computed for pixel at position "offset" relative to
%| the "center" pixel associated with an FFT psf(nx/2+1, ny/2+1)
%| in
%|	C			can be C matrix, or C'Cej, or R object (also: C = {C_R, C_I})
%|	beta [1,2]	beta = [beta_R, beta_I];
%|
%| options
%|	'offset'	coordinates of pixel of interest (default in 2d: [0 0])
%|	'dx'		pixel size, to report fwhm in physical units (def: 1)
%|	'dz'		slice spacing, report z.fwhm in physical units (def: 1)
%|	'yb'		substitute this data for A e (default: uses A*e)
%|			this is useful if recon model differs from "truth"!
%|	'shift0'	PSF centered at 0, even if 'offset' (default: 0)
%|	'psftype'	'reale' (default) or 'abs' determines which function
%|			is used to convert psf from complex-valued to real-valued
%|			if psf spectrum is not hermitian symmetric
%|	'fwhmtype'	'slice' (default) or 'profile' (new) or 'fwhm3' (not)
%|			or '2z' to return fwhm of each plane of "3d" psf
%|			or 'none' to prevent fwhm computation
%|	'gtwg'		user-provided A'*W*A
%|	'ihes_gtwg'	user-provided (A'WA + beta C'C)^{-1} A'WA
%|	'loop'	0|1	if 1, then loop to prompt for other betas.  (default 0)
%|
%| out
%|	psf
%|	var	approximate relative variance for the center pixel
%|	fwhm	FWHM of psf
%|	con	approximate condition number
%|
%| caution: fwhm likely inaccurate if the psf is centered near the image edge
%| yet is wide enough to overlap the edge.  'shift0' may help in such cases.
%|
%| Copyright 2009-8-14, Valur Olafsson, The University of Michigan

% todo: option for separating "x" and "y" spatial resolutions in 2d

if nargin == 1 && streq(A, 'test'), qpwls_psf_slow_st_test; return, end
if nargin < 4, help(mfilename), error(mfilename), end

arg.chat = 0;
arg.dx = 1; % 1 pixel distance unit
arg.dz = 1; % 1 pixel distance unit
arg.iwarn = 'disp';
arg.offset = zeros(size(size(mask)));
arg.yb = [];
arg.gtwg = [];
arg.ihes_gtwg = [];
arg.shift0 = false;
arg.psftype = 'reale';
arg.fwhmtype = 'slice';
arg.loop = false;

% support old-style argument
if length(varargin) && isnumeric(varargin{1})
	arg.offset = varargin{1};
	varargin = {varargin{2:end}};
end

arg = vararg_pair(arg, varargin);

if ~isvar('W') || isempty(W)
	W = 1;	% cheap identity
end

if length(arg.offset) ~= ndims(mask), error 'bad offset dim', end

if ~islogical(mask), warn 'not logical mask', mask = logical(mask); end

% support running with multiple beta values
betaRemaining= beta(2:end,:);
beta = beta(1,:);

%
% unit vector at a pixel
% relative to 'center' pixel in FFT sense
%
dim = size(mask);
ndim = length(dim);
np = sum(mask(:));
ej = zeros(dim);
ej( subv2ind(dim, floor(dim/2) + 1 + arg.offset) ) = 1;
if ~sum(ej(:) .* mask(:))
	warn 'offset outside mask'
	psf = nan(size(mask));
	var = NaN;
	fwhm = NaN;
return
end

if isempty(arg.ihes_gtwg)
	% gram matrix: F = A'WA
	if isempty(arg.gtwg)
		% Fisher information: A'WA
		gtg = A' * (W * A);

	else % user provided (to save time for repeated calls)
		gtg = arg.gtwg;
	end

	%
	% regularization matrix: R = C'C
	%
	% make C into a cellarray
	if iscell(C)
		if length(C) > 2
			error 'can only have 2 different regularizers'
		end
	else
		C = {C};
	end
	nC = length(C);
	ctc = cell(1,nC);
	for ii = 1:nC
		ctc{ii} = form_ctc(C{ii}, beta(ii), ej, mask, size(gtg));
	end


	% stacked "full-matrix" Fisher information
	gtg_S = [real(gtg) -imag(gtg); ...
			 imag(gtg)  real(gtg)];

	% stacked "full-matrix" of penalty
	if nC == 1
		ctc{2} = ctc{1};
	end

	if ~isempty(arg.yb)
		gtwy = A' * (W * arg.yb(:));
	end
end

% now loop to allow user to interactively pick betas
beta_original = beta;
repeater = true;
beta_list = [];
fwhm_list = [];
if ~isempty(betaRemaining)
	psf_list = zeros([dim 2 length(betaRemaining)+1]);
	var_list = zeros(length(betaRemaining)+1, 2);
	iBeta = 1;
end

while repeater
	if ~isempty(arg.ihes_gtwg)
		psfR = arg.ihes_gtwg * [ej(mask); zeros(np,1)];
		psfI = arg.ihes_gtwg * [zeros(np,1); ej(mask)];
	else
		% hessian
		hes = gtg_S + blkdiag(beta(1)*ctc{1}, beta(2)*ctc{2});
		%hes = [real(gtg)+beta(1)*ctc{1} -imag(gtg); ...
		%	   imag(gtg) real(gtg)+beta(2)*ctc{2}];

		con = 1/rcond(hes);
		printm('approximate condition number %g', con)

		if ~isempty(arg.yb)
			psfR = hes \ [real(gtwy); imag(gtwy)];		% faster and more accurate (according to MATLAB help)
			%psfR = inv(hes) * [real(gtwy); imag(gtwy)];	% slower calculation
			psfI = hes \ [-imag(gtwy); real(gtwy)];		% faster and more accurate (according to MATLAB help)
			%psfI = inv(hes) * [-imag(gtwy); real(gtwy)];	% slower calculation
		else

			% psf of real part, where e_{R,j} = [0 ... 0 1 0 ... 0],  e_{I,j} = zeros
			psfR = hes \ (gtg_S * [ej(mask); zeros(np,1)]);		% faster and more accurate (according to MATLAB help)
			%psfR = inv(hes) * (gtg_S * [ej(mask); zeros(np,1)]);
			% psf of imag part, where  e_{R,j} = zeros, e_{I,j} = [0 ... 0 1 0 ... 0]
			psfI = hes \ (gtg_S * [zeros(np,1); ej(mask)]);		% faster and more accurate (according to MATLAB help)
			%psfI = inv(hes) * (gtg_S * [zeros(np,1); ej(mask)]);
		end
	end

	% form psf for real part (psfR) and for imaginary part (psfI)
	psfR = embed(psfR(1:end/2) + 1i*psfR(end/2+1:end), mask);
	psfI = embed(psfI(1:end/2) + 1i*psfI(end/2+1:end), mask);

	% note: psfR and psfI could be complex if spectrum is not hermitian symmetric
	if streq(arg.psftype, 'reale')
		psfR = reale(psfR, arg.iwarn);
		psfI = reale(-1i * psfI, arg.iwarn);	% grab imaginary part
	elseif streq(arg.psftype, 'abs')
		psfR = abs(psfR);
		psfI = abs(psfI);
	else
		error 'unknown psftype'
	end

	psf = cat(ndim + 1, psfR, psfI);

	if arg.dx == 1
		units = '[pixel]';
	else
		units = '[dis]';
	end

	if streq(arg.fwhmtype, 'none')
		fwhm = [];
	else
		fwhm = qpwls_psf_fwhm(psfR, psfI, ndim, arg, units); % report fwhm
	end

	% predicted relative variance for white noise based on local Fourier analysis
	if nargout > 2
		ao_R = cat(1, ones(np,1), zeros(np,1));
		ao_I = cat(1, zeros(np,1), ones(np,1));
		var_R = 1/np .* (ao_R' * (hes \ (gtg_S * (hes \ ao_R))));
		var_I = 1/np .* (ao_I' * (hes \ (gtg_S * (hes \ ao_I))));
		var = [var_R, var_I];
	end

	% if desired, display images, primarily for debug
	if arg.chat && im
		clf, im pl 3 3
		ix = [-dim(1)/2:dim(1)/2-1] - arg.offset(1);
		iy = [-dim(2)/2:dim(2)/2-1] - arg.offset(2);
	%	im(1, ix, iy, abs(gtg), '|A''WA|'), cbar h
		im(1, ix, iy, real(embed(gtg*ej(mask),mask)), 'real(A''WA)'), cbar h
		im(2, ix, iy, imag(embed(gtg*ej(mask),mask)), 'imag(A''WA)'), cbar h
		im(3, ix, iy, cat(3, embed(ctc{1}*ej(mask),mask), embed(ctc{2}*ej(mask),mask)), 'real(C''C) | imag(C''C)'), cbar h
		im(4, ix, iy, psf(:,:,1), 'psf_R'), cbar h
		im(5, ix, iy, psf(:,:,2), 'psf_I'), cbar h
	end

	if arg.loop
		beta_list(end+1,:) = beta(:)';
		fwhm_list(end+1,:) = fwhm(:)';
		[beta_list_tmp tmp] = sort(beta_list(:,1));
		beta_list = beta_list(tmp,:);
		fwhm_list = fwhm_list(tmp,:);
		pr beta_list
		pr fwhm_list
		if 1 && im
			clamp = @(x,a,b) min(max(x, a), b);
			for id=1:ndims(psf)-1
				ii{id} = ceil(dim(id)/2+1) + [-10:10] + arg.offset(id);
				ii{id} = ii{id}(ii{id} >= 1 & ii{id} <= dim(id));
			end
			clf, im pl 1 2, im(psf(ii{:},1)), im(psf(ii{:},2))
		end
		tmp = input('enter new beta scales [q]: ', 's');
		if isempty(tmp) || streq(tmp, 'q')
			repeater = false;
		else
			beta_scale = eval(['[' tmp ']']);
			pr beta_scale
			beta = beta_original .* beta_scale;
		end
	elseif ~isempty(betaRemaining)
		if ndim == 2
			psf_list(:,:,:,iBeta) = psf;
		elseif ndim == 3
			psf_list(:,:,:,:,iBeta) = psf;
		end
		if (nargout > 2)
			var_list(iBeta,:) = var(:).';
		end
		fwhm_list(end+1,:) = fwhm(:)';
		beta = betaRemaining(1,:);
		betaRemaining = betaRemaining(2:end,:);
	else
		if exist('psf_list', 'var')
			if ndim == 2
				psf_list(:,:,:,iBeta) = psf;
			elseif ndim == 3
				psf_list(:,:,:,:,iBeta) = psf;
			end
			psf = psf_list;
			if (nargout > 2)
				var = cat(1, var_list, var(:).');
			end
			fwhm = cat(1, fwhm_list, fwhm(:)');
		end
		repeater = false;
	end
end



% form_ctc()
% determine C' * C
function ctc = form_ctc(C, beta, ej, mask, size_gtg)

dim = size(mask);

if isnumeric(C) % sparse (or full) matrix
	if isequal(size(C), dim)
		ctc = C; % here "C" is impulse response of R for center pixel
	else
		ctc = C' * C;
	end
elseif streq(class(C), 'strum') ... % new Reg1 strum
	|| isstruct(C) % Robject

	error 'C has to be a numeric matrix'

elseif isa(C, 'fatrix2')
	error 'C has to be a numeric matrix'

elseif isa(C, 'Fatrix')
	error 'C has to be a numeric matrix'
else
	error 'unknown regularizer'
end


%
% qpwls_psf_fwhm()
%
function fwhm = qpwls_psf_fwhm(psfR, psfI, ndim, arg, units)

switch ndim
case 2
	try
		fwhmR = fwhm2(psfR, 'dx', arg.dx, ...
			'type', 'user', ...
			'center', floor(size(psfR)/2)+1+arg.offset);
		fwhmI = fwhm2(psfI, 'dx', arg.dx, ...
			'type', 'user', ...
			'center', floor(size(psfI)/2)+1+arg.offset);
		fwhm = [fwhmR, fwhmI];
		if arg.dx == 1
			printm('fwhm = [%g %g] [pixel]', fwhm)
		else
			printm('fwhm = [%g %g] [pixel] or [%g %g] [dis]', ...
				fwhm./arg.dx, fwhm)
		end
	catch
		warn 'fwhm2 failed, returning NaN for fwhm'
		fwhm = [NaN, NaN];
	end
case 3
	warn '''fwhm'' is not implemented for 3d, returning NaN for fwhm'
	fwhm = [NaN NaN];
	return

	%% CANNOT RUN THE REST OF THIS SINCE NOT SETUP FOR 3D RECON %%
	if streq(arg.fwhmtype, '2z') % each slice of psf separately
		for ii = 1:size(psf,3)
			fwhm(ii) = qpwls_psf_fwhm(psf(:,:,ii), 2, arg, units);
		end
	return
	end

	if ~arg.shift0 && any(arg.offset) && ~streq(arg.fwhmtype, 'fwhm3')
		warn('3d "fwhm" probably wrong without shift0 option')
	end

	try
		dd = [arg.dx arg.dx arg.dz];
		switch arg.fwhmtype

		case 'none'
			fwhm = [];

		case 'profile'
			% new: use fwhm1 for 1d profiles through psf
			tmp = num2cell(floor(size(psf)/2+1));
			[ix iy iz] = deal(tmp{:});
			fwhm = [
				fwhm1(psf(:,iy,iz), 'dx', dd(1))
				fwhm1(psf(ix,:,iz), 'dx', dd(2))
				fwhm1(psf(ix,iy,:), 'dx', dd(3))
			];

		case 'slice' % old: report separate yz, xz, xy fwhm!
			% old: use fwhm2 for 2d slices through psf
			tmp = @(psf, i1, i2) ...
				fwhm2(squeeze(psf), 'dx', dd(i1), 'dy', dd(i2));
			warn 'this is the default, but you probably want the ''profile'' option!'
			fwhm = [
				tmp(psf(floor(end/2)+1,:,:), 2, 3)
				tmp(psf(:,floor(end/2)+1,:), 1, 3)
				tmp(psf(:,:,floor(end/2)+1), 1, 2)
			];

		case 'oldslice'
			fwhm = [
				fwhm2(squeeze(psf(floor(end/2)+1,:,:)))
				fwhm2(squeeze(psf(:,floor(end/2)+1,:)))
				fwhm2(squeeze(psf(:,:,floor(end/2)+1)))];

		case 'fwhm3'
			fwhm = fwhm3(psf);
%			im(psf(:,:,floor(end/2)-1+[-3:3])), cbar

		otherwise
			error 'unknown fwhmtype'
		end
		printm('fwhm = %g %g %g %s', fwhm, units)
	catch
		warn 'fwhm(3d) failed, returning NaN for fwhm'
		fwhm = [NaN NaN];
	end

otherwise
	fwhm = [NaN NaN];
	warn 'fwhm not implemented'
end


%
% self test; also illustrates use
%
function psf = qpwls_psf_slow_st_test
runMri = true;
if runMri
	% form simulation maps
	dir = [path_find_dir('mri') '/phase-data/'];
	wtrue = 2*pi * fld_read([dir 'fieldmap128.fld']); % "true" fieldmap
	mag = fld_read([dir 'mag128.fld']); % "true" magnitude
	dim = size(mag);
	[X Y] = meshgrid([0:dim(1)-1]-dim(1)/2+1/2, [0:dim(2)-1]-dim(2)/2+1/2);
	Xnew = X(1:2:end,1:2:end)+1/2; Ynew = Y(1:2:end,1:2:end)+1/2;
	wtrue = interp2(X, Y, wtrue, Xnew, Ynew);
	mag = interp2(X, Y, mag, Xnew, Ynew);
	r2strue = mag ./ max(mag(:)) .* 25;

	[nx ny] = size(mag);
	ig = image_geom('nx', nx, 'ny', ny, 'dx', 22/nx, 'dy', 22/nx);
	ig.mask = abs(mag) > 0.05 * max(abs(mag(:))); % stress it with mask
	z = (r2strue + 1i*wtrue) .* ig.mask;
	[kspace, omega] = mri_kspace_spiral('N', ig.nx, 'fov', ig.fov, 'dt', 4e-6);
	nk = size(kspace,1);
	kspace = kspace ./ ig.fovs(ones(nk,1),:);
	tk = 20e-3+[0:nk-1]*4e-6;

	%sg.plot(ig), prompt

	A = Gmri(kspace, ig.mask, 'fov', ig.fovs, 'basis', {'rect', [ig.dx ig.dy]}, ...
			 'ti', tk, 'zmap', z, 'L', {1, 0.01});

	beta = [2^3 2^6];
	offset = [3 -3];
else
	ig = image_geom('nx', 64, 'ny', 32, 'dx', 1, 'dy', 1);
	tmp = ig.mask; tmp(1) = logical(0); ig.mask = tmp; % stress it with mask
	sg = sino_geom('par', 'nb', 80, 'na', 40, 'dr', 1, 'strip_width', 'dr');

	%sg.plot(ig), prompt

	A = Gtomo2_strip(sg, ig);

	beta = [2^3 2^6];
	offset = [10 3];
end
printm('converting system matrix A into full matrix');
A = A(:,:);

%R1 = Reg1(ig.mask, 'offsets', [ig.nx-1]); % single offset for Hugo.  bad!
R1 = Reg1(ig.mask);
Cfull = R1.C;
printm('converting roughness penalty matrix C into full matrix');
Cfull = Cfull(:,:);
C = {Cfull Cfull};
%C = Csparse('maskbar', ig.mask); % test with sparse C

[psf0 fwhm var con0] = ...
	qpwls_psf_slow_st(A, C, beta, ig.mask, [], 'offset', offset, 'chat', 1);
printm('approximate stddev = %g', sqrt(var))

if 1 % check matched yb version.
	ej = ig.unitv(ig.nx/2+1+offset(1), ig.ny/2+1+offset(2));
	yb1 = A * ej(ig.mask);
	psf1 = qpwls_psf_slow_st(A, C, beta, ig.mask, [], 'offset', offset, ...
		'fwhmtype', 'none', ...
		'yb', yb1);
	max_percent_diff(psf0, psf1)
prompt
end

if 0
	ell = [offset+0.5 0.5 0.5 0 1];
	yb2 = ellipse_sino(sg, ell, 'oversample', 3);
%	im clf, im(121, yb1), im(122, yb2)
%	max_percent_diff(yb1, yb2)
	psf2 = qpwls_psf_slow_st(A, C, beta, ig.mask, [], 'offset', offset, ...
		'loop', 0, ...
		'chat', 0, 'shift0', 0, 'yb', yb2);
	max_percent_diff(psf0, psf2)
	im clf, im(cat(3, psf0, psf2, psf0-psf2))
prompt
end
