%%
H1 = load('results/H1_param.mat'); % obtained from cmd_H1.m

%%
rng(20180629);
c = struct();
c.reg = 1e-8; % firing rate
c.targetFR = 2.233;
c.monReg_enc = 1e-3; % monotonicity of nonlinearity
c.monReg_dec = 1e2; % monotonicity of decoder
c.minFR = 0.1; %0.02;
c.pymass = 10;
%loss = @(x) (x.^2);
%loss = @(x) abs(x);
maxRate = c.targetFR * 3;
xr = linspace(-3, 3, 32)';
sr = (0:ceil(maxRate * 3))';
ff_emp = H1.fr * H1.slen / H1.Fs;
enc00 = interp1(H1.xr, ff_emp, xr, 'spline');

%%
for kInit = 1:10
    
switch(kInit)
    case 1
        enc0 = maxRate / 20 * log(1 + exp(4 * (xr+1.8))); enc0(xr < -2) = enc0(xr < -2) .* (1e-2 * (xr(xr < -2) + 3.1));
        enc0 = (enc0 - min(enc0)) + c.minFR;
    case 2
        enc0 = enc00;
    case 3
        enc0 = 0.3 * enc00;
    case 4
        enc0 = c.targetFR * ones(size(xr));
    case 5
        enc0 = c.minFR * ones(size(xr)) + (xr > 0) .* (2.5 * (c.targetFR - c.minFR) * ones(size(xr)));
    case 6
        enc0 = linspace(c.minFR, c.targetFR*2, numel(xr));
    case 7
        enc0 = linspace(-c.targetFR*2, c.targetFR*2, numel(xr));
        enc0(xr <= 0) = 1e-5;
    case 8
        enc0 = enc00;
        enc0(end) = enc0(end-1) + 1e-3;
    case 9
        enc0 = enc00;
        enc0(xr <= 0) = 1e-5;
        enc0 = sort(enc0 + 0.01 * randn(size(enc00)));
        enc0(enc0 < 1e-5) = 1e-5;
    case 10
        enc0 = enc00;
        enc0 = sort(enc0 + 0.01 * randn(size(enc00)));
        enc0(enc0 < 1e-5) = 1e-5;
end
%enc0 = sort(20 * rand(size(xr)), 'ascending');
dec0 = log(sr/2+0.1);
px = normpdf(xr, 0, 1); px = px / sum(px);

%% optimization parameters
opts_enc = optimoptions('fmincon','Display','iter-detailed','MaxIterations',2,'OptimalityTolerance',1e-5,'MaxFunctionEvaluations',10 * numel(xr));
opts_dec = optimoptions('fmincon','Display','iter-detailed','MaxIterations',4,'OptimalityTolerance',1e-5,'MaxFunctionEvaluations',6 * (numel(sr)+3));
isPlotIter = true;
maxIter = 300;

%%
%pList = 2.^(-3:4);
%pList = logspace(log10(2), log10(8), 16);
pList = [0.5, 1, 2, 3, 4, 5.5, 7, 0];
pLabel = cell(numel(pList), 1);
results = struct();
cmap = cool(numel(pList));

%%
fig0 = figure(8400); clf; hold all;
dxr = linspace(0, 1);
for kp = 1:numel(pList)
    p = pList(kp);
    if p == 0
        pLabel{kp} = 'MI';
    else
        pLabel{kp} = sprintf('%2.3f', p);
        loss = @(x) abs(x).^p;
        plot(dxr, loss(dxr), 'Color', cmap(kp, :));
    end
end
lh = legend(pLabel{:}); set(lh, 'box', 'off', 'Location', 'NorthWest');
title('loss function comparison');
xlabel('absolute reconstruction error');
ylabel('loss');

%%
for kp = 1:numel(pList)
    p = pList(kp);
    if p == 0
        fprintf('max MI (min conditional entropy)\n');
        loss = 'entropy';
    else
        fprintf('[p = %g]\n', p);
        loss = @(x) abs(x).^p;
        
        xz = loss(xr(:) - xr(:)');
    end
    [L0, mFR_f0, extra] = LNP_BEC_loss_joint(enc0, dec0, loss, xr, sr, px);
    
    %%
    enc = enc0;
    dec = dec0;
    
    if isPlotIter
        fig = figure(1941); clf;
        subplot(2,2,1); hold all; title('encoding nonlinearity'); xlabel('stimulus'); ylabel('firing rate'); plot(xr, enc0, 'k'); plot(xr, enc00, 'g', 'LineWidth', 3);
        subplot(2,2,2); hold all; title('decoding'); xlabel('spike count'); ylabel('stimulus'); plot(sr, dec0, 'k');
        subplot(2,2,3); hold all; title('loss'); xlabel('optimization step'); ylabel('log10(loss)');
        subplot(2,2,4); hold all; title('mean firing rate'); line([0, maxIter], c.targetFR * [1,1], 'Color', 'k');
        fig2 = figure(471); imagesc(xr, sr, extra.pxy'); axis xy; colorbar;
        plot(xr, enc, 'w');
        plot(dec, sr, 'r');
    end
    
    %% optimization loop
    encOld = zeros(size(enc));
    
    for k = 1:maxIter
        %%
        if kp ~= 0 && mod(k, 1) == 0
            fprintf('re-initializing dec');
            px_y = bsxfun(@times, extra.pxy, 1 ./ sum(extra.pxy));
            [~, idx] = min(xz * px_y);
            xhat_y = xr(idx);
            dec = xhat_y;
        end
        satConst = true;
        
        if kp ~= 0 % for MI, decoder doesn't change
            [dec, Ldec] = fmincon(@(dec) LNP_BEC_loss_joint(enc, dec, loss, xr, sr, px, c), dec, [], [], [], [], min(xr) * ones(size(dec)), max(xr) * ones(size(dec)), [], opts_dec);
        end
        
        if any(diff(dec) < -1e-4)
            satConst = false;
            c.monReg_dec = c.monReg_dec * 2;
            fprintf('## increasing decoding monotonicity regularizer to [%g]\n', c.monReg_dec);
        else
            if c.monReg_dec > 1e-10
                c.monReg_dec = c.monReg_dec / 1.5;
                fprintf('@@ decreasing decoding monotonicity regularizer to [%g]\n', c.monReg_dec);
            end
        end
        
        encOld = enc;
        [enc, Lenc] = fmincon(@(enc) LNP_BEC_loss_joint(enc, dec, loss, xr, sr, px, c), enc, [], [], [], [], c.minFR * ones(size(enc)), [], [], opts_enc);
        [L, mFR_f, extra] = LNP_BEC_loss_joint(enc, dec, loss, xr, sr, px); % compute marginal statistics
        
        if abs(mFR_f - c.targetFR) > 1e-3 % if constraint is not within tolerence
            satConst = false;
            c.reg = c.reg * 1.1; % increase regularizer if conditions are not met
            fprintf('## increasing FR regularizer to [%g]\n', c.reg);
            
            if c.targetFR - mFR_f > 1e-3 % if target firing rate is higher
                enc = enc * (1 + 0.05 * (c.targetFR - mFR_f));
                fprintf('++++++++\n');
            end
            if c.targetFR - mFR_f < -1e-3
                enc = enc * (1 + 0.05 * (c.targetFR - mFR_f));
                fprintf('--------\n');
            end
        else
            if c.reg > 1e-10
                fprintf('@@ decreasing FR regularizer to [%g]\n', c.reg);
                c.reg = c.reg / 1.5;
            end
        end
        
        if any(diff(enc) < -1e-3)
            satConst = false;
            c.monReg_enc = c.monReg_enc * 2;
            fprintf('## increasing encoding monotonicity regularizer to [%g]\n', c.monReg_enc);
            if mod(k, 1) == 0
                fprintf('################ reordering encoding nonlinearity\n');
                enc = sort(enc, 'ascend');
            end
        else
            if c.monReg_enc > 1e-10
                c.monReg_enc = c.monReg_enc / 2;
                fprintf('@@ decreasing encoding monotonicity regularizer to [%g]\n', c.monReg_enc);
            end
        end
        
        if max(enc) > sr(end) % rate too high for the grid
            fprintf('firing rate is too high for the grid!\n');
            break;
        end
        
        if isPlotIter
            figure(fig);
            subplot(2,2,1); plot(xr, enc);
            subplot(2,2,2); plot(sr, dec);
            subplot(2,2,3); plot(k, log10(Lenc), 'rx'); plot(k, log10(Ldec), 'bo');
            subplot(2,2,4); plot(k, mFR_f, 'ro');
            
            fig2 = figure(471); clf; hold all;
            imagesc(xr, sr, extra.pxy'); axis xy;
            plot(xr, enc, 'w');
            plot(dec, sr, 'r');
            
            title(sprintf('joint distribution [%g]', p)); xlabel('stimulus'); ylabel('spike count');
            axis tight;
            colorbar
            drawnow;
        end
        
        if norm(encOld - enc) < 1e-5 && satConst
            fprintf('>>> minimal change in encoding weights! <<< [%d]\n', k);
            [dec, Ldec] = fmincon(@(dec) LNP_BEC_loss_joint(enc, dec, loss, xr, sr, px, c), dec, [], [], [], [], min(xr) * ones(size(dec)), max(xr) * ones(size(dec)), [], opts_dec);
            [L, mFR_f, extra] = LNP_BEC_loss_joint(enc, dec, loss, xr, sr, px);
            break;
        end
    end
    
    %% save optimization results
    results(kp).L0 = L0;
    results(kp).satConst = satConst;
    results(kp).L = L;
    results(kp).mFR = mFR_f;
    results(kp).enc = enc;
    results(kp).dec = dec;
    results(kp).extra = extra;
    
    savefig([fig, fig2], sprintf('%s_%02.fig', mfilename, p))
end % kp:pList

save(sprintf('results/%s_results_unipy_init_%d', mfilename, kInit), 'results', 'pList', 'pLabel', 'xr', 'ff_emp', 'enc0', 'dec0', 'sr', 'px');

%%
end % kInit

%% generate summary figure
fig = figure(887); clf;
subplot(1, 3, 2); hold on;
fsf = 1; % spike count version
%fsf = 1 / (H1.slen / H1.Fs); % firing rate scaling factor
line([xr(1), xr(end)], mean(H1.y2) * fsf * [1,1], 'Color', 0.8 * [1,1,1]);
plot(H1.xr, fsf * ff_emp, '-', 'LineWidth', 2, 'Color', 0.5 * [1,1,1]);
%plot(H1.xr, exp(H1.fm1+2*H1.fs1), '-', 'LineWidth', 1, 'Color', 0.3 * [1,1,1]);
%plot(H1.xr, exp(H1.fm1-2*H1.fs1), '-', 'LineWidth', 1, 'Color', 0.3 * [1,1,1]);
plot(xr, fsf * enc0, ':', 'LineWidth', 2, 'Color', 0.4 * [0.2,1,0.2]);
title('optimal nonlinearities'); xlabel('stimulus projection'); ylabel('firing rate');
subplot(1, 3, 3); hold on;
y2h = histcounts(H1.y2, [sr; Inf], 'Normalization', 'pdf');
plot(sr, cumsum(y2h), '.-', 'LineWidth', 2, 'Color', 0.5 * [1,1,1]);
title('response distributions'); xlabel('spike count (25 ms bin)'); ylabel('P(response)');

mse = nan(numel(pList), 1);
for kp = 1:numel(pList)
    mse(kp) = px' * (results(kp).enc(:) - enc00).^2;
end
[~, midx] = min(mse);
fprintf('MMSE loss p = [%g]\n', pList(midx));

cmap = parula(numel(pList));

for kp = 1:numel(pList)
    figure(fig); subplot(1, 3, 2); hold on
    if kp == midx
        plot(xr, fsf * results(kp).enc, '-', 'Color', 'r', 'LineWidth', 2);
        figure(fig); subplot(1, 3, 3); hold on
        plot(sr, cumsum(results(kp).extra.py), '-', 'Color', 'r', 'LineWidth', 2);
    else
        plot(xr, fsf * results(kp).enc, '-', 'Color', cmap(kp, :), 'LineWidth', 1);
        figure(fig); subplot(1, 3, 3); hold on
        plot(sr, cumsum(results(kp).extra.py), '-', 'Color', cmap(kp, :));
    end
end
subplot(1, 3, 2);
lh = legend('mean firing rate', 'empirical', 'initialization', pLabel{:}, 'Location', 'NorthWest');
set(lh, 'box', 'off');
set(lh, 'FontSize', 11);
subplot(1, 3, 2); xlim([-3, 3]);
subplot(1, 3, 3); xlim([0, 15]);

% %%
% set(fig, 'PaperUnits', 'inches', 'PaperSize', [7, 3.5], 'PaperPosition', [0, 0, 7, 3.5]);
% saveas(fig, 'LNP_BEC_nl.pdf');
% saveas(fig, 'LNP_BEC_nl.png');
% 
% %%
% savefig(fig, sprintf('cmd_LNP_BEC_summary_%s.fig', datestr(now, 30)));

%% scaled gaussian-CDF with the same mean firing rate constraint
ncdf = 4.4615 * normcdf(xr, 0, 1);
for k = 1:1000
    [L, mFR_f, extra] = LNP_BEC_loss_joint(ncdf, dec0, loss, xr, sr, px);
    
    fprintf('%g\n', mFR_f);
    if abs(c.targetFR - mFR_f) < 1e-3
        break;
    end
    if c.targetFR - mFR_f > 1e-3 % if target firing rate is higher
        ncdf = ncdf * (1 + 0.21 * (c.targetFR - mFR_f));
    end
    if c.targetFR - mFR_f < -1e-3
        ncdf = ncdf * (1 + 0.22 * (c.targetFR - mFR_f));
    end
end

%%
subplot(1, 3, 1); cla; hold on;
title('blowfly H1 data')
ridx = randperm(numel(H1.y2)); ridx = ridx(1:6000);
fprintf('plotting %f%% of the data\n', numel(ridx)/numel(H1.y2)*100);
line([xr(1), xr(end)], mean(H1.y2) * fsf * [1,1], 'Color', 0.8 * [1,1,1]);
plot(H1.z2(ridx), (H1.y2(ridx) + 0.2*abs(randn(numel(ridx), 1))) * fsf, '.');
ylabel('spike count (25 ms bin)');
xlabel('z-scored stimulus projected onto STA');
plot(H1.xr, fsf * ff_emp, '-', 'LineWidth', 2, 'Color', 0.1 * [1,1,1]);
plot(xr, fsf * ncdf, 'g', 'LineWidth', 2); % stimulus CDF
%plot(H1.xr, fsf * exp(H1.fm1+2*H1.fs1), '-', 'LineWidth', 1, 'Color', 0.3 * [1,1,1]);
%plot(H1.xr, fsf * exp(H1.fm1-2*H1.fs1), '-', 'LineWidth', 1, 'Color', 0.3 * [1,1,1]);
%ylim([-40, 600]);
plot(xr, fsf * results(end).enc, '-', 'Color', cmap(end, :), 'LineWidth', 2);
ylim([0, 14]);
set(gca, 'TickDir', 'out');
%grid on
%set(fig, 'PaperUnits', 'inches', 'PaperSize', [3.5, 3.5], 'PaperPosition', [0, 0, 7, 3.5]);
%xlim([-5, 5]);
set(gca, 'XTick', -5:5)
xlim([-3, 3]);
lh = legend('mean firing rate', '1% of data', 'empirical nonlinearity', 'scaled stimulus CDF', 'MI solution', 'Location', 'NorthWest');
set(lh, 'box', 'off');
set(lh, 'FontSize', 11);

%%
subplot(1,3,2); %ylim([-40, 600]);
%plot(xr, fsf * ncdf, 'g', 'LineWidth', 2); % stimulus CDF
ylim([0, 14]);
set(gca, 'YTickLabel', []); ylabel([]);
set(gca, 'TickDir', 'out');
%grid on

%%
set(fig, 'PaperUnits', 'inches', 'PaperSize', [9, 3.5], 'PaperPosition', [0, 0, 9, 3.5]);
saveas(fig, 'LNP_BEC_nl3.pdf');
saveas(fig, 'LNP_BEC_nl3.png');