%parameters: image1: 2, 32; image2: 8, 32, image3: 8, 32

function flash()
  compute('image1', 2.0, 32.0);
  compute('image2', 8.0, 32.0);
  compute('image3', 8.0, 32.0);
end

function compute(imageFile, r, s)
  A = imread(strcat(imageFile, '_noflash.bmp'));    
  F = imread(strcat(imageFile, '_flash.bmp'));
  A = double(A);
  F = double(F);
  
  ABase = jointbilateralfilter(A, A, r, s);
  imwrite(uint8(ABase), strcat(imageFile, '_base.bmp'));
  ANR = jointbilateralfilter(A, F, r, s);
  imwrite(uint8(ANR), strcat(imageFile, '_nr.bmp'));
  FBase = jointbilateralfilter(F, F, r, s);  
  FDetail = 255.0 * (F + 1.0) ./ (FBase + 1.0);
  imwrite(uint8(FDetail), strcat(imageFile, '_fdetail.bmp'));  
  
  [AH, AS, AV] = rgb2hsv(A);
  [FH, FS, FV] = rgb2hsv(F);
  MShadow = shadowmask(AV, FV);
  MSpecular = specularmask(FV);
  Mask = max(MShadow, MSpecular);
  imwrite(MShadow, strcat(imageFile, '_shadow.bmp'));
  imwrite(MSpecular, strcat(imageFile, '_specular.bmp'));
  imwrite(Mask, strcat(imageFile, '_mask.bmp'));
  
  AFinal(:, :, 1) = (1.0 - Mask) .* ANR(:, :, 1) .* FDetail(:, :, 1) / 255.0 + Mask .* ABase(:, :, 1);        
  AFinal(:, :, 2) = (1.0 - Mask) .* ANR(:, :, 2) .* FDetail(:, :, 2) / 255.0 + Mask .* ABase(:, :, 2);        
  AFinal(:, :, 3) = (1.0 - Mask) .* ANR(:, :, 3) .* FDetail(:, :, 3) / 255.0 + Mask .* ABase(:, :, 3);        
  imwrite(uint8(AFinal), strcat(imageFile, '_output.bmp'));  
end

function MShadow = shadowmask(A, F)
%{
  MShadow = A;
  threshold = 0.01;
  diff = abs(F - A) ./ (F + 0.0001);
  for i = 1 : size(A, 1) * size(A, 2)    
    if diff(i) < threshold
      MShadow(i) = 1.0;
    else
      MShadow(i) = 0.0;
    end
  end
%}
  MShadow = F;
  threshold = 0.10;
  maxV = max(max(F));
  for i = 1 : size(F, 1) * size(F, 2)
    if F(i) / maxV < threshold
      MShadow(i) = 1.0;
    else
      MShadow(i) = 0.0;
    end
  end
end

function MSpecular = specularmask(F)
  MSpecular = F;
  threshold = 0.90;
  maxV = max(max(F));
  for i = 1 : size(F, 1) * size(F, 2)
    if F(i) / maxV > threshold
      MSpecular(i) = 1.0;
    else
      MSpecular(i) = 0.0;
    end
  end
end

function ANR = jointbilateralfilter(A, F, r, s)  
  h = fspecial('gaussian', 2.0 * r + 1.0, (2.0 * r + 1.0) / 3.0);
  tt = 2.0 * s * s;
  t = 1.0 / (pi * tt);
  ANR = A;
  for i = 1 : size(A, 1)
    for j = 1 : size(A, 2)     
      total = zeros(3);
      totalWeight = 0.0;    
      for k = -r : r
        for l = -r : r
          spaceWeight = h(k + r + 1, l + r + 1);
          ik = min(max(i + k, 1), size(A, 1));
          jl = min(max(j + l, 1), size(A, 2));
          diffR = F(i, j, 1) - F(ik, jl, 1);
          diffG = F(i, j, 2) - F(ik, jl, 2);
          diffB = F(i, j, 3) - F(ik, jl, 3);
          diff = diffR * diffR + diffG * diffG + diffB * diffB;
          rangeWeight =  t * exp(-diff / tt);
          weight = spaceWeight * rangeWeight;
          totalWeight = totalWeight + weight;
          total(1) = total(1) + weight * A(ik, jl, 1);
          total(2) = total(2) + weight * A(ik, jl, 2);
          total(3) = total(3) + weight * A(ik, jl, 3);
        end
      end
      ANR(i, j, 1) = total(1) / totalWeight;
      ANR(i, j, 2) = total(2) / totalWeight;
      ANR(i, j, 3) = total(3) / totalWeight;
    end
  end
end