function jsdnss = jsd_nss(img1, img2, C, shift_idx, plot_nss)
% -----------COPYRIGHT NOTICE STARTS WITH THIS LINE------------
% Copyright (c) 2020 The University of Texas at Austin
% All rights reserved. Permission is hereby granted, without written agreement and without license or royalty fees, to use, copy, modify, and distribute this code (the source files) and its documentation for any purpose, provided that the copyright notice in its entirety appear in all copies of this code, and the original source of this code, Laboratory for Image and Video Engineering (LIVE, http://live.ece.utexas.edu) and Center for Perceptual Systems (CPS, http://www.cps.utexas.edu) at the University of Texas at Austin (UT Austin, http://www.utexas.edu), is acknowledged in any publication that reports research using this code.
% IN NO EVENT SHALL THE UNIVERSITY OF TEXAS AT AUSTIN BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OF THIS DATABASE AND ITS DOCUMENTATION, EVEN IF THE UNIVERSITY OF TEXAS AT AUSTIN HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. THE UNIVERSITY OF TEXAS AT AUSTIN SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE DATABASE PROVIDED HEREUNDER IS ON AN "AS IS" BASIS, AND THE UNIVERSITY OF TEXAS AT AUSTIN HAS NO OBLIGATION TO PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
% -----------COPYRIGHT NOTICE ENDS WITH THIS LINE------------
%
% The research is to be cited in the bibliography as:
% L.-H. Chen, C. G. Bampis, Z. Li, and A. C. Bovik, "Learning to Distort 
% Images Using Generative Adversarial Networks," in IEEE Signal Processing 
% Letters, vol. 24, no. 9, pp. 1333-1337, Sept. 2020.
%
% ========================================================================
%Input :
%             img1: the first input image
%             img2: the second input image
%             C: constants in the denominator of MSCN calculation
%             shift_idx: Paired product index (1: right-shift, 2: bot-shift
%             plot_nss: visualization
%
%Output:     
%             jsdnss: the JSD_{NSS} value between 2 images.
%
%Basic Usage:
%   Given 2 test images img1 and img2,ranging [0,255]
%   jsdnss = jsd_nss(img1, img2);
%
%   A self demo is also included in this function, simply use
%   jsd_nss
%========================================================================

if nargin <  1, selfdemo; return; end
if (nargin < 2 || nargin > 5), jsdnss = -Inf; return; end
if nargin == 2, C = 1; shift_idx = 3; plot_nss = false; end
if nargin == 3,        shift_idx = 3; plot_nss = false; end
if nargin == 4,                       plot_nss = false; end

pair1 = compute_pair_prod(img1, shift_idx, C);
pair2 = compute_pair_prod(img2, shift_idx, C);

bins = 150; % number of bins forming histograms
[N1,edges] = histcounts(pair1, bins, 'Normalization', 'probability');
x1 = linspace(edges(1), edges(end), length(N1));
[N2,edges] = histcounts(pair2, bins, 'Normalization', 'probability');
x2 = linspace(edges(1), edges(end), length(N2));
xbins = linspace(-max(abs([x1,x2])), max(abs([x1,x2])), bins);
jsdnss = compute_jsd(pair1(:), pair2(:), xbins);

if plot_nss == true
    % Cleaning 0's for better visualization of log scale
    N1(N1==0) = min(N1(N1~=0));
    N2(N2==0) = min(N2(N2~=0));
    % Begin plot
    plot(x1, N1, '-' , 'LineWidth', 1.2, 'Color', '#0F4C81'); hold on;
    plot(x2, N2, '-.', 'LineWidth', 1.9, 'Color', '#77202F');
    
    xlabel('Paired product MSCN coefficient', 'interpreter', 'latex')
    ylabel('Probability', 'interpreter', 'latex')
    xlim([-max(abs([x1,x2]))+0.02, max(abs([x1,x2]))-0.02])
    ylim([0, 1])
    leg = legend('Img1', 'Img2');
    set(leg, 'FontSize', 16, 'interpreter', 'latex')
    grid on

    aa = gca;
    aa.YScale = 'log';
    aa.FontSize = 14;
    aa.XRuler.Axle.LineWidth = 1.4;
    aa.YRuler.Axle.LineWidth = 1.4;
end

function pair = compute_pair_prod(img, itr_shift, c)
    img_double = im2double(img);
    shifts = [0 1; 1 0; 1 1; -1 1];
    window = fspecial('gaussian',7,7/6);
    window = window/sum(sum(window));
    pair = [];
    for i = 1:3
        mu                = filter2(window, img_double(:,:,i), 'same');
        mu_sq             = mu.*mu;
        sigma             = sqrt(abs(filter2(window, img_double(:,:,i).*img_double(:,:,i), 'same') - mu_sq));
        structdis         = (img_double(:,:,i)-mu)./(sigma+c);
        shifted_structdis = circshift(structdis,shifts(itr_shift,:));
        pair              = [pair structdis(:).*shifted_structdis(:)];
    end
end

function jsd = compute_jsd(coeff1, coeff2, binsize)
    [h, ~] = hist(coeff1, binsize);
    [y, ~] = hist(coeff2, binsize);
    h = h./sum(h);
    y = y./sum(y);
    M = 0.5*(h+y);
    M = M./sum(M);
    jsd = 0.5*sum(y.*(log(eps+y)-log(eps+M)))+ 0.5*sum(h.*(log(eps+h)-log(eps+M)));
end

function selfdemo
    % MATLAB builtin demo images
    I      = imread('lighthouse.png');
    Inoise = imnoise(I,'salt & pepper',0.02);
    jsdnss_value = jsd_nss(I, Inoise, 1, 3, true);
    display(["Resulting JSD NSS value:", jsdnss_value]);
end

end