%% find the linear colour spill
% written by Natalie Ko-Ferrigno

prefix = "/Users/natalie/Documents/year 3/IP/code/temp/25.03.2024/images 2/";
% where the red, green and blue files are stored in that order
paths = ["%sR/%s" "%sG/%s" "%sB/%s"];
% the images to use, in the order red, green, blue
names = [{dir(sprintf(paths(1),prefix,  "*.tiff")).name}; 
    {dir(sprintf(paths(2),prefix,  "*.tiff")).name};
    {dir(sprintf(paths(3),prefix,  "*.tiff")).name}];

%paths = ["%sRED/RED_F8.0_25ms/%s"];
%names = ["142339.tiff"];

labels = ["r" "g" "b"];
 
coeffs = zeros(3, 3);

for color = 1:3
    intensities = [];

    for i = 1:size(names(1,:), 2)
        current_path = sprintf(paths(color), prefix, names{color, i});
        temp = get_correlation(current_path);
        intensities = cat(1, intensities, temp);
    end
    
    
    x = linspace(0, 1);

    hold on;
        
    c_r = fit(intensities(:, color), intensities(:, 1), 'p1*x');
    c_g = fit(intensities(:, color), intensities(:, 2), 'p1*x');
    c_b = fit(intensities(:, color), intensities(:, 3), 'p1*x');

    fprintf("done %d\n", color);
    scatter(intensities(:, color), intensities(:, 1), 1, "r");
    scatter(intensities(:, color), intensities(:, 2), 1, "g");
    scatter(intensities(:, color), intensities(:, 3), 1, "b");
    

    plot(x, c_r.p1 * x, "m");
    plot(x, c_b.p1 * x, "c");
    plot(x, c_g.p1 * x, "y");

    ylim([0 1]);
    xlim([0 1]);

    
    legend("R", "G", "B", "R fit", "G fit", "B fit");
    title("correlation for ", labels(color))
    hold off;
    coeffs(color, :) = coeffs(color, :) + [c_r.p1, c_g.p1, c_b.p1];
    figure();
end
fprintf("done\n");


function [intensities] = get_correlation(path)
    mask = "bggr"; 
    image = imread(path);
    
    rgb_int = demosaic(image, mask);
    
    [i, j, brightness, color] = rgb_find_maxima(rgb_int, [50 50 40]);
     %show the places identified as particles
    mask = zeros(3040, 4064, 'uint16');
    peaks = zeros(length(i), 3);
    for point = 1:length(i)
        mask(i(point), j(point)) = 1;
    
        peaks(point, 1) = i(point);
        peaks(point, 2) = j(point);
        peaks(point, 3) = brightness(point);
    end
    
    rgb = mask .* rgb_int;
    
    intensities = [];
    
    % store the intensities at each peak
    
    for channel = 1:3
        [~, ~, k] = find(rgb(:, :, channel));
        k = cast(k, "double")/bitshift(1, 16); % normalise to 0 to 1 for fun
        intensities = [intensities, k]; % dont think i need to preallocate, its only 3 loops
    end
    
    
end

function [i,j,val,b_color] = rgb_find_maxima(img, threshold)
	% [i,j] = find_maxima(im, threshold)
	%
	% find indices i,j of entries in im which are local maxima
	% and above a given threshold
    % shamelessly ripped from rgb_find_particles.m
	
	%% find indices above lowest threshold, sufficiently far from edges
	im_size			= size(img);
	min_thresh		= min(threshold);
	[i,j]			= find(img >= min_thresh);
	ok				= i > 2 & i < im_size(1)-1 ...
					& j > 2 & j < im_size(2)-1;
    
	i				= i(ok);
	j				= j(ok);
	
	%% categories as red, blue or green
	% if the local maximum occurs on a red pixel, it is a "red" particle
	% likewise for green and blue
	b_blue			= mod(i,2) == 1 & mod(j, 2) == 1;
	b_red			= mod(i,2) == 0 & mod(j, 2) == 0;
	b_green			= ~b_blue & ~b_red;
	val				= img(sub2ind(im_size, i  ,j  ));
	% apply red/green/blue threshold independently
	ok				= false(size(i));
	ok(b_red)		= val(b_red ) > threshold(1);
	ok(b_green)		= val(b_green)> threshold(2);
	ok(b_blue)		= val(b_blue) > threshold(3);
	i				= i(ok);
	j				= j(ok);
	
	%% of these, find indices that are local maxima
	val				= img(sub2ind(im_size, i  ,j  ));
	is_max			= img(sub2ind(im_size, i  ,j-2)) < val ...
					& img(sub2ind(im_size, i  ,j+2)) < val ...
					& img(sub2ind(im_size, i-2,j  )) < val ...
					& img(sub2ind(im_size, i+2,j  )) < val;
	val				= val(is_max);
	i				= i(is_max);
	j				= j(is_max);
	b_blue			= mod(i,2) == 1 & mod(j, 2) == 1;
	b_red			= mod(i,2) == 0 & mod(j, 2) == 0;
	b_green			= ~b_blue & ~b_red;
	b_color			= [b_red b_green b_blue];
	
	%% filter close neighbours of green pixels
	% green pixels must be additionally maxima in a local neighbourhood
	% to avoid identifying the same particle twice close together 
	i_green			= i(b_green);
	j_green			= j(b_green);
	v_green			= val(b_green);
	i_green			= i_green(:);
	j_green			= j_green(:);
	v_green			= v_green(:);
	is_green_max	= img(sub2ind(im_size, i_green+1, j_green+1)) < v_green ...
					& img(sub2ind(im_size, i_green+1, j_green-1)) < v_green ...
					& img(sub2ind(im_size, i_green-1, j_green+1)) < v_green ...
					& img(sub2ind(im_size, i_green-1, j_green-1)) < v_green;
	n_ok			= nnz(is_green_max);
	% cut and replace
	i(b_green)		= [];
	j(b_green)		= [];
	val(b_green)	= [];
	b_color(b_green,:)=[];
	i				= [i; i_green(is_green_max)];
	j				= [j; j_green(is_green_max)];
	val				= [val; v_green(is_green_max)];
	b_color			= [b_color; false(n_ok,1), true(n_ok,1), false(n_ok,1)];
end