function [imgColorRestore1stRGB, imgColorRestore2ndRGB, imgColorRestore1stLAB, imgColorRestore1stRGBMult] = ...
    restoreCD(imgColor, frame, desc, imgColorDB, frameDB, descDB)
% [imgColorRestore1stRGB, imgColorRestore2ndRGB, imgColorRestore1stLAB, ...
%   imgColorRestore1stRGBMult] = ...
%   restoreCD(imgColor, frame, desc, imgColorDB, frameDB, descDB)
% Restores a single CD using different models.

% -------------------------------------------------------------------------
% Initialize
% -------------------------------------------------------------------------

imgColor = double(imgColor);
imgColorDB = double(imgColorDB);

% Match query and database features
minPostRANSAC = 6;
maxTries = 5;
numTry = 1;
while numTry <= maxTries
    [flag, preRANSACMatches, postRANSACMatches, model, matches, definiteInliers] = ...
        SURFPairwiseMatchFeatureSets(frame, desc, [], frameDB, descDB, [], 'euclidean', 0);    
    if postRANSACMatches >= minPostRANSAC
        break;
    end % if
    numTry = numTry + 1;
end % while
if ~flag || (postRANSACMatches < minPostRANSAC)
    disp('Not enough point matches to model distortion.');
    imgColorRestore1stRGB = [];
    imgColorRestore2ndRGB = [];
    imgColorRestore1stLAB = [];
    imgColorRestore1stRGBMult = [];
    return;
end
frame = frame(definiteInliers(1,:), :);
frameDB = frameDB(definiteInliers(2,:), :);

% Scale query and database frame locations
[height, width, channels] = size(imgColor);
[heightDB, widthDB, channels] = size(imgColorDB);
heightSmall = 480;
widthSmall = 640;
heightStretch = height / heightSmall;
widthStretch = width / widthSmall;
frame(:,1) = round( frame(:,1) * widthStretch );
frame(:,2) = round( frame(:,2) * heightStretch );
heightStretch = heightDB / heightSmall;
widthStretch = widthDB / widthSmall;
frameDB(:,1) = round( frameDB(:,1) * widthStretch );
frameDB(:,2) = round( frameDB(:,2) * heightStretch );

% -------------------------------------------------------------------------
% Interpolate between distant feature points
% -------------------------------------------------------------------------

interpolateDistant = 1;
desiredPostRANSAC = 40;
desiredNumAdded = min( 30, floor(size(frame,1)/3) );

if interpolateDistant
    
    % Generate pairs of matches
    frameNew = frame;
    frameDBNew = frameDB;
    availableMatches = [];
    for nMatch1 = 1:size(frame,1)
        for nMatch2 = nMatch1+1:size(frame,1)
            availableMatches = [availableMatches; nMatch1 nMatch2];
        end % nMatch2
    end % nMatch1
    
    numAdded = 0;
    while ( (size(frameNew,1) < desiredPostRANSAC) || (numAdded < desiredNumAdded) ) && (size(availableMatches,1) > 0)
    
        % Find most distant feature pair
        maxDist = 0;
        for nPair = 1:size(availableMatches,1)
            dist = norm(frame(availableMatches(nPair,1),1:2) - frame(availableMatches(nPair,2),1:2));
            if dist > maxDist
               maxDist = dist;
               bestPair = availableMatches(nPair,1:2);
               bestPairInd = nPair;
            end
        end % nPair

        % Sample midpoint of pair
        availableMatches(bestPairInd,:) = [];
        % Query frame        
        coord1 = frame(bestPair(1),1:2);
        coord2 = frame(bestPair(2),1:2);
        midPoint = round( 0.5*(coord1 + coord2) );
        frameNew = [frameNew; midPoint 0 0 0 0];
        % Database frame
        coord1 = frameDB(bestPair(1),1:2);
        coord2 = frameDB(bestPair(2),1:2);
        midPoint = round( 0.5*(coord1 + coord2) );
        frameDBNew = [frameDBNew; midPoint 0 0 0 0];
        numAdded = numAdded + 1;
        
    end % while
    clear frame frameDB;
    frame = frameNew;
    frameDB = frameDBNew;
    
end % if
disp(sprintf('Number of points used for modeling: %d', size(frame,1)));

% showSURF(uint8(imgColorDB), frameDB); title('Database with SURF Points');
% pause(2);

% -------------------------------------------------------------------------
% Use 1st order independent RGB model
% -------------------------------------------------------------------------

addConstant = 1;

% Form matrix X
X = frame(:, 1:2); % (x, y)
if addConstant
    X = [X ones(size(X,1), 1)];
end
XDB = frameDB(:, 1:2);

% Form error vector delta-C
dC = zeros(size(X,1), channels);
for nMatch = 1:size(dC,1)
    for nChannel = 1:channels
        dC(nMatch, nChannel) = imgColorDB(XDB(nMatch,2),XDB(nMatch,1),nChannel) - imgColor(X(nMatch,2),X(nMatch,1),nChannel);
    end % nChannel
end % nMatch

% Calculate least-squares solution
if addConstant
    a = zeros(3, channels);
else
    a = zeros(2, channels);
end
for nChannel = 1:channels
   a(:, nChannel) = pinv(X) * dC(:, nChannel);
end % nChannel
disp('1st order RGB model coefficients:');
disp(a);

% Apply correction to image
imgColorRestore1stRGB = imgColor;
[xMat, yMat] = meshgrid(1:width, 1:height);
for nChannel = 1:channels
    imgColorRestore1stRGB(:,:,nChannel) = imgColorRestore1stRGB(:,:,nChannel) + a(1,nChannel)*xMat + a(2,nChannel)*yMat;
    if addConstant
        imgColorRestore1stRGB(:,:,nChannel) = imgColorRestore1stRGB(:,:,nChannel) + a(3,nChannel);
    end
end % nChannel

% -------------------------------------------------------------------------
% Use 2nd order independent regularized RGB model
% -------------------------------------------------------------------------

addConstant = 1;

% Form matrix X
X = frame(:, 1:2);
X = [X X.^2 X(:,1).*X(:,2)]; % (x, y, x^2, y^2, xy)
if addConstant
    X = [X ones(size(X,1), 1)];
end
lambda = [0 0 10000 10000 10000];
if addConstant
   lambda = [lambda 0]; 
end
X(size(frame,1)+1,:) = lambda;
XDB = frameDB(:, 1:2);

% Form error vector delta-C
dC = [dC; zeros(1, channels)];

% Calculate least-squares solution
if addConstant
    a = zeros(6, channels);
else
    a = zeros(5, channels);
end
for nChannel = 1:channels
   a(:, nChannel) = pinv(X) * dC(:, nChannel); 
end % nChannel
disp('2nd order RGB model coefficients:');
disp(a);

% Apply correction to image
imgColorRestore2ndRGB = imgColor;
for nChannel = 1:channels
    imgColorRestore2ndRGB(:,:,nChannel) = imgColorRestore2ndRGB(:,:,nChannel) + a(1,nChannel)*xMat + a(2,nChannel)*yMat;
    imgColorRestore2ndRGB(:,:,nChannel) = imgColorRestore2ndRGB(:,:,nChannel) + a(3,nChannel)*xMat.^2 + a(4,nChannel)*yMat.^2;
    imgColorRestore2ndRGB(:,:,nChannel) = imgColorRestore2ndRGB(:,:,nChannel) + a(5,nChannel)*xMat.*yMat;
    if addConstant
        imgColorRestore2ndRGB(:,:,nChannel) = imgColorRestore2ndRGB(:,:,nChannel) + a(6,nChannel);
    end
end % nChannel

% -------------------------------------------------------------------------
% Use 1st order independent LAB model
% -------------------------------------------------------------------------

addConstant = 1;

% Form matrix X
X = frame(:, 1:2); % (x, y)
if addConstant
   X = [X ones(size(X,1),1)]; 
end
XDB = frameDB(:, 1:2);

% Form error vector delta-C
dC = zeros(size(X,1), channels);
color1 = zeros(1,channels);
color2 = zeros(1,channels);
whiteRGB = [255 255 255];
for nMatch = 1:size(dC,1)
    % Convert from RGB to LAB
    color1(1:channels) = imgColorDB(XDB(nMatch,2),XDB(nMatch,1),1:channels);
    color2(1:channels) = imgColor(X(nMatch,2),X(nMatch,1),1:channels);
    [color1(1:channels), whiteXYZ] = rgb2xyz(color1, whiteRGB);
    [color2(1:channels), whiteXYZ] = rgb2xyz(color2, whiteRGB);
    color1(1:channels) = vcXYZ2lab(color1, whiteXYZ);
    color2(1:channels) = vcXYZ2lab(color2, whiteXYZ);
    for nChannel = 1:channels
        dC(nMatch, nChannel) = color1(nChannel) - color2(nChannel);
    end % nChannel
end % nMatch

% Calculate least-squares solution
if addConstant
    a = zeros(3, channels);
else
    a = zeros(2, channels);
end
for nChannel = 1:channels
   a(:, nChannel) = pinv(X) * dC(:, nChannel);
end % nChannel
disp('1st order LAB model coefficients:');
disp(a);

% Convert from RGB to LAB
imgColorRestore1stLAB = imgColor;
xMatFlat = reshape(xMat, width*height, 1);
yMatFlat = reshape(yMat, width*height, 1);
imgColorRestore1stLABFlat = reshape(imgColorRestore1stLAB, width*height, channels);
imgColorRestore1stLABFlat2 = rgb2xyz(imgColorRestore1stLABFlat, whiteRGB);
clear imgColorRestore1stLABFlat;
imgColorRestore1stLABFlat = vcXYZ2lab(imgColorRestore1stLABFlat2, whiteXYZ);
clear imgColorRestore1stLABFlat2;

% Apply correction
for nChannel = 1:channels
    imgColorRestore1stLABFlat(:,nChannel) = imgColorRestore1stLABFlat(:,nChannel) + a(1,nChannel)*xMatFlat + a(2,nChannel)*yMatFlat;
    if addConstant
        imgColorRestore1stLABFlat(:,nChannel) = imgColorRestore1stLABFlat(:,nChannel) + a(3,nChannel);
    end
end % nChannel

% Convert from LAB to RGB
imgColorRestore1stLABFlat2 = lab2xyz(imgColorRestore1stLABFlat, whiteXYZ);
clear imgColorRestore1stLABFlat;
imgColorRestore1stLABFlat = xyz2rgb(imgColorRestore1stLABFlat2);
clear imgColorRestore1stLABFlat2;
imgColorRestore1stLAB = reshape(imgColorRestore1stLABFlat, height, width, channels);
clear imgColorRestore1stLABFlat;

% -------------------------------------------------------------------------
% Use 1st order independent RGB multiplicative model
% -------------------------------------------------------------------------

addConstant = 1;

% Form matrix X
X = frame(:, 1:2); % (x, y)
if addConstant
    X = [X ones(size(X,1), 1)];
end
XDB = frameDB(:, 1:2);

% Form error vector delta-C
dC = zeros(size(X,1), channels);
for nMatch = 1:size(dC,1)
    for nChannel = 1:channels
        dC(nMatch, nChannel) = log(imgColor(X(nMatch,2),X(nMatch,1),nChannel) + eps) - log(imgColorDB(XDB(nMatch,2),XDB(nMatch,1),nChannel) + eps);
    end % nChannel
end % nMatch

% Calculate least-squares solution
if addConstant
    a = zeros(3, channels);
else
    a = zeros(2, channels);
end
for nChannel = 1:channels
   a(:, nChannel) = pinv(X) * dC(:, nChannel);
end % nChannel
disp('1st order RGB multiplicative model coefficients:');
disp(a);

% Apply correction to image
imgColorRestore1stRGBMult = zeros(height, width, channels);
[xMat, yMat] = meshgrid(1:width, 1:height);
for nChannel = 1:channels
    imgLog = a(1,nChannel)*xMat + a(2,nChannel)*yMat;
    if addConstant
        imgLog = imgLog + a(3,nChannel);
    end
    imgMult = exp(imgLog);
    imgColorRestore1stRGBMult(:,:,nChannel) = imgColor(:,:,nChannel) ./ (imgMult + eps);
    clear imgMult imgLog;
end % nChannel

% -------------------------------------------------------------------------
% Post scaling to match grayscale means
% -------------------------------------------------------------------------

postScale = 1;

if postScale
    
    % Convert color to grayscale
    warning off;
    imgGrayDB = double( rgb2gray(uint8(imgColorDB)) );
    imgGray1stRGB = double( rgb2gray(uint8(imgColorRestore1stRGB)) );
    imgGray2ndRGB = double( rgb2gray(uint8(imgColorRestore2ndRGB)) );
    imgGray1stLAB = double( rgb2gray(uint8(imgColorRestore1stLAB)) );
    imgGray1stRGBMult = double( rgb2gray(uint8(imgColorRestore1stRGBMult)) );
    warning on;

    % Extract convex region of feature points
    featureMapDB = zeros(heightDB, widthDB);
    for nFeature = 1:size(frameDB,1)
        featureMapDB(frameDB(nFeature,2), frameDB(nFeature,1)) = 1;
    end % nFeature
    shapeProps = regionprops(featureMapDB, 'BoundingBox', 'ConvexImage');
    ci = shapeProps(1).ConvexImage;
    bb = round( shapeProps(1).BoundingBox );
    featureMapDB(bb(2):bb(2)+bb(4)-1, bb(1):bb(1)+bb(3)-1) = ci;
    indMapDB = find(featureMapDB > 0);
    featureMapQ = zeros(height, width);
    for nFeature = 1:size(frame,1)
        featureMapQ(frame(nFeature,2), frame(nFeature,1)) = 1;
    end % nFeature
    shapeProps = regionprops(featureMapQ, 'BoundingBox', 'ConvexImage');
    ci = shapeProps(1).ConvexImage;
    bb = round( shapeProps(1).BoundingBox );
    featureMapQ(bb(2):bb(2)+bb(4)-1, bb(1):bb(1)+bb(3)-1) = ci;
    indMapQ = find(featureMapQ > 0);
    
    % Sample means from convex region
    meanDB = mean( imgGrayDB(indMapDB) );
    mean1stRGB = mean( imgGray1stRGB(indMapQ) );
    mean2ndRGB = mean( imgGray2ndRGB(indMapQ) );
    mean1stLAB = mean( imgGray1stLAB(indMapQ) );
    mean1stRGBMult = mean (imgGray1stRGBMult(indMapQ) );
    scale1stRGB = meanDB / mean1stRGB;
    scale2ndRGB = meanDB / mean2ndRGB;
    scale1stLAB = meanDB / mean1stLAB;
    scale1stRGBMult = meanDB / mean1stRGBMult;
    disp('Post scale factors:'); 
    disp([scale1stRGB scale2ndRGB scale1stLAB scale1stRGBMult]);
    imgColorRestore1stRGB = imgColorRestore1stRGB * scale1stRGB;
    imgColorRestore2ndRGB = imgColorRestore2ndRGB * scale2ndRGB;
    imgColorRestore1stLAB = imgColorRestore1stLAB * scale1stLAB;
    imgColorRestore1stRGBMult = imgColorRestore1stRGBMult * scale1stRGBMult;

end % if

% -------------------------------------------------------------------------
% Convert output images to uint8 precision
% -------------------------------------------------------------------------

warning off; imgColorRestore1stRGB = uint8( imgColorRestore1stRGB ); warning on;

warning off; imgColorRestore2ndRGB = uint8( imgColorRestore2ndRGB ); warning on;

warning off; imgColorRestore1stLAB = uint8( imgColorRestore1stLAB ); warning on;

warning off; imgColorRestore1stRGBMult = uint8( imgColorRestore1stRGBMult ); warning on;
