function [origIm resultIm Kr Kb] = colorDifference(filename, edgeType)
% function [origIm resultIm Kr Kb] = colorDifference(filename, edgeType)
% 
% Computes the color difference interpolated image. Takes in a filename
% for an image, applies the Bayer color filter array to it, then demosaicks
% the image.
%
% origIm = the original image
% resultIm = the demosaicked image
% Kr = the color difference channel Kr (Green minus Red)
% Kb = the color difference channel Kb (Green minus Blue)
% filename = input image filename
% edgeType = 0 for vertical, 1 for horizontal, or 2 for adaptive

t1=cputime;
VERTICAL = 0;
HORIZONTAL = 1;
ADAPTIVE = 2;


%assuming image size is even in both dimensions (divisible by 2)


im = cast(imread(filename), 'double');
imCFA = bayerCFA(filename);

imSize = size(imCFA);
im = im(1:floor(imSize(1)/2)*2, 1:floor(imSize(2)/2)*2,:);
imCFA = imCFA(1:floor(imSize(1)/2)*2, 1:floor(imSize(2)/2)*2,:);

imSize = size(imCFA);

R = imCFA(:,:,1);
G = imCFA(:,:,2);
B = imCFA(:,:,3);


Kr = zeros(imSize(1), imSize(2));

% find Kr values at R pixels (row,col)
for i=1:2:imSize(1)

    for j=1:2:imSize(2)
        
        Rcenter = R(i,j);

        Rup1 = 0; Rup2 = 0; Rdown1 = 0; Rdown2 = 0;
        Rleft1 = 0; Rleft2 = 0; Rright1 = 0; Rright2 = 0;
        Gup1 = 0; Gup2 = 0; Gdown1 = 0; Gdown2 = 0;
        Gleft1 = 0; Gleft2 = 0; Gright1 = 0; Gright2 = 0;
        
        % Find vertical coefficients
        if (i == 1) 
            Rdown1 = R(i+2,j);
            Rdown2 = R(i+4,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        elseif (i == 3)
            Rup1 = R(i-2,j);
            Rdown1 = R(i+2,j);
            Rdown2 = R(i+4,j);
            Gup1 = G(i-1,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        elseif (i == (imSize(1)-1))
            Rup1 = R(i-2,j);
            Rup2 = R(i-4,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
        elseif (i == (imSize(1)-3))
            Rup1 = R(i-2,j);
            Rup2 = R(i-4,j);
            Rdown1 = R(i+2,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        else
            Rup1 = R(i-2,j);
            Rup2 = R(i-4,j);
            Rdown1 = R(i+2,j);
            Rdown2 = R(i+4,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        end
        
        
        % Find horizontal coefficents
        if (j == 1)
            Rright1 = R(i,j+2);
            Rright2 = R(i,j+4);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        elseif (j == 3)
            Rleft1 = R(i,j-2);
            Rright1 = R(i,j+2);
            Rright2 = R(i,j+4);
            Gleft1 = G(i,j-1);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        elseif (j == imSize(2)-1)
            Rleft1 = R(i,j-2);
            Rleft2 = R(i,j-4);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
        elseif (j == imSize(2)-3)
            Rleft1 = R(i,j-2);
            Rleft2 = R(i,j-4);
            Rright1 = R(i,j+2);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        else
            Rleft1 = R(i,j-2);
            Rleft2 = R(i,j-4);
            Rright1 = R(i,j+2);
            Rright2 = R(i,j+4);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        end
  
        % Check for smoothness of edges
        
        if (edgeType == ADAPTIVE)
            V = 0.5*abs( (Gup2-Gdown1)+0.5*(-Rup2-Rup1+Rcenter+Rdown1) ) + 0.5*abs( (Gup1-Gdown2)+0.5*(-Rup1-Rcenter+Rdown1+Rdown2) );
            H = 0.5*abs( (Gleft2-Gright1)+0.5*(-Rleft2-Rleft1+Rcenter+Rright1) ) + 0.5*abs( (Gleft1-Gright2)+0.5*(-Rleft1-Rcenter+Rright1+Rright2) );

        elseif (edgeType == VERTICAL)
            V = 0; H = 1;
        else
            V = 1; H = 0;
        end
        
        
        if (V < H) % Kr is smoother in vertical than horizontal direction
            Kr(i,j) = 0.5*(Gup1+Gdown1) - 0.25*(Rup1+2*Rcenter+Rdown1);
        else
            Kr(i,j) = 0.5*(Gleft1+Gright1) - 0.25*(Rleft1+2*Rcenter+Rright1);
        end

     end
   
end


Kb = zeros(imSize(1),imSize(2));

% find Kb values at B pixels (row,col)
for i=2:2:imSize(1)

    for j=2:2:imSize(2)
        
        Bcenter = B(i,j);

        Bup1 = 0; Bup2 = 0; Bdown1 = 0; Bdown2 = 0;
        Bleft1 = 0; Bleft2 = 0; Bright1 = 0; Bright2 = 0;
        Gup1 = 0; Gup2 = 0; Gdown1 = 0; Gdown2 = 0;
        Gleft1 = 0; Gleft2 = 0; Gright1 = 0; Gright2 = 0;
        
        % Find vertical coefficients
        if (i == 2) 
            Bdown1 = B(i+2,j);
            Bdown2 = B(i+4,j);
            Gup1 = G(i-1,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        elseif (i == 4)
            Bup1 = B(i-2,j);
            Bdown1 = B(i+2,j);
            Bdown2 = B(i+4,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        elseif (i == imSize(1))
            Bup1 = B(i-2,j);
            Bup2 = B(i-4,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
        elseif (i == imSize(1)-2)
            Bup1 = B(i-2,j);
            Bup2 = B(i-4,j);
            Bdown1 = B(i+2,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
        else
            Bup1 = B(i-2,j);
            Bup2 = B(i-4,j);
            Bdown1 = B(i+2,j);
            Bdown2 = B(i+4,j);
            Gup1 = G(i-1,j);
            Gup2 = G(i-3,j);
            Gdown1 = G(i+1,j);
            Gdown2 = G(i+3,j);
        end
        
        
        % Find horizontal coefficents
        if (j == 2)
            Bright1 = B(i,j+2);
            Bright2 = B(i,j+4);
            Gleft1 = G(i,j-1);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        elseif (j == 4)
            Bleft1 = B(i,j-2);
            Bright1 = B(i,j+2);
            Bright2 = B(i,j+4);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        elseif (j == imSize(2))
            Bleft1 = B(i,j-2);
            Bleft2 = B(i,j-4);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
        elseif (j == imSize(2)-2)
            Bleft1 = B(i,j-2);
            Bleft2 = B(i,j-4);
            Bright1 = B(i,j+2);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
        else
            Bleft1 = B(i,j-2);
            Bleft2 = B(i,j-4);
            Bright1 = B(i,j+2);
            Bright2 = B(i,j+4);
            Gleft1 = G(i,j-1);
            Gleft2 = G(i,j-3);
            Gright1 = G(i,j+1);
            Gright2 = G(i,j+3);
        end
  
        % Check for smoothness of edges
        
        if (edgeType == ADAPTIVE)
            V = 0.5*abs( (Gup2-Gdown1)+0.5*(-Bup2-Bup1+Bcenter+Bdown1) ) + 0.5*abs( (Gup1-Gdown2)+0.5*(-Bup1-Bcenter+Bdown1+Bdown2) );
            H = 0.5*abs( (Gleft2-Gright1)+0.5*(-Bleft2-Bleft1+Bcenter+Bright1) ) + 0.5*abs( (Gleft1-Gright2)+0.5*(-Bleft1-Bcenter+Bright1+Bright2) );

        elseif (edgeType == VERTICAL)
            V = 0; H = 1;
        else
            V = 1; H = 0;
        end


        if (V < H) % Kb is smoother in vertical than horizontal direction
            Kb(i,j) = 0.5*(Gup1+Gdown1) - 0.25*(Bup1+2*Bcenter+Bdown1);
        else
            Kb(i,j) = 0.5*(Gleft1+Gright1) - 0.25*(Bleft1+2*Bcenter+Bright1);
        end 
    end
end
%%
% Now we have Kr pixels at R locations and Kb pixels at B locations
% Interpolate the remaining pixels in each Kr and Kb by taking the average
% of four corners

% Fill in Kr pixels at B locations
for i=2:2:imSize(1)
    for j=2:2:imSize(2)
        
        Kr1 = 0; Kr2 = 0; Kr3 = 0; Kr4 = 0;
        
        if (i == imSize(1) && j == imSize(2))
            Kr1 = Kr(i-1,j-1);
        elseif (i == imSize(1))
            Kr1 = Kr(i-1,j-1);
            Kr2 = Kr(i-1,j+1);
        elseif (j == imSize(2))
            Kr1 = Kr(i-1,j-1);
            Kr3 = Kr(i+1,j-1);
        else
            Kr1 = Kr(i-1,j-1);
            Kr2 = Kr(i-1,j+1);
            Kr3 = Kr(i+1,j-1);
            Kr4 = Kr(i+1,j+1);
        end
        
        Kr(i,j) = 0.25*(Kr1 + Kr2 + Kr3 + Kr4);
    end
end
        

% Fill in Kb pixels at R locations
for i=1:2:imSize(1)
    for j=1:2:imSize(2)
        
        Kb1 = 0; Kb2 = 0; Kb3 = 0; Kb4 = 0;
        
        if (i == 1 && j == 1)
            Kb4 = Kb(i+1,j+1);
        elseif (i == 1)
            Kb3 = Kb(i+1,j-1);
            Kb4 = Kb(i+1,j+1);
        elseif (j == 1)
            Kb2 = Kb(i-1,j+1);
            Kb4 = Kb(i+1,j+1);
        else
            Kb1 = Kb(i-1,j-1);
            Kb2 = Kb(i-1,j+1);
            Kb3 = Kb(i+1,j-1);
            Kb4 = Kb(i+1,j+1);
        end
        
        Kb(i,j) = 0.25*(Kb1 + Kb2 + Kb3 + Kb4);
    end
end

%%
% Fill in Kr and Kb pixels at G locations (pt1)
for i=1:2:imSize(1)
    for j=2:2:imSize(2)
        Krup1 = 0; Krup2 = 0; Krdown1 = 0; Krdown2 = 0;
        Krleft1 = 0; Krleft2 = 0; Krright1 = 0; Krright2 = 0;
        
        Kbup1 = 0; Kbup2 = 0; Kbdown1 = 0; Kbdown2 = 0;
        Kbleft1 = 0; Kbleft2 = 0; Kbright1 = 0; Kbright2 = 0;
        
        % Find vertical coefficients
        if (i == 1)
            Krdown1 = Kr(i+1,j);
            Krdown2 = Kr(i+3,j);
            
            Kbdown1 = Kb(i+1,j);
            Kbdown2 = Kb(i+3,j);
        elseif (i == 3)
            Krup1 = Kr(i-1,j);
            Krdown1 = Kr(i+1,j);
            Krdown2 = Kr(i+3,j);
            
            Kbup1 = Kb(i-1,j);
            Kbdown1 = Kb(i+1,j);
            Kbdown2 = Kb(i+3,j);
        elseif (i == imSize(1)-1)
            Krup1 = Kr(i-1,j);
            Krup2 = Kr(i-3,j);
            Krdown1 = Kr(i+1,j);
            
            Kbup1 = Kb(i-1,j);
            Kbup2 = Kb(i-3,j);
            Kbdown1 = Kb(i+1,j);
        else
            Krup1 = Kr(i-1,j);
            Krup2 = Kr(i-3,j);
            Krdown1 = Kr(i+1,j);
            Krdown2 = Kr(i+3,j);
            
            Kbup1 = Kb(i-1,j);
            Kbup2 = Kb(i-3,j);
            Kbdown1 = Kb(i+1,j);
            Kbdown2 = Kb(i+3,j);
        end
        
        % Find horizontal coefficients
        if (j == 2)
            Krleft1 = Kr(i,j-1);
            Krright1 = Kr(i,j+1);
            Krright2 = Kr(i,j+3);
        
            Kbleft1 = Kb(i,j-1);
            Kbright1 = Kb(i,j+1);
            Kbright2 = Kb(i,j+3);
        elseif (j == imSize(2))
            Krleft1 = Kr(i,j-1);
            Krleft2 = Kr(i,j-3);
            
            Kbleft1 = Kb(i,j-1);
            Kbleft2 = Kb(i,j-3);
        elseif (j == imSize(2)-2)
            Krleft1 = Kr(i,j-1);
            Krleft2 = Kr(i,j-3);
            Krright1 = Kr(i,j+1);
            
            Kbleft1 = Kb(i,j-1);
            Kbleft2 = Kb(i,j-3);
            Kbright1 = Kb(i,j+1);
        else
            Krleft1 = Kr(i,j-1);
            Krleft2 = Kr(i,j-3);
            Krright1 = Kr(i,j+1);
            Krright2 = Kr(i,j+3);
            
            Kbleft1 = Kb(i,j-1);
            Kbleft2 = Kb(i,j-3);
            Kbright1 = Kb(i,j+1);
            Kbright2 = Kb(i,j+3);
        end

        
        if (edgeType == ADAPTIVE) 
            Vr = 0.5*abs( (Krdown1-Krup2) ) + 0.5*abs( (Krup1-Krdown2) );
            Vb = 0.5*abs( (Kbdown1-Kbup2) ) + 0.5*abs( (Kbup1-Kbdown2) );

            Hr = 0.5*abs( (Krright1-Krleft2) ) + 0.5*abs( (Krleft1-Krright2) );
            Hb = 0.5*abs( (Kbright1-Kbleft2) ) + 0.5*abs( (Kbleft1-Kbright2) );

        elseif (edgeType == VERTICAL)
            Vr = 0; Vb = 0;
            Hr = 1; Hb = 1;
        else
            Vr = 1; Vb = 1;
            Hr = 0; Hb = 0;
        end
        
        
        if (Vr < Hr) 
            Kr(i,j) = (Krup1+Krdown1)/2;
        else
            Kr(i,j) = (Krleft1+Krright1)/2;
        end
        
        if (Vb < Hb)
            Kb(i,j) = (Kbup1+Kbdown1)/2;
        else
            Kb(i,j) = (Kbleft1+Kbright1)/2;
        end
        
    end
end



% Fill in Kr and Kb pixels at G locations (pt2)
for i=2:2:imSize(1)
    for j=1:2:imSize(2)
        Krup1 = 0; Krup2 = 0; Krdown1 = 0; Krdown2 = 0;
        Krleft1 = 0; Krleft2 = 0; Krright1 = 0; Krright2 = 0;
        
        Kbup1 = 0; Kbup2 = 0; Kbdown1 = 0; Kbdown2 = 0;
        Kbleft1 = 0; Kbleft2 = 0; Kbright1 = 0; Kbright2 = 0;
        
        
        % Find vertical coefficients
        if (i == 2)
            Krup1 = Kr(i-1,j);
            Krdown1 = Kr(i+1,j);
            Krdown2 = Kr(i+3,j);
        
            Kbup1 = Kb(i-1,j);
            Kbdown1 = Kb(i+1,j);
            Kbdown2 = Kb(i+3,j);
        elseif (i == imSize(1))
            Krup1 = Kr(i-1,j);
            Krup2 = Kr(i-3,j);
            
            Kbup1 = Kb(i-1,j);
            Kbup2 = Kb(i-3,j);
        elseif (i == imSize(1)-2)
            Krup1 = Kr(i-1,j);
            Krup2 = Kr(i-3,j);
            Krdown1 = Kr(i+1,j);
            
            Kbup1 = Kb(i-1,j);
            Kbup2 = Kb(i-3,j);
            Kbdown1 = Kb(i+1,j);
        else
            Krup1 = Kr(i-1,j);
            Krup2 = Kr(i-3,j);
            Krdown1 = Kr(i+1,j);
            Krdown2 = Kr(i+3,j);
            
            Kbup1 = Kb(i-1,j);
            Kbup2 = Kb(i-3,j);
            Kbdown1 = Kb(i+1,j);
            Kbdown2 = Kb(i+3,j);
        end
        
        % Find horizontal coefficients
        if (j == 1)
            Krright1 = Kr(i,j+1);
            Krright2 = Kr(i,j+3);
            
            Kbright1 = Kb(i,j+1);
            Kbright2 = Kb(i,j+3);
        elseif (j == 3)
            Krleft1 = Kr(i,j-1);
            Krright1 = Kr(i,j+1);
            Krright2 = Kr(i,j+3);
            
            Kbleft1 = Kb(i,j-1);
            Kbright1 = Kb(i,j+1);
            Kbright2 = Kb(i,j+3);
        elseif (j == imSize(2)-1)
            Krleft1 = Kr(i,j-1);
            Krleft2 = Kr(i,j-3);
            Krright1 = Kr(i,j+1);
            
            Kbleft1 = Kb(i,j-1);
            Kbleft2 = Kb(i,j-3);
            Kbright1 = Kb(i,j+1);
        else
            Krleft1 = Kr(i,j-1);
            Krleft2 = Kr(i,j-3);
            Krright1 = Kr(i,j+1);
            Krright2 = Kr(i,j+3);
            
            Kbleft1 = Kb(i,j-1);
            Kbleft2 = Kb(i,j-3);
            Kbright1 = Kb(i,j+1);
            Kbright2 = Kb(i,j+3);
        end
        

        if (edgeType == ADAPTIVE) 
            Vr = 0.5*abs( (Krdown1-Krup2) ) + 0.5*abs( (Krup1-Krdown2) );
            Vb = 0.5*abs( (Kbdown1-Kbup2) ) + 0.5*abs( (Kbup1-Kbdown2) );

            Hr = 0.5*abs( (Krright1-Krleft2) ) + 0.5*abs( (Krleft1-Krright2) );
            Hb = 0.5*abs( (Kbright1-Kbleft2) ) + 0.5*abs( (Kbleft1-Kbright2) );
            
        elseif (edgeType == VERTICAL)
            Vr = 0; Vb = 0;
            Hr = 1; Hb = 1;
        else
            Vr = 1; Vb = 1;
            Hr = 0; Hb = 0;
        end
        
        
        if (Vr < Hr) 
            Kr(i,j) = (Krup1+Krdown1)/2;
        else
            Kr(i,j) = (Krleft1+Krright1)/2;
        end
        
        if (Vb < Hb)
            Kb(i,j) = (Kbup1+Kbdown1)/2;
        else
            Kb(i,j) = (Kbleft1+Kbright1)/2;
        end
        
    end
end

%%

% Now we have fully interpolated Kr and Kb, so we need to go back to RGB

Rnew = zeros(imSize(1),imSize(2));
Gnew = zeros(imSize(1),imSize(2));
Bnew = zeros(imSize(1),imSize(2));

% Pixels at original R locations
for i=1:2:imSize(1)
    for j=1:2:imSize(2)
        Rnew(i,j) = R(i,j);
        Gnew(i,j) = Kr(i,j) + R(i,j);
        Bnew(i,j) = Gnew(i,j) - Kb(i,j);
    end
end

% Pixels at original B locations
for i=2:2:imSize(1)
    for j=2:2:imSize(2)
        Bnew(i,j) = B(i,j);
        Gnew(i,j) = Kb(i,j) + B(i,j);
        Rnew(i,j) = Gnew(i,j) - Kr(i,j);
    end
end

% Pixels at original G locations (pt1)
for i=1:2:imSize(1)
    for j=2:2:imSize(2)
        Gnew(i,j) = G(i,j);
        Rnew(i,j) = G(i,j) - Kr(i,j);
        Bnew(i,j) = G(i,j) - Kb(i,j);
    end
end

% Pixels at original G locations (pt2)
for i=2:2:imSize(1)
    for j=1:2:imSize(2)
        Gnew(i,j) = G(i,j);
        Rnew(i,j) = G(i,j) - Kr(i,j);
        Bnew(i,j) = G(i,j) - Kb(i,j);
    end
end

newIm(:,:,1) = Rnew;
newIm(:,:,2) = Gnew;
newIm(:,:,3) = Bnew;

% Cap the values that are out of limits
newIm(find(newIm > 255))=255;
newIm(find(newIm < 0)) = 0;

time = cputime - t1

% mse = mean(mean((im-newIm).^2)) 
% psnr = 10*log10(((max(max(newIm))).^2)./mse)

% figure
origIm = im/255;
% image(im/255);
% axis equal;
% title('Original Image');

% figure
resultIm = newIm/255;
% image(newIm/255);
% axis equal;
% title('Processed Image');