clear;
% Define the directory where the binary frames are located
videoDir = 'FilePath';

% Load binary frames from the pre-saved .mat file
load(fullfile(videoDir, 'binaryFrames.mat'), 'binaryFrames');


% Ensure startFrame and endFrame are within the valid range of frames
startFrame = 1; %max(1, startFrame);  % At least the first frame
endFrame = max(1, length(binaryFrames));  % At most the total number of frames

% Set parameters for tracking
frameRate = 10;  % Assuming 10 frames per second (adjust as needed)
timeBetweenFrames = 1 / frameRate;  % Time between frames in seconds

minObjectSize = 40; % Minimum size to filter out noise
maxDistance = 15;  % Maximum distance to consider object tracked in next frame
areaChangeThreshold = 0.20;  % Area change threshold for splitting/merging
areaBinSize = 100;  % Bin size for area
aspectRatioBinSize = 0.1;  % Bin size for aspect ratio

% Initialize tracks structure with unique IDs
tracks = struct('id', {}, 'centroids', {}, 'areas', {}, 'aspectRatios', {}, 'frameStart', {}, 'frameEnd', {}, 'active', {});

uniqueID = 1;  % Start unique ID counter
objectCount = 0; % Initialize object count

% Arrays for storing data for plots
velocities = [];
areas = [];
aspectRatios = [];
trackDurations = [];

% Sequential tracking loop
figure(1);  % Create a figure for visualization

% Pre-allocate an array for storing valid frame indices and their velocities
validFrameNumbers = [];
validFrameVelocities = [];

% Loop through frames between startFrame and endFrame
for frameNumber = startFrame:endFrame
    % Read pre-loaded binary frame
    binaryImage = binaryFrames{frameNumber};

    % Remove small objects
    binaryImage = bwareaopen(binaryImage, minObjectSize);

    % Display the current frame number in the command window
    disp(['Processing Frame: ' num2str(frameNumber)]);

    % Get centroids, areas, and bounding boxes for aspect ratio calculation
    stats = regionprops(binaryImage, 'Centroid', 'Area', 'BoundingBox');
    centroids = cat(1, stats.Centroid);
    areasCurrent = [stats.Area];
    aspectRatiosCurrent = arrayfun(@(x) max(x.BoundingBox(3:4)) / min(x.BoundingBox(3:4)), stats);

    % Handle objects tracked from previous frames
    if frameNumber > startFrame
        frameVelocities = [];  % Store velocities for this frame
        
        % Get active tracks (objects being tracked)
        activeTrackIndices = find([tracks.active]);
        
        % Check for matches with centroids in the current frame
        if ~isempty(activeTrackIndices) && ~isempty(centroids)
            previousCentroids = cell2mat(arrayfun(@(x) x.centroids(end, :), tracks(activeTrackIndices), 'UniformOutput', false)');
            distances = pdist2(previousCentroids, centroids);  % Pairwise distances between old and new centroids
            
            % For each tracked object, find the closest object in the current frame
            [minDistances, closestIdx] = min(distances, [], 2);

            % Update tracks with matched centroids
            for i = 1:length(activeTrackIndices)
                trackIdx = activeTrackIndices(i);
                if minDistances(i) <= maxDistance
                    % Check area change
                    previousArea = tracks(trackIdx).areas(end);
                    currentArea = areasCurrent(closestIdx(i));
                    areaChange = abs(currentArea - previousArea) / previousArea;

                    if areaChange <= areaChangeThreshold
                        % Update track with new centroid, area, and aspect ratio
                        tracks(trackIdx).centroids = [tracks(trackIdx).centroids; centroids(closestIdx(i), :)];
                        tracks(trackIdx).areas = [tracks(trackIdx).areas; currentArea];
                        tracks(trackIdx).aspectRatios = [tracks(trackIdx).aspectRatios; aspectRatiosCurrent(closestIdx(i))];
                        tracks(trackIdx).frameEnd = frameNumber;

                        % Calculate velocity
                        if size(tracks(trackIdx).centroids, 1) > 1
                            distanceTravelled = sqrt(sum((tracks(trackIdx).centroids(end, :) - tracks(trackIdx).centroids(end-1, :)).^2));
                            velocity = distanceTravelled / timeBetweenFrames;
                            frameVelocities = [frameVelocities; velocity];
                            velocities = [velocities; velocity];
                            areas = [areas; currentArea];
                            aspectRatios = [aspectRatios; aspectRatiosCurrent(closestIdx(i))];
                        end

                        % Mark centroid as processed
                        centroids(closestIdx(i), :) = NaN;
                        areasCurrent(closestIdx(i)) = NaN;
                    else
                        % Keep track inactive but retain its history
                        tracks(trackIdx).active = false;
                    end
                else
                    % Track is inactive, but retain its history
                    tracks(trackIdx).active = false;
                end
            end
        end
    end

    % Initialize new tracks for unmatched objects
    newObjectIndices = find(~isnan(centroids(:, 1)));  % Objects that are new or unmatched
    for i = 1:length(newObjectIndices)
        idx = newObjectIndices(i);
        tracks(end+1) = struct('id', uniqueID, 'centroids', centroids(idx, :), ...
                               'areas', areasCurrent(idx), 'aspectRatios', aspectRatiosCurrent(idx), ...
                               'frameStart', frameNumber, 'frameEnd', frameNumber, 'active', true);
        uniqueID = uniqueID + 1;  % Increment unique ID for the next new object
    end

    % Store average velocity for valid frames
    if frameNumber > startFrame && ~isempty(frameVelocities)
        validFrameNumbers = [validFrameNumbers, frameNumber];  % Store valid frame number
        validFrameVelocities = [validFrameVelocities, mean(frameVelocities)];  % Store the average velocity
    end

    % Visualization of active tracks
    imshow(binaryImage);  % Update the image data for current frame
    hold on;
    for j = 1:length(tracks)
        if tracks(j).active && ~isempty(tracks(j).centroids)  % Only visualize active tracks
            validCentroids = tracks(j).centroids(~isnan(tracks(j).centroids(:, 1)), :);
            if size(validCentroids, 1) > 1
                plot(validCentroids(:,1), validCentroids(:,2), '-', 'MarkerSize', 3, 'LineWidth', 1.5);
            end
        end
    end
    hold off;
    title(['Frame ' num2str(frameNumber)]);
    drawnow;
    pause(0.1);  % Optional pause to control visualization speed
end


% Save the tracking data to a .mat file
save(fullfile(videoDir, 'trackedData.mat'), 'tracks', 'velocities', 'areas', 'aspectRatios');
disp('Tracking data saved successfully.');


% Plot average instantaneous velocity vs frame number
figure;
plot(validFrameNumbers, validFrameVelocities, '-o');
xlabel('Frame Number');
ylabel('Average Instantaneous Velocity');
title('Average Instantaneous Velocity vs Frame Number');

saveas(gcf, fullfile(videoDir, 'avg_velocity_vs_frame.png'));  % Save as PNG

% Bin frame numbers by 50 frames and calculate average velocity for each bin
binSize = 50;
binnedFrameNumbers = [];
binnedVelocities = [];
for i = 1:binSize:length(validFrameNumbers)
    if i+binSize-1 <= length(validFrameNumbers)
        binnedFrameNumbers = [binnedFrameNumbers, mean(validFrameNumbers(i:i+binSize-1))];
        binnedVelocities = [binnedVelocities, mean(validFrameVelocities(i:i+binSize-1))];
    else
        disp('Not enough frames to bin');
        break;
    end
end

% Plot binned average instantaneous velocity
figure;
plot(binnedFrameNumbers, binnedVelocities, '-o');
xlabel('Binned Frame Number');
ylabel('Average Instantaneous Velocity');
title('Binned (50 frames) Average Instantaneous Velocity');

saveas(gcf, fullfile(videoDir, 'binned_avg_velocity.png'));  % Save as PNG