% function z=reconstruct(x,K,e,ev,emod,SPC)
%
% Replaces the epitome (e,ev) of an image by another epitome image (emod)
% keeping the mapping intact. The routine can be used to edit images or to
% preprocess by coloring the image accoring to the image (see the segmentation
% part of the web page for more info). K specifies the patch size to use in 
% mapping inference and SPC defines the patch spacing.
%
% Copyright  Brendan Frey and Nebojsa Jojic

function z=reconstruct(x,K,e,ev,emod,SPC)


[Nx Ny NoColors]=size(x);

if NoColors==1
    temp=x;
    x=zeros(Nx,Ny,NoColors);
    x(:,:,1)=temp;
    x(:,:,2)=temp;
    x(:,:,3)=temp;
end

[N Ny NoColors]=size(e);
onemat=ones(K,K);
SY=size(x,1); SX=size(x,2);


ewrap=zeros(N+K-1,N+K-1,NoColors); evwrap=zeros(N+K-1,N+K-1,NoColors);
MINP=1e-6; MINV=0.01;

ewrap(1:N,1:N,:)=e; ewrap(N+1:end,:,:)=ewrap(1:K-1,:,:);
ewrap(:,N+1:end,:)=ewrap(:,1:K-1,:);
evwrap(1:N,1:N,:)=ev; evwrap(N+1:end,:,:)=evwrap(1:K-1,:,:);
evwrap(:,N+1:end,:)=evwrap(:,1:K-1,:);
logevwrap=log(evwrap); evwrapi=1./evwrap;
tmp3=zeros(N,N,3); tmp4=zeros(N,N,3);
for i=1:3
    tmp3(:,:,i)=conv2(ewrap(:,:,i).^2.*evwrapi(:,:,i),onemat,'valid');
    tmp4(:,:,i)=conv2(logevwrap(:,:,i),onemat,'valid');
end;
emodwrap=ewrap;
emodwrap(1:N,1:N,:)=emod; emodwrap(N+1:end,:,:)=emodwrap(1:K-1,:,:);
emodwrap(:,N+1:end,:)=emodwrap(:,1:K-1,:);

% For each patch in the input image, find the distribution over the epitome posn,
% and reconstruct the patch in the input image, using the modified epitome.
z=zeros(size(x)); cnt=zeros(SY,SX);
for iy=1:SPC:SY-K+1
    for ix=1:SPC:SX-K+1
        yc=x(iy:iy+K-1,ix:ix+K-1,:);
        yct=yc(end:-1:1,end:-1:1,:);
        lP(:,:)=0; %log(p);
        for i=1:3
            tmp1=conv2(evwrapi(:,:,i),yct(:,:,i).^2,'valid');
            tmp2=conv2(ewrap(:,:,i).*evwrapi(:,:,i),yct(:,:,i),'valid');
            lP=lP-0.5*(tmp4(:,:,i)+tmp1-2*tmp2+tmp3(:,:,i));
        end;
        mxlP=max(lP(:)); P=exp(lP-mxlP)+MINP; nrmP=sum(P(:)); P=P/nrmP;
        [mxp mxpj]=max(P,[],2); [mxp mxpi]=max(mxp); mxpj=mxpj(mxpi);
        newp=emodwrap(mxpi:mxpi+K-1,mxpj:mxpj+K-1,:);
        oldp=ewrap(mxpi:mxpi+K-1,mxpj:mxpj+K-1,:);
        if prod(double(oldp(:)==newp(:)))==0
            z(iy:iy+K-1,ix:ix+K-1,:)=z(iy:iy+K-1,ix:ix+K-1,:)+newp;
        else
            z(iy:iy+K-1,ix:ix+K-1,:)=z(iy:iy+K-1,ix:ix+K-1,:)+ ...
                x(iy:iy+K-1,ix:ix+K-1,:);
        end;
        cnt(iy:iy+K-1,ix:ix+K-1,:)=cnt(iy:iy+K-1,ix:ix+K-1,:)+1;
    end;
    iy
end;
z=z(1:iy+K-1,1:ix+K-1,:); cnt=cnt(1:iy+K-1,1:ix+K-1);
for i=1:3 z(:,:,i)=z(:,:,i)./cnt; end;