% =========================================================================
% Nonlocal Autoregressive Modeling based image interpolation, Version 1.0
% Copyright(c) 2013 Weisheng Dong, Lei Zhang, Guangming Shi
% All Rights Reserved.
%
% ----------------------------------------------------------------------
% Permission to use, copy, or modify this software and its documentation
% for educational and research purposes only and without fee is here
% granted, provided that this copyright notice and the original authors'
% names appear on all copies and supporting documentation. This program
% shall not be used, rewritten, or adapted as the basis of a commercial
% software or hardware product without first obtaining permission of the
% authors. The authors make no representations about the suitability of
% this software for any purpose. It is provided "as is" without express
% or implied warranty.
%----------------------------------------------------------------------
%
% This is an implementation of the algorithm for image interpolation
% 
% Please cite the following paper if you use this code:
%
% Weisheng Dong, Lei Zhang, R. Lukac, and G. Shi,"Sparse representation based
% image interpolation with nonlocal autoregressive modeling", IEEE Trans. on
% Image Processing, vol. 22, no. 4, pp. 1382-1394, Apr. 2013.
% 
%--------------------------------------------------------------------------
function [im_out PSNR SSIM bic_out bic_PSNR bic_SSIM]   =  NARM_Superresolution( par )
s           =   par.scale;
lr_im       =   par.LR;
[lh lw ch]  =   size(lr_im);
hh          =   lh*s-s+1;
hw          =   lw*s-s+1;
[X,Y]       =   meshgrid(1:1/s:lw, 1:1/s:lh);
[CX,CY]     =   meshgrid(1:lw, 1:lh);
opts1.win       =   5;
opts1.cls_num   =   60;
opts1.step      =   2;
opts1.m_blks    =   160;
opts1.nblk      =   23;
opts1.hp        =   120;
par.s1          =   opts1;

hrim            =   uint8(zeros(hh, hw, ch));
ori_im          =   zeros(hh,hw);

if  ch == 3
    lrim           =  rgb2ycbcr( uint8(lr_im) );
    lrim           =  double( lrim(:,:,1));
    b_im(:,:,1)    =   interp2(CX,CY,lr_im(:,:,1),X,Y, 'spline');
    b_im(:,:,2)    =   interp2(CX,CY,lr_im(:,:,2),X,Y, 'spline');
    b_im(:,:,3)    =   interp2(CX,CY,lr_im(:,:,3),X,Y, 'spline');
    b_im2          =   rgb2ycbcr( uint8(b_im) );
    hrim(:,:,2)    =   b_im2(:,:,2);
    hrim(:,:,3)    =   b_im2(:,:,3);
    
    if isfield(par, 'I')
        ori_im         =   rgb2ycbcr( uint8(par.I) );
        ori_im         =   double( ori_im(:,:,1));
    end
else
    lrim           =   lr_im;
    
    if isfield(par, 'I')
        ori_im             =   par.I;
    end
end
par.lrim    =   lrim;    
hr_im0      =   interp2(CX,CY,lrim,X,Y, 'spline');
hr_im       =   Superresolution( lrim, par, ori_im,  hr_im0, 0, 2, 0.03);
hr_im       =   Superresolution( lrim, par, ori_im,  hr_im, 1, 4, par.beta);


bic_im       =   interp2(CX,CY,lrim,X,Y, 'cubic');
bicim        =   hrim;
if ch==3
    hrim(:,:,1)  =  uint8(hr_im);
    im_out       =  double(ycbcr2rgb( hrim ));
    
    bicim(:,:,1)  =  uint8(bic_im);
    bic_out       =  double(ycbcr2rgb( bicim ));        
else
    im_out    =  hr_im;
    bic_out   =  bic_im;
end
if (par.extrapol==1)
    im_out = padarray(im_out,[s-1 s-1],'replicate','post');
end

if isfield(par,'I')
   PSNR      =  csnr( hr_im(1:hh,1:hw), ori_im(1:hh,1:hw), 0, 0 );
   SSIM      =  cal_ssim( hr_im(1:hh,1:hw), ori_im(1:hh,1:hw), 0, 0 );
   
   [h1 w1]     =  size(bic_im);
   bic_PSNR    =  csnr( bic_im, ori_im(1:h1,1:w1), 0, 0 );
   bic_SSIM    =  cal_ssim( bic_im, ori_im(1:h1,1:w1), 0, 0 );   
end
return;


function hr_im  = Superresolution( lr_im, par, ori_im, hr_im0, flag, L, beta )
s           =   par.scale;
[lh lw]     =   size(lr_im);
[X,Y]       =   meshgrid(1:1/s:lw, 1:1/s:lh);
[CX,CY]     =   meshgrid(1:lw, 1:lh);
hr_im       =   interp2(CX,CY,lr_im,X,Y, 'spline');

[h1 w1]     =   size(hr_im);

y           =   lr_im;
B           =   par.B;
BTB         =   B'*B;
PSNR        =   csnr( hr_im(1:h1,1:w1), ori_im(1:h1,1:w1), 5, 5 );

fprintf( '\nInitial PSNR : PSNR = %f\n', PSNR );
cnt       =   0;
Tau1      =   zeros(0);
Tau2      =   zeros(0);
im0       =   hr_im0;
if flag==0
    r       =  0.78;
else
    r       =  0.72;
end
Z       =   zeros( size(y(:)) );
mu      =   0.7;

tau1        =   par.tau1;
tau2        =   par.tau2;


for k   =  1:L
    
    H         =   Cal_Matrixes_LS( im0, 5, par );
    HTY       =   H'*y(:);

    Dict                 =    Clustering_PCA_New( im0, par.s1 );    
    [blk_arr wei_arr]    =    find_blks_fast( im0, par.s1);    
    if  flag==1
        [Tau1, Tau2]     =    Cal_Parameters( im0, Dict, par, blk_arr, wei_arr );
    end        
        
    par.tau1    =    tau1/beta;
    par.tau2    =    tau2/beta;    
    S           =    @(x) PCA_IS( x, Dict, par, Tau1./beta, Tau2./beta, blk_arr, wei_arr, flag );
    f           =    hr_im;
        
    for  iter = 1 : par.iters
        cnt      =   cnt  +  1;
   
        if (mod(cnt, 10) == 0)
            if isfield(par,'I')
                PSNR     =  csnr( f(1:h1,1:w1), ori_im(1:h1,1:w1), 5, 5 );
                fprintf( 'NARM, Iter %d : PSNR = %f \n', cnt, PSNR );
            end
        end    
        [f, T1, P1]    =   S( f );
        b      =    beta*T1(:) + r*HTY + mu*B'*y(:) + 0.5*B'*Z;
        f1     =    reshape( cgsolve2( @(x) Afun(x, H, r, beta, P1, size(f), mu, BTB), b), h1, w1);
        Z      =    Z + mu*( y(:)-B*f1(:) );
        mu     =    mu*1.2;
        f      =    f1;       
    end    
    hr_im   =  f;
    im0     =  f;
    beta    =  beta*1.25;
end


function  [Tau1, Tau2]   =  Cal_Parameters( im, Dict, par, blk_arr, wei_arr )
PCA_idx    =   Dict.cls_idx;
PCA_D      =   Dict.PCA_D;
s_idx      =   Dict.s_idx;
seg        =   Dict.seg;
A          =   Dict.D0;

b        =  par.s1.win;
s        =  par.s1.step;
b2       =  b*b;
[h  w]   =  size(im);

N       =  h-b+1;
M       =  w-b+1;
r       =  [1:s:N];
r       =  [r r(end)+1:N];
c       =  [1:s:M];
c       =  [c c(end)+1:M];
X0      =  zeros(b*b, N*M);
X_m     =  zeros(b*b,length(r)*length(c),'single');

N       =  length(r);
M       =  length(c);
L       =  N*M;

k    =  0;
for i  = 1:b
    for j  = 1:b
        k        =  k+1;        
        blk      =  im(i:end-b+i,j:end-b+j);
        X0(k,:)  =  blk(:)';
    end
end
idx        =   s_idx(seg(1)+1:seg(2));
set        =   1:size(X_m,2);
set(idx)   =   [];

X_m        =   zeros(length(r)*length(c),b*b,'single');
X = X0';
for i = 1:par.nblk
   v            =  wei_arr(set,i); 
   X_m(set,:)   =  X_m(set,:) + X(blk_arr(set,i),:) .*v(:, ones(1,b2));
end
X_m=X_m';

C0          =   zeros(b2, L, 'single' );
Cu0         =   zeros(b2, L, 'single' );
a0          =   zeros(b2, L, 'single');
b0          =   zeros(b2, L, 'single');
idx         =   s_idx(seg(1)+1:seg(2));    
for  k  =  1 : length(idx)
    i           =   idx(k);
    coe         =   A*X0(:, blk_arr(i, :));    
    C0(:,i)     =   mean(coe.^2, 2);
    a0(:,i)     =   A*X0(:,i);
    coe         =   A*(X0(:, blk_arr(i, :)) - repmat(X_m(:, i), 1, par.s1.nblk) );
    Cu0(:,i)    =   mean(coe.^2, 2);    
    b0(:,i)     =   A*(X0(:,i)-X_m(:,i));
end

set        =   1:L;
set(idx)   =   [];
L          =   length(set);
for  k  =  1 : L
    i         =   set(k);    
    cls       =   PCA_idx(i);
    P         =   reshape(PCA_D(:, cls), b2, b2);    
    coe       =   P*X0(:, blk_arr(i,:));
    C0(:,i)   =   mean(coe.^2, 2);
    a0(:,i)   =   P*X0(:,i);    
    coe       =   P*(X0(:, blk_arr(i,:)) - repmat(X_m(:, i), 1, par.s1.nblk) );
    Cu0(:,i)  =   mean(coe.^2, 2);   
    b0(:,i)   =   P*(X0(:,i)-X_m(:,i));
end

C0       =   max(0, C0-par.nSig^2);
Cu0      =   max(0, Cu0-par.nSig^2);
a0       =   (C0.*abs(a0)).^(1/3);
b0       =   (Cu0.*(b0.^2)).^(1/4);
% Tau1     =   par.c1./(sqrt(C0) + par.eps1);
% Tau2     =   par.c2./(Cu0 + par.eps2);

Tau1     =   par.c1./(abs(a0) + par.eps1);
Tau2     =   par.c2./(b0.^2 + par.eps2);
return;
