% Load reference path
data = load('parkingLotReferenceData.mat');

% Set reference trajectory of the ego vehicle
refPosesX = data.refPosesX;
refPosesY = data.refPosesY;
refPosesT = data.refPosesT;

% Set poses of the parked vehicles
parkedPoses = data.parkedPoses([18 21],:);

% Display the reference path and the parked vehicle locations
sceneName = 'LargeParkingLot';
hScene = figure('Name', 'Large Parking Lot', 'NumberTitle', 'off');
helperShowSceneImage(sceneName);
hold on
plot(refPosesX(:,2),refPosesY(:,2),'LineWidth',2,'DisplayName','Reference Path');
scatter(parkedPoses(:,1),parkedPoses(:,2),[],'filled','DisplayName','Parked Vehicles');
xlim([-60 40])
ylim([10 60])
hScene.Position = [100,100,1000,500]; % Resize figure
legend
hold off

modelName = 'GenerateLidarDataOfParkingLot';
open_system(modelName)
snapnow

helperAddParkedVehicles(modelName,parkedPoses);

close(hScene)

if ~ispc
    error("Unreal Engine Simulation is supported only on Microsoft" + char(174) + " Windows" + char(174) + ".");
end

% Run simulation
simOut = sim(modelName);

close_system(modelName,0);

% Extract lidar data
ptCloudArr = helperGetPointClouds(simOut);

% Extract ground truth for the lidar data as an array of rigid3d objects
groundTruthPosesLidar = helperGetLidarGroundTruth(simOut);

ptCloud = ptCloudArr(1);

% Remove outliers in the point cloud
ptCloudDenoised = pcdenoise(ptCloud);

% Clip the point cloud. This is done to improve the processing performance
% and also to include only the areas of interest

% Set the limits to select the point cloud
selectLimitX = [-40 40];
selectLimitY = [-40 40];
selectLimitZ = [0 5];

roi = [selectLimitX selectLimitY selectLimitZ];
indices = findPointsInROI(ptCloudDenoised,roi);
ptCloudClipped = select(ptCloudDenoised,indices);

% Segment and remove ego vehicle

% Set location of the sensor and vehicle radius
sensorLocation = [0 0 0];
vehicleRadius = 3.5;

% Find the indices of the points other than the ego vehicle and create a
% point cloud with these points
egoIndices = findNeighborsInRadius(ptCloudClipped,sensorLocation,vehicleRadius);

egoFixed = false(ptCloudClipped.Count,1);
egoFixed(egoIndices) = true;
ptCloudNonEgo = select(ptCloudClipped,~egoFixed);

% Since you use the Z-values to determine the occupancy in a grid image,
% move the point cloud by sensor height so that more points are included to
% calculate occupancy

% Set the sensor height. This information can be obtained from the lidar
% point cloud
sensorHeight = groundTruthPosesLidar(1).Translation(3);

locationPts = ptCloudNonEgo.Location;
locationPts(:,3) = locationPts(:,3) + sensorHeight;
ptCloudProcessed = pointCloud(locationPts);

% Visualize and compare the point cloud before and after preprocessing.
figure('Name', 'Processed Point Clouds', 'NumberTitle','off')
pcViewAxes = pcshowpair(ptCloud, ptCloudProcessed);
title('Point cloud before and after preprocessing')
xlabel(pcViewAxes, 'X (m)')
ylabel(pcViewAxes, 'Y (m)')
zlabel(pcViewAxes, 'Z (m)')

% Set the occupancy grid size to 100 m with a resolution of 0.2 m
gridSize = 100;
gridStep = 0.2;

% The occupancy grid is created by scaling the points from 0-5 m in
% height to the probability values of [0 1]
zLimits = [0 5];

% Calculate the number of bins
spatialLimits = [-gridSize/2 gridSize/2; -gridSize/2 gridSize/2; ptCloudProcessed.ZLimits];

gridXBinSize = round(abs(spatialLimits(1,2) - spatialLimits(1,1)) / gridStep);
gridYBinSize = round(abs(spatialLimits(2,2) - spatialLimits(2,1)) / gridStep);

numBins = [gridXBinSize gridYBinSize 1];

% Find bin indices
binIndices = pcbin(ptCloudProcessed,numBins,spatialLimits,'BinOutput',false);

% Preallocate occupancy grid
occupancyGrid = zeros(gridXBinSize,gridYBinSize,'like',ptCloudProcessed.Location);
gridCount = zeros(gridXBinSize,gridYBinSize);

% Scale the Z values of the points to the probability range [0 1]
zValues = rescale(ptCloudProcessed.Location(:,3),'InputMin',zLimits(1),'InputMax',zLimits(2));

for idx = 1:numel(binIndices)
    binIdx = binIndices(idx);
    if ~isnan(binIdx)
        occupancyGrid(binIdx) = occupancyGrid(binIdx) + zValues(idx);
        gridCount(binIdx) = gridCount(binIdx) + 1;
    end
end

gridCount(gridCount == 0) = 1;

occupancyGrid = occupancyGrid ./ gridCount;

% Visualize the created occupancy grid
figure
subplot(1,2,1)
pcshow(ptCloudProcessed)
view(2)
title('Point Cloud Birds Eye View')
subplot(1,2,2)
imshow(imrotate(occupancyGrid,90))
tHandle = title('Occupancy Grid Image');
tHandle.Color = [1 1 1];

% Pose graph for registration poses
regEstPoseGraph = poseGraph;
relTrans = groundTruthPosesLidar(1).T(4,1:2);
relYaw = -atan2(groundTruthPosesLidar(1).T(2,1),groundTruthPosesLidar(1).T(1,1));

% Get ground truth pose graph
gTPoseGraph = helperGetGTPoseGraph(groundTruthPosesLidar);

% Obtain the nodes from the pose graphs
gTNodes = gTPoseGraph.nodes();

% Plot the ground truth trajectory
figure('Name', 'Vehicle Trajectory', 'NumberTitle',"off")
compareTrajAxes = axes;
compareTrajAxes.XLim = [-10 60];
compareTrajAxes.YLim = [-20 5];
compareTrajAxes.XGrid = 'on';
compareTrajAxes.YGrid = 'on';
title('Compare Trajectories')

hold on
plot(gTNodes(:,1),gTNodes(:,2),'Color','blue','LineWidth',3,'DisplayName','Ground Truth Trajectory')

% Visualize the estimated trajectory and the location of the vehicle when
% running the algorithm
estimatedTraj = plot(gTNodes(1,1),gTNodes(1,2),'Color','green','LineWidth',3 ,'DisplayName','Estimated Trajectory');
currentLoc = scatter(gTNodes(1,1),gTNodes(1,2),50,'red','filled','DisplayName','Vehicle Location');

legend;

% Process every fourth frame. This is done to speed up the processing
% without affecting much accuracy.
skipFrames = 3;

numFrames = numel(groundTruthPosesLidar);

occGridFixed = occupancyGrid';

for frameIdx = 1+skipFrames:skipFrames:numFrames
    
    ptCloud = ptCloudArr(frameIdx);
    
    ptCloudProcessed = helperPreprocessPtCloud(ptCloud);
    
    occGridMoving = helperCreateOccupancyGrid(ptCloudProcessed,gridSize,gridStep,zLimits);
    
    % Registration
    
    % Because imregcorr reports the transformation from the origin, you must
    % provide a reference such that the sensor is at the center of each image
    Rgrid = imref2d(size(occGridMoving));
    offsetX = mean(Rgrid.XWorldLimits);
    Rgrid.XWorldLimits = Rgrid.XWorldLimits - offsetX;
    offsetY = mean(Rgrid.YWorldLimits);
    Rgrid.YWorldLimits = Rgrid.YWorldLimits - offsetY;    
    
    % Find relative pose in image coordinates
    [relPoseImage, peak] = imregcorr(occGridMoving,Rgrid,occGridFixed,Rgrid);
    
    % Convert translation to grid coordinates
    transInGrid = relPoseImage.T(3,1:2) .* gridStep;
    
    % The tranformation is a rigid transformation. Since relative pose is
    % an affine2d object, convert to rigid2d object
    rotations = relPoseImage.T(1:2,1:2);
    [u,~,v] = svd(rotations);
    relPose = rigid2d(u*v',transInGrid);
    
    % Add pose to registration estimate pose graph
    relTrans = relPose.Translation(1:2);
    relYaw = -atan2(relPose.T(2,1),relPose.T(1,1));     
    regEstPoseGraph.addRelativePose([relTrans relYaw]);  
    
    occGridFixed = occGridMoving;
    
    % Update the estimated trajectory and vehicle location
    estimatedNode = regEstPoseGraph.nodes(regEstPoseGraph.NumNodes);
    estimatedTraj.XData(end + 1) = estimatedNode(1);
    estimatedTraj.YData(end + 1) = estimatedNode(2);
    set(currentLoc,'XData',estimatedNode(1),'YData',estimatedNode(2));
    drawnow
end
hold off

% Set lidarSLAM properties
mapResolution = 2;
maxLidarRange = 50;
loopClosureThreshold = 2450;
loopClosureRadius = 3;

% Initialize lidarSLAM object
slamAlg = lidarSLAM(mapResolution,maxLidarRange);
slamAlg.ScanRegistrationMethod = 'PhaseCorrelation';
slamAlg.LoopClosureThreshold = loopClosureThreshold;
slamAlg.LoopClosureSearchRadius = loopClosureRadius;

occGridFixed = occupancyGrid';

% Create a lidarScan object from an occupancy grid image
[rows,cols,values] = find(occGridFixed);
xLocs = -gridSize/2 + rows * gridStep;
yLocs = -gridSize/2 + cols * gridStep;

scan = lidarScan([yLocs xLocs]);

% Visualize the grid image and lidar scan
figure('Name','Occupancy image and lidar scan','NumberTitle','off')
subplot(1,2,1)
imshow(imrotate(occupancyGrid,-180))
title('Occupancy Grid Image')
subplot(1,2,2)
plot(scan)

% Plot the ground truth trajectory
figure('Name','Trajectory traversal','NumberTitle','off') 
slamAlgAxes = axes;
plot(gTNodes(:,1),gTNodes(:,2),'Color','blue','LineWidth',3,'DisplayName','Ground Truth Trajectory')
hold on 
slamAlgAxes.XLim = [-10 60];
slamAlgAxes.YLim = [-20 5];
slamAlgAxes.XGrid = 'on';
slamAlgAxes.YGrid = 'on';
title('Trajectory Traversal and Loop Closure Detections')
legend

% Plot the current location of the vehicle as the algorithm progresses.
currentLoc = scatter(gTNodes(1,1),gTNodes(1,2),50,'red','filled','DisplayName','Vehicle Location');

% Also plot any location where a loop closure was detected.
loopLocation = scatter([],[],50,'black','filled','DisplayName','Loop Detected');

for frameIdx = 1:skipFrames:numFrames
    
    ptCloud = ptCloudArr(frameIdx);
    ptCloudProcessed = helperPreprocessPtCloud(ptCloud);
    occGridMoving = helperCreateOccupancyGrid(ptCloudProcessed,gridSize,gridStep,zLimits);
    
    [rows,cols,values] = find(occGridMoving);
    xLocs = -gridSize/2 + rows * gridStep;
    yLocs = -gridSize/2 + cols * gridStep;
    
    scan = lidarScan([yLocs xLocs]);  
    
    % Add lidar scan to algorithm
    [isScanAccepted,loopClosureInfo] = addScan(slamAlg,scan); 
    
    % Update the loop closure detected location
    if isScanAccepted && size(loopClosureInfo.EdgeIDs,1) > 0
        loopLocation.XData(end+1) = gTNodes(frameIdx,1); 
        loopLocation.YData(end+1) = gTNodes(frameIdx,2);
    end
    
    % Update the vehicle location
    set(currentLoc,'XData',gTNodes(frameIdx,1),'YData',gTNodes(frameIdx,2));
    drawnow
end
hold off

% Obtain the poses from the pose graph
slamAlgPoses = slamAlg.PoseGraph.nodes();
estNodes = regEstPoseGraph.nodes();

% Plot the trajectories from the nodes
figure('Name','Vehicle Trajectory','NumberTitle',"off"); 
slamTrajAxes = axes;
slamTrajAxes.XLim = [-10 60];
slamTrajAxes.YLim = [-20 5];
slamTrajAxes.XGrid = 'on';
slamTrajAxes.YGrid = 'on';
title('Compare Trajectories')
legend

hold on; 

plot(gTNodes(:,1),gTNodes(:,2),'Color','blue','LineWidth',3,'DisplayName','Ground Truth Trajectory'); 
plot(estNodes(:,1),estNodes(:,2),'Color','green','LineWidth',3 ,'DisplayName','Estimated Trajectory');
plot(slamAlgPoses(:,1),slamAlgPoses(:,2),'Color','magenta','LineWidth',3 ,'DisplayName','Optimized Trajectory');

hold off;

slamAlg.PoseGraph.LoopClosureEdgeIDs

figure('Name','Loop Closure Results','NumberTitle',"off")
pGraphAxes = slamAlg.PoseGraph.show('IDs','off');
pGraphAxes.XLim = [0 10];
pGraphAxes.YLim = [-5 3];
title('Loop Closure Edges')

[scans,optimizedPoses]  = scansAndPoses(slamAlg);

map = buildMap(scans,optimizedPoses,1,50);
figure('Name','Occupancy Map','NumberTitle',"off") 
show(map);
hold on
show(slamAlg.PoseGraph,'IDs','off');
xlim([-20 100])
ylim([-50 50])
hold off;

sensorHeight = groundTruthPosesLidar(1).Translation(3);
ptCloudAccum = ptCloud.empty;

% Configure display
xlimits = [ -30  100];
ylimits = [-50  60];
zlimits = [-3  30];
player = pcplayer(xlimits,ylimits,zlimits);

% Define rigid transformation between the lidar sensor mounting position
% and the vehicle reference point. 
lidarToVehicleTform = helperPoseToRigidTransform(single([0 0 -1.57 0 0 0]));

% Specify vehicle dimensions
centerToFront = 1.104;
centerToRear  = 1.343;
frontOverhang = 0.828;
rearOverhang  = 0.589;
vehicleWidth  = 1.653;
vehicleHeight = 1.513;
vehicleLength = centerToFront + centerToRear + frontOverhang + rearOverhang;
hatchbackDims = vehicleDimensions(vehicleLength,vehicleWidth,vehicleHeight,...
'FrontOverhang',frontOverhang,'RearOverhang',rearOverhang);

vehicleDims   = [hatchbackDims.Length,hatchbackDims.Width,hatchbackDims.Height];
vehicleColor  = [0.85 0.325 0.098];

% Estimation trajectory handle
markerSize = 3;
hold(player.Axes,'on')
estTrajHandle   = scatter3(player.Axes,NaN,NaN,NaN,markerSize,vehicleColor,'filled');
hold(player.Axes,'off')

for frameIdx = 1:skipFrames:numFrames
    
    % Obtain the point clouds
    ptCloud = ptCloudArr(frameIdx);
    ptCloudProcessed = helperPreprocessPtCloud(ptCloud);
    
    % Obtain the pose
    poseIdx = floor((frameIdx-1)/skipFrames) + 1;
    pose = optimizedPoses(poseIdx,:);
    
    % Construct the rigid3d object required to concatenate point clouds.
    % This object holds the rigid transformation
    yaw = pose(3);
    rot = [cos(yaw) sin(yaw) 0;...
       -sin(yaw) cos(yaw) 0;...
       0 0 1];
    trans = [pose(1:2) sensorHeight];
    
    absPose = rigid3d(rot,trans);
    
    % Accumulated point cloud map
    ptCloudAccum = pccat([ptCloudAccum,pctransform(ptCloudProcessed,absPose)]);
    
    % Update accumulated point cloud map
    view(player,ptCloudAccum);
    
    % Set viewing angle to top view
    view(player.Axes,2);
    
    % Convert current absolute pose of sensor to vehicle frame 
    absVehiclePose = rigid3d( lidarToVehicleTform.T * absPose.T );
    
    % Draw vehicle at current absolute pose
    helperDrawVehicle(player.Axes,absVehiclePose,vehicleDims,'Color',vehicleColor);
    
    estTrajHandle.XData(end+1) = trans(1);
    estTrajHandle.YData(end+1) = trans(2);
    estTrajHandle.ZData(end+1) = trans(3);
    
end

hold(player.Axes,'off')

% Display the scene and the parked vehicle locations
figure;
sceneAxes = axes;
sceneName = 'LargeParkingLot';
[sceneImage, sceneRef] = helperGetSceneImage(sceneName);
helperShowSceneImage(imrotate(sceneImage, 180), sceneRef);
hold on
parkedPoses = [-17.255 -44.789 90; -0.7848, -44.6161, 90];
scatter(parkedPoses(:,1), parkedPoses(:,2), [], [0.8500 0.3250 0.0980], 'filled', 'DisplayName', 'Parked Vehicles');
ylim([-90 20])
sceneAxes.XTick = [];
sceneAxes.YTick = [];
sceneAxes.XLabel = [];
sceneAxes.YLabel = [];
legend
hold off;

function ptCloudArr = helperGetPointClouds(simOut)

% Extract signal
ptCloudData = simOut.ptCloudData.signals.values;

% Create a pointCloud array
ptCloudArr = pointCloud(ptCloudData(:,:,:,2)); % Ignore first frame
for n = 3 : size(ptCloudData,4)
    ptCloudArr(end+1) = pointCloud(ptCloudData(:,:,:,n));  %#ok<AGROW>
end

end

function gTruth = helperGetLidarGroundTruth(simOut)
numFrames = size(simOut.lidarLocation.time,1);

gTruth = repmat(rigid3d,numFrames-1,1);

for i = 2:numFrames
    gTruth(i-1).Translation = squeeze(simOut.lidarLocation.signals.values(:,:,i));
    % Ignore the roll and the pitch rotations since the ground is flat
    yaw = double(simOut.lidarOrientation.signals.values(:,3,i));
    gTruth(i-1).Rotation = [cos(yaw),sin(yaw), 0; ...
        -sin(yaw),cos(yaw),0; ...
        0,         0,    1];
end
end

function gTPoseGraph = helperGetGTPoseGraph(gTruthLidar)

gTPoseGraph = poseGraph;
gTPoseGraph.addRelativePose([0 0 0]);

for idx = 3:numel(gTruthLidar)
    relTform = gTruthLidar(idx).T * inv(gTruthLidar(idx-1).T);
    
    relTrans = relTform(4,1:2);
    relYaw = -atan2(relTform(2,1),relTform(1,1));     
    
    gTPoseGraph.addRelativePose([relTrans relYaw]);
    
end

end