%% Add ground truth labels to the current point cloud
% Functions samples specified number of ground truth labels from the
% database and performs collision test and returns the samples that are
% non-overlapping.
function data = groundTruthDataAugmenation(data,processedData, groundTruth,samplesToAdd)
    
    groundTruthData = generateGTDataForAugmentation(processedData, groundTruth)
    ptCloud = data{1,1};
    
    % Convert to Mx4 format from point cloud object.
    ptCloud = cat(2,ptCloud.Location,ptCloud.Intensity);
    bboxesGt = data{1,2};
    bboxesClasses = data{1,3};
    
    fieldNames = fieldnames(samplesToAdd);

    for i = 1:numel(fieldNames)            

        % To remove the pitch and roll from the box coordinates.
        if ~isempty(bboxesGt)
            bboxesGt = bboxesGt(:,[1,2,3,4,5,6,9]);
        end

        % Calculate the number of boxes to add.
        numPCToAdd = samplesToAdd.(fieldNames{i}) - sum(data{1,3} == fieldNames{i});
        numPCToAdd = max(numPCToAdd,0); 
        bboxesToAdd = [];
        pcToAdd = [];

        % Sample the required number of ground truth objects.
        samples = groundTruthData.(fieldNames{i}).sample(numPCToAdd);

        for ii = 1:numel(samples)
            difficulty = samples(ii,1).difficulty; 
            if difficulty == -1 && ~isempty(bboxesGt)
                continue;
            end
            bboxesToAdd = [bboxesToAdd;samples(ii,1).boxDims(:,[1,2,3,4,5,6,7])];
            pcToAdd = [pcToAdd;{samples(ii,1).lidarpoints}];
        end

        % Find overlap within the selected boxes.
        if ~isempty(bboxesToAdd)
            bboxesBEV = bboxesToAdd(:,[1,2,4,5,7]);
            boxscores  = rand(size(bboxesBEV,1),1);
            [~,~,idx] = selectStrongestBbox(bboxesBEV,boxscores,'OverlapThreshold',0);
            bboxesToAdd = bboxesToAdd(idx,:);
            pcToAdd = pcToAdd(idx);
        end

       % Condition wherein both ground truth boxes and boxes to add
       % are present.
       if ~isempty(bboxesGt) && ~isempty(bboxesToAdd)
            overlapRatio = bboxOverlapRatio(bboxesToAdd(:,[1,2,4,5,7]),bboxesGt(:,[1,2,4,5,7]));
            maxOverlap = max(overlapRatio,[],2);
            bboxesToAddFinal = bboxesToAdd(maxOverlap == 0,:);
            pcToAddFinal = pcToAdd(maxOverlap == 0);
            finalGtBoxes = [bboxesGt;bboxesToAddFinal];
            pcToAddClasses = cell(size(bboxesToAddFinal,1),1);
            pcToAddClasses(:) = {fieldNames{i}};
            pcToAddClasses = categorical(pcToAddClasses);
            finalClasses = cat(1,bboxesClasses,pcToAddClasses);

        % Condition wherein there are no ground truth boxes.
        elseif isempty(bboxesGt)
            bboxesToAddFinal = bboxesToAdd;
            pcToAddFinal = pcToAdd;
            finalGtBoxes = bboxesToAdd;
            pcToAddClasses = cell(size(finalGtBoxes,1),1);
            pcToAddClasses(:) = {fieldNames{i}};
            pcToAddClasses = categorical(pcToAddClasses);
            finalClasses = pcToAddClasses;

        % Condition wherein there are no bboxes to add to the
        % existing ground truths.
        elseif isempty(bboxesToAdd)
            finalGtBoxes = bboxesGt;
            bboxesToAddFinal = [];
            finalClasses = bboxesClasses;
        end

        % Extract the points inside the boxes and append to point
        % cloud.
        if ~isempty(bboxesToAddFinal)
           for ii = 1:size(bboxesToAddFinal,1)
                pointsToAdd = pcToAddFinal(ii);
                pointsToAdd = pointsToAdd{1,1};
                if isempty(pointsToAdd)
                    continue;
                end
                ptCloud = [ptCloud;pointsToAdd];
            end
        end

        finalBoxes = zeros([size(finalGtBoxes,1),9]);
        finalBoxes(:,1:6) = finalGtBoxes(:,1:6);
        finalBoxes(:,9) = finalGtBoxes(:,7);
        bboxesGt = finalBoxes;
        bboxesClasses = finalClasses;
    end

    data{1,1} = ptCloud;
    if isempty(finalBoxes)
        fprintf("Error in finalBoxes Augmentation");
    elseif size(finalBoxes,1) ~= size(finalClasses,1)
        fprintf("Error in final Classes Augmentation");
    end
    data{1,2} = finalBoxes;
    data{1,3} = finalClasses;
end
function groundTruthData = generateGTDataForAugmentation(processedData, groundTruth)
% This function is used to extract all the ground truth labels from the
% provided input data.

    groupDbInfos = struct;
    numFiles = size(processedData,1);

    classNames = groundTruth.Properties.VariableNames;
    for j = 1:numel(classNames)
        Data = {};
        for i = 1:numFiles     
            finalPC = processedData{i,1};   
            labels = groundTruth(i,j).Variables;
            labels = labels{1};
               for ii = 1:size(labels,1)
                   label = labels(ii,:);
                   
                   % Find points that are present inside the cuboid and set
                   % the difficulty level.
                   pointsIdx = findPointsInsideCuboid(cuboidModel(label),finalPC);
                   data = struct;
                   data.numPointsInGt = size(pointsIdx,1);
                   loc = cat(2,finalPC.Location,finalPC.Intensity);
                   data.lidarpoints = loc(pointsIdx,:);
                   if size(pointsIdx,1) < 20
                       data.difficulty = -1;
                   else
                       data.difficulty = 1;
                   end
                   data.boxDims = labels(ii,[1,2,3,4,5,6,9]);
                   data.className = classNames{j};
                   Data{end+1} = data;
                end 
        end
        groupDbInfos.(classNames{j}) =  Data;
    end
    fnames = fieldnames(groupDbInfos);
    groundTruthData = struct;
    for i = 1:numel(fnames)
        fname = fnames{i};
        sampledList = groupDbInfos.(fname);
        groundTruthData.(fname) = batchSampler(sampledList, fname);
    end
end