function [im_SOS,p,c] = SOS(im)
% Input: im - original image or its filename
% Output: im_SOS - croped image with a watermark of +1
%         p      - watermark (carrier)
%         c      - detector output at the payload of 1/1600 bit per pixel 

% min{Distortion}  s.t. Robustness >= strength 
% i.e.
% min{w'*A*w}  s.t. p'*w >= strength
% where w - addtive watermark, 
%       p - watermark modulated by the spread spectrum carrier
%       A - kernel of the distortion metric
% Solution w = strength/(p'*inv(A)*p)*inv(A)*p
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LICENSE
% Copyright (C) 2010 ZHANG Fan, with HUST, NTU, & CUHK
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if ischar(im)
    im = imread(im);
end
if ndims(im) >= 3
    im = rgb2gray(im(:,:,1:3));
end
[ROW,COL] = size(im);
row = floor(ROW/40)*40; col = floor(COL/40)*40;
row1 = floor((ROW - row)/2); col1 = floor((COL-col)/2);
im0 = double(im(row1+1:row1+row,col1+1:col1+col));

% load p.mat;
rand('twister', sum(100*clock));
p = double(round(rand(row,col))*2 - 1);

strength = row*col; % Unit watermark strength

% Estimate A

%% SOS metric
blocksize = 5; % Only an ODD block size is supported by this code.
[sI] = BlockInd(row,col,blocksize);

X = im0(sI);
stdX = std(X);
x = mean(X,2);
Xp = X - repmat(x,[1 size(X,2)]);
C = Xp*Xp'/(size(X,2)-1); %Covariance matrix of the cover image

[V,D] = eig(C);
D = diag(D);
D = D / sum(D);

P = p(sI);
VP = V'*P;
D = repmat(D,[1 length(stdX)]);
stdX = repmat((stdX+20),[25,1]);
P = V * (VP ./ D .* stdX);
Ap = zeros(row,col);
Ap(sI) = P;
pAp = sum(sum(VP.^2 ./ D .* stdX));
w_SOS = strength/pAp*Ap;
im_SOS = double(uint8(im0 + w_SOS)); 
c_temp = sum((im_SOS(:) - mean(im_SOS(:))).*p(:))/(row*col);
im_SOS = uint8(im0 + w_SOS*1.0/c_temp); % To control the watermark strength here
figure,subplot(1,2,1),imshow(uint8(im0)),title('Cover image'),
       subplot(1,2,2),imshow(im_SOS),title('Watermarked image');
clear Xp; clear X; clear D; clear P; clear VP; clear Ap; clear stdX; clear C; clear V;

im_attack = attack(im_SOS);

% KLT filtering-based detector
% Detecting the watermarked image
X = single(im_attack(sI));
x = mean(X,2);
Xp = X - repmat(x,[1 size(X,2)]);
C = Xp*Xp'/(size(X,2)-1); %Covariance matrix of the recieved image
[V,D] = eig(C);
D = diag(D);
D = 1./D;
D(1:8) = D(9);
X = V*diag(D)*V'*X; % KLT filtering
P = p(sI);
X = X - mean(X(:));
P = P - mean(P(:));
sI40 = BlockInd(row,col,40); % watermark payload: 40x40 blocking
X = X(sI40);
P = P(sI40);
c1 = sum(X.*P)./sqrt(sum(X.^2).*sum(P.^2)); % normalized LCD
clear X; clear x; clear Xp; clear V; clear D;

% Detecting the cover image (watermark innocent)
X = single(im0(sI));
x = mean(X,2);
Xp = X - repmat(x,[1 size(X,2)]);
C = Xp*Xp'/(size(X,2)-1); %Covariance matrix of the recieved image
[V,D] = eig(C);
D = diag(D);
D = 1./D;
D(1:8) = D(9);
X = V*diag(D)*V'*X; % KLT filtering
P = p(sI);
X = X - mean(X(:));
P = P - mean(P(:));
sI40 = BlockInd(row,col,40); % watermark payload: 40x40 blocking
X = X(sI40);
P = P(sI40);
c0 = sum(X.*P)./sqrt(sum(X.^2).*sum(P.^2)); % normalized LCD

% To show distributions of the detector outputs
bins = min([c1 c0]):range([c1 c0]/10):max([c1 c0]);
hist_x = hist(c0,bins);
hist_z = hist(c1,bins);
figure,plot(bins,hist_x,'b',bins,hist_z,'r');

end

function [sI] = BlockInd(row,col,blocksize)
nsizeh = (blocksize-1)/2;
sCI = 0:row*blocksize:(col/blocksize-1)*row*blocksize;
sRI = 0:blocksize:(row/blocksize-1)*blocksize;
[sCI,sRI] = meshgrid(sCI,sRI);
centerind = sCI + sRI + nsizeh*row + nsizeh + 1;
[sCI,sRI] = meshgrid(-nsizeh*row:row:row*nsizeh,-nsizeh:1:nsizeh);
bo = sRI + sCI;
[sCI,sRI] = meshgrid(centerind(:),bo(:));
sI = uint32(sCI + sRI);
end

function [im_attack] = attack(im)
% To simulate attack here
im_attack = im;
end