function [L, mFR_f, extra] = LNP_BEC_loss(ff, loss, xr, sr, xpmf, c)
% Bayesian Efficient Coding Loss function for Linear-NonLinear-Poisson 
%
% Input
%   ff: (NX x 1) nonlinearity lambda = f(x) discritized over range of x
%   loss: loss function or p > 0 for p-th power loss
%   xr: (NX x 1) range of stimulus values
%   sr: (NS x 1) range of possible spike counts
%   xpmf: (NX x 1) marginal distribution of x (sums to 1)
%   c: (struct/optional) combine the loss function with constraint on
%                  firing rate to produce a single objective function for
%                  easy optimization. c.targetFR is the FR target and c.reg
%                  is the regularizer such that
%                               c.reg * (c.targetFR - mFR)^2
%                  is minimized.
% Output
%   L: (1) BEC loss function E_{X,Y} loss(X - xhat(Y))
%   mFR_f: mean firing rate of this nonlinearity
%   extra: intermediate stuff in case you need them handy

if ~isa(loss, 'function_handle')
    if isempty(loss)
        p = 1; % default
    else
        p = loss;
    end
    assert(p > 0, 'p-th power loss, p > 0');
    loss = @(x) abs(x).^p;
end

if nargin < 5 || isempty(xpmf)
    fprintf('missing x probability mass function, using standard normal\n');
    xpmf = normpdf(xr, 0, 1); xpmf = xpmf / sum(xpmf);
end

%% numerically compute p(y|x), p(x|y), and p(x,y)
% p(y|x,f) ~ Poiss(f(x))
py_x = zeros(numel(sr), numel(xr));
for kr = 1:numel(xr)
    py_x(:, kr) = poisspdf(sr, ff(kr));
end
% p(y) = int p(y|x) * p(x) dx
py = py_x * xpmf';
% p(x|y) = mean(exp(log(p(x)) + log(p(y|x))))
px_y = bsxfun(@times, py_x, xpmf)';
px_y = bsxfun(@times, px_y, 1 ./ sum(px_y));
% p(x,y) = p(x|y) * p(y)
pxy = bsxfun(@times, px_y, py'); % joint distribution

%% mean firing rate
mFR_f = sum(sr .* py'); % constraint acts on this

%% BEC loss function
xz = loss(xr(:) - xr(:)');
[~, idx] = min(xz * px_y);
xhat_y = xr(idx);
L = loss(xr(:) - xhat_y(:)') .* pxy;
L = sum(L(:));

%% combine loss and constrain
if nargin > 5
    L = L + c.reg * (c.targetFR - mFR_f)^2 + c.monReg * sum(diff(ff) < 0);
end

%%
if nargout > 2
    extra.py_x = py_x;
    extra.py = py;
    extra.pxy = pxy;
    extra.px_y = px_y;
end