clear
clc
close all

%% make the sky blue
pic = imread('Vienna.jpg');
witney = imread('Witney.jpg');
image(witney)  % First figure is original Witney sky
figure

% make a perfect sky by curve fitting the shapes
% of each red and green row (blues will all be 255)
[rows, cols, clrs] = size(pic);
known = 600;  % the rows where we have data
needed = 1200;  % the rowes we want to extrapolate
ps = uint8(ones(600, cols, 3) * 255); % space for the perfect sky
gc = zeros(needed, 3);  % storage for coefficients of x^2, x and const
rc = zeros(needed, 3);
x = 1:cols;
for row = 1:600
    rd = double(witney(row,:,1));
    gn = double(witney(row,:,2));
    % extract original sky values
    % compute and save coefficients green row
    gcoeff = polyfit(x, gn, 2);
    gc(row,:) = gcoeff;
    % write a new green row
    g1 = polyval(gcoeff, x);
    % put it into the perfect sky
    ps(row, x, 2) = uint8(g1);
    % repeat for the red
    rcoeff = polyfit(x, rd, 2);
    rc(row,:) = rcoeff;
    r1 = polyval(rcoeff, x);
    ps(row, x, 1) = uint8(r1);
end
image(ps);  % second figure is the perfect sky
figure
xk = 1:known;
xn = 1:needed;
% construct storage for the coefficients of x^2, x and c
% for each color
% these are aligned by row, so we transpose to think about
% fitting curves to them 
g2 = zeros(1,needed);
g1 = zeros(1,needed);
g0 = zeros(1,needed);
r2 = zeros(1,needed);
r1 = zeros(1,needed);
r0 = zeros(1,needed);
% We have to scale the values so they all appear on one plot
g2(xk) = gc(xk,1)' * 10^4;
g1(xk) = gc(xk,2)';
g0(xk) = gc(xk,3)'/1000;
r2(xk) = rc(xk,1)' * 10^4;
r1(xk) = rc(xk,2)';
r0(xk) = rc(xk,3)'/1000;
plot(xn, g2, xn, g1, xn, g0, xn, r2, xn, r1, xn, r0);
lgnd = {'green coeffs of x^2 * 10^4', 'green coeffs of x', 'green const/1000' ...
    'red coeffs of x^2 * 10^4', 'red coeffs of x', 'red const/1000' };
legend (lgnd,'Location','SouthWest');
% now, we do a 4th order polynomial fit for each coefficient
% and plot each out to row 1200, not just 600
order = 4;
hold on
fg2 = polyval(polyfit(xk,g2(xk),order), xn);
plot(xn, fg2, 'b');
fg1 = polyval(polyfit(xk,g1(xk),order), xn);
plot(xn, fg1, 'g');
fg0 = polyval(polyfit(xk,g0(xk),order), xn);
plot(xn, fg0, 'r');
fr2 = polyval(polyfit(xk,r2(xk),order), xn);
plot(xn, fr2, 'c');
fr1 = polyval(polyfit(xk,r1(xk),order), xn);
plot(xn, fr1, 'm');
fr0 = polyval(polyfit(xk,r0(xk),order), xn);
plot(xn, fr0, 'y');
% plots of the extrapolated coefficients
    
figure
% Now, we construct and show the new, extrapolated sky
for row = known+1:needed
    ngc(1) = fg2(row)/10^4;
    ngc(2) = fg1(row);
    ngc(3) = fg0(row) * 1000;
    nrc(1) = fr2(row)/10^4;
    nrc(2) = fr1(row);
    nrc(3) = fr0(row) * 1000;
    g1 = polyval(ngc, x);
    r1 = polyval(nrc, x);
    ps(row, x, 1) = uint8(r1);
    ps(row, x, 2) = uint8(g1);
    ps(row, x, 3) = 255;
end    
image(ps)  % extrapolated sky
    
figure
% Fortunately, there is just enough well behaved sky to
% do the sky replacement
pp = pic;
blue_lim = 200;
green_lim = 200;
red_lim = 200;
for r = 1:800
    for c = 1:cols
        if (pp(r,c,1) > red_lim ) ...
                & (pp(r,c,2) > green_lim)...
                & (pp(r,c,3) > blue_lim)
            pp(r,c,:) = ps(r,c,:) ;
        end
    end
end
image(pp)
imwrite(pp,'blueVienna.jpg')