diff --git a/Mapping.py b/Mapping.py
index 4af5e32..42b6c16 100644
--- a/Mapping.py
+++ b/Mapping.py
@@ -1,35 +1,125 @@
+from matplotlib import pyplot as plt
 import numpy as np
 import scipy as sp
 
-from parseData import MAX_RANGE_CLIP_DEFAULT
+from parseData import MAX_RANGE_CLIP_DEFAULT, RANGE_RESOLUTION_CART_M, convertPolarImageToCartesian
 from trajectoryPlotting import Trajectory
-import m2dp
+# import m2dp
 from getPointCloud import getPointCloudPolarInd
+from utils import getRotationMatrix
+from motionDistortion import MotionDistortionSolver
 
 # Thresholds
 ROT_THRESHOLD = 0.2  # radians
 TRANS_THRESHOLD = 2.0  # meters
 TRANS_THRESHOLD_SQ = TRANS_THRESHOLD * TRANS_THRESHOLD  # meters^2
 
+RADAR_CART_CENTER = None
+
 
 # Keyframe class
 class Keyframe():
 
-    def __init__(self, pose: np.ndarray, featurePoints: np.ndarray,
-                 radarPolarImg: np.ndarray) -> None:
+    def __init__(self, globalPose: np.ndarray, featurePointsLocal: np.ndarray,
+                 radarPolarImg: np.ndarray, velocity: np.ndarray) -> None:
         '''
         @brief Keyframe class. Contains pose, feature points and point cloud information
-        @param[in] pose (3 x 1) Pose information [x, y, th] in (m, m, rad) # TODO: Confirm these units
-        @param[in] featurePoints (K x 2) Tracked feature points from previous keyframe
-        @param[in] radarPolarImg (M x N) Radar polar (range-azimuth) image 
+        @param[in] globalPose           (3 x 1) Pose information [x, y, th] in global coordinates, 
+                                                        units of [m, m, rad] # TODO: Confirm these units
+        @param[in] featurePointsLocal   (K x 2) Tracked feature points from previous keyframe,
+                                                in local coordinates (pixels)
+        @param[in] radarPolarImg        (M x N) Radar polar (range-azimuth) image
+
+        @see updateInfo() 
+        '''
+        self.updateInfo(globalPose, featurePointsLocal, radarPolarImg, velocity)
+
+    def updateInfo(self, globalPose: np.ndarray,
+                   featurePointsLocal: np.ndarray,
+                   radarPolarImg: np.ndarray,
+                   velocity: np.ndarray) -> None:
+        '''
+        @brief Update internal information: pose, feature points and point cloud information
+        @param[in] globalPose           (3 x 1) Pose information [x, y, th] in global coordinates, 
+                                                        units of [m, m, rad] # TODO: Confirm these units
+        @param[in] featurePointsLocal   (K x 2) Tracked feature points from previous keyframe,
+                                                in local coordinates (pixels)
+        @param[in] radarPolarImg        (M x N) Radar polar (range-azimuth) image 
         '''
-        self.pose = pose
-        self.featurePoints = featurePoints  # set of (tracked) feature points
+        self.pose = globalPose
         self.radarPolarImg = radarPolarImg  # radar polar image
 
+        # Figure out and cache the center of the Cartesian image, needed for converting coordinates
+        global RADAR_CART_CENTER
+        if RADAR_CART_CENTER is None:
+            radarCartImg = convertPolarImageToCartesian(radarPolarImg)
+            RADAR_CART_CENTER = np.array(radarCartImg.shape) / 2
+
+        self.featurePointsLocal = featurePointsLocal  # set of original feature points, in local (px coordinates)
+        self.prunedFeaturePoints = self.featurePointsLocal  # set of pruned feature points (tracked until the next keyframe)
+
         # TODO: Not sure if needed/useful
         self.pointCloud = getPointCloudPolarInd(radarPolarImg)
 
+        self.velocity = velocity
+        self.featurePointsLocalUndistorted = MotionDistortionSolver.undistort(velocity, featurePointsLocal)[:, :2]
+        self.prunedUndistortedLocals = self.featurePointsLocalUndistorted
+
+    def copyFromOtherKeyframe(self, keyframe) -> None:
+        self.updateInfo(keyframe.pose, keyframe.featurePointsLocal,
+                        keyframe.radarPolarImg)
+
+    def convertFeaturesLocalToGlobal(
+            self, featurePointsLocal: np.ndarray) -> np.ndarray:
+        '''
+        @brief Local feature points to convert into global coordinates, given internal pose
+        @param[in] featurePointsLocal   (K x 2) Tracked feature points from previous keyframe,
+                                        in local coordinates (pixels)
+        @return featurePointsGlobal     (K x 2) Feature points in global coordinates, meters
+        '''
+        x, y, th = self.pose
+
+        # First translate local points to origin at center
+        featurePointsGlobal = featurePointsLocal - RADAR_CART_CENTER
+
+        # Then we need to convert to meters
+        featurePointsGlobal *= RANGE_RESOLUTION_CART_M  # px * (m/px) = m
+
+        # Center origin at pose center
+
+        # Rotate and translate to make into global coordinate system
+        R = getRotationMatrix(th)
+        t = np.array([x, y]).reshape(2, 1)
+        featurePointsGlobal = (R @ (featurePointsGlobal.T) + t).T
+
+        return featurePointsGlobal
+    
+    def getPrunedFeaturesGlobalPosition(self):
+        x, y, th = self.pose
+
+        # First translate local points to origin at center
+        featurePointsGlobal = self.prunedUndistortedLocals
+
+        # Then we need to convert to meters
+
+        # Center origin at pose center
+
+        # Rotate and translate to make into global coordinate system
+        R = getRotationMatrix(th)
+        t = np.array([x, y]).reshape(2, 1)
+        featurePointsGlobal = (R @ (featurePointsGlobal.T) + t).T
+
+        return featurePointsGlobal
+
+    def pruneFeaturePoints(self, corrStatus: np.ndarray) -> None:
+        '''
+        @brief Prune feature points based on correspondence status
+        @param[in] corrStatus 
+        @note In place changing of `self.prunedFeaturePoints` function, which aims to track and prune away the feature points that are part of this keyframe
+        '''
+        self.prunedFeaturePoints = self.prunedFeaturePoints[corrStatus.flatten().astype(bool)]
+        self.prunedUndistortedLocals = self.prunedUndistortedLocals[corrStatus.flatten().astype(bool)]
+
     # def isVisible(self, keyframe):
     #     '''
     #     @brief Return points that are visible from keyframe
@@ -51,8 +141,15 @@ def __init__(self, sequenceName: str, estTraj: Trajectory,
         self.filePaths = filePaths
 
         self.estTraj = estTraj
+
+        # TODO: Instead of storing set of all keyframes, only store the last keyframe, and then a big array of map points in global coordinates
+        self.mapPoints = [
+        ]  # should be a large np.array of global feature points
         self.keyframes = []
 
+    def updateInternalTraj(self, traj: Trajectory):
+        self.estTraj = traj
+
     # TODO: might not want to make keyframe before adding it
     def isGoodKeyframe(self, keyframe: Keyframe) -> bool:
         '''
@@ -92,6 +189,45 @@ def addKeyframe(self, keyframe: Keyframe) -> None:
     def bundleAdjustment(self) -> None:
         '''
         @brief Perform bundle adjustment on the last 2 keyframes
+        # TODO: Should actually bundle adjust on all keyframes within range
         @return None
         '''
-        pass
\ No newline at end of file
+        # Cannot do BA on only 1 KF
+        if len(self.keyframes) <= 1:
+            return
+
+        # TODO: Perform bundle adjustment on last 2 keyframes
+        old_kf = self.keyframes[-2]
+        new_kf = self.keyframes[-1]
+
+        # Obtain relevant information
+        pose_old = old_kf.pose
+        pose_new = new_kf.pose
+        points_old_local = old_kf.prunedFeaturePoints
+        points_new_local = new_kf.featurePointsLocal
+
+        # NOTE: remember ot convert to global coordinates before doing anything with the keyframes
+        points_old = old_kf.convertFeaturesLocalToGlobal(points_old_local)
+        points_new = new_kf.convertFeaturesLocalToGlobal(points_new_local)
+
+        # TODO: iterative optimisation and related solvers here
+        # 2d bundle adjustment
+
+        pass
+
+    def plot(self, fig: plt.figure, show: bool = False):
+        '''
+        @brief Plot map points on plt figure
+        @param[in] fig plt figure to plot on @todo do this
+        '''
+        # TODO: For now, plot all the keyframe feature points
+        for kf in self.keyframes:
+            points_local = kf.featurePointsLocal
+            points_global = kf.convertFeaturesLocalToGlobal(points_local)
+            
+            plt.scatter(points_global[:, 0], points_global[:, 1], marker='+', color='red',label='Map Points')
+
+        if show:
+            plt.show()
+
+        return
\ No newline at end of file
diff --git a/RawROAMSystem.py b/RawROAMSystem.py
index 8cd71db..44c3097 100644
--- a/RawROAMSystem.py
+++ b/RawROAMSystem.py
@@ -1,15 +1,20 @@
 import os
 import shutil
 from matplotlib import pyplot as plt
-
 import numpy as np
-from getFeatures import appendNewFeatures
-from parseData import convertPolarImageToCartesian, getCartImageFromImgPaths, getPolarImageFromImgPaths, getRadarImgPaths
+
+from Mapping import Keyframe, Map
+from getFeatures import N_FEATURES_BEFORE_RETRACK, appendNewFeatures
+from parseData import convertPolarImageToCartesian, getCartImageFromImgPaths, getPolarImageFromImgPaths, getRadarImgPaths, RANGE_RESOLUTION_CART_M
 from trajectoryPlotting import Trajectory, getGroundTruthTrajectory, plotGtAndEstTrajectory
 from utils import convertRandHtoDeltas, f_arr, getRotationMatrix, plt_savefig_by_axis, radarImgPathToTimestamp
 from Tracker import Tracker
+from motionDistortion import MotionDistortionSolver
+from utils import *
 
-
+# Bad solution. better solution is to save in config between mapping and this file
+RADAR_CART_CENTER = np.array([1012, 1012])
+wantToPlot = -1
 class RawROAMSystem():
 
     def __init__(self,
@@ -46,7 +51,7 @@ def __init__(self,
         self.sequenceSize = len(self.imgPathArr)
 
         # Create Save paths for imaging
-        imgSavePath = os.path.join(".", "img", "roam_rejectOutliers_1m",
+        imgSavePath = os.path.join(".", "img", "roam_mapping",
                                    sequenceName).strip(os.path.sep)
         trajSavePath = imgSavePath + '_traj'
 
@@ -73,6 +78,8 @@ def __init__(self,
                                self.filePaths, self.paramFlags)
 
         # TODO: Initialize mapping
+        self.map = Map(self.sequenceName, self.estTraj, self.imgPathArr,
+                       self.filePaths)
 
         pass
 
@@ -121,6 +128,15 @@ def run(self, startSeqInd: int = 0, endSeqInd: int = -1) -> None:
         self.gtTraj = gtTraj
         self.estTraj = estTraj
 
+        # Initialialize Motion Distortion Solver
+        # Covariance matrix, point errors
+        cov_p = np.diag([4, 4]) # sigma = 2 pixels
+        # Covariance matrix, velocity errors
+        cov_v = np.diag([1, 1, (5 * np.pi / 180) ** 2]) # 1 pixel/s, 1 pixel/s, 5 degrees/s
+        MDS = MotionDistortionSolver(cov_p, cov_v)
+        # Prior frame's pose
+        prev_pose = convertPoseToTransform(initPose)
+
         # Actually run the algorithm
         # Get initial polar and Cartesian image
         prevImgPolar = getPolarImageFromImgPaths(imgPathArr, startSeqInd)
@@ -130,32 +146,137 @@ def run(self, startSeqInd: int = 0, endSeqInd: int = -1) -> None:
         blobCoord = np.empty((0, 2))
         blobCoord, _ = appendNewFeatures(prevImgCart, blobCoord)
 
+        # Initialize first keyframe
+        metricCoord = (blobCoord - RADAR_CART_CENTER) * RANGE_RESOLUTION_CART_M
+        zero_velocity = np.zeros((3,))
+        old_kf = Keyframe(initPose, metricCoord, prevImgPolar, zero_velocity) # pointer to previous kf
+        self.map.addKeyframe(old_kf)
+
+        possible_kf = Keyframe(initPose, metricCoord, prevImgPolar, zero_velocity)
+
         for seqInd in range(startSeqInd + 1, endSeqInd + 1):
             # Obtain polar and Cart image
             currImgPolar = getPolarImageFromImgPaths(imgPathArr, seqInd)
             currImgCart = convertPolarImageToCartesian(currImgPolar)
 
             # Perform tracking
-            good_old, good_new, rotAngleRad = tracker.track(
+            # TODO: Figure out how to integrate the keyframe addition when creation of new features
+            good_old, good_new, rotAngleRad, corrStatus = tracker.track(
                 prevImgCart, currImgCart, prevImgPolar, currImgPolar,
                 blobCoord, seqInd)
+            '''
+            if seqInd == wantToPlot:
+                plt.figure()
+                plt.subplot(1, 2, 1)
+                plt.scatter(good_old[:,0], good_old[:,1])
+                plt.subplot(1, 2, 2)
+                plt.title("Good old")
+                #applied = homogenize(centered_new) @ new_transform.T
+                plt.scatter(good_new[:,0], good_new[:,1])
+                plt.title(f"Good new")
+                plt.show()
+            '''
+            # Keyframe updating
+            old_kf.pruneFeaturePoints(corrStatus)
+            
             print("Detected", np.rad2deg(rotAngleRad), "[deg] rotation")
             estR = getRotationMatrix(-rotAngleRad)
-
-            R, h = tracker.getTransform(good_old, good_new)
-
+            
+            R, h = tracker.getTransform(good_old, good_new, pixel = False)
             # R = estR
+            
+            # Solve for Motion Compensated Transform
+            p_w = old_kf.getPrunedFeaturesGlobalPosition() # centered
+
+            #TODO: Scatter p_w, then try the transform on the centered new points
+            # Scatter that on the same plot
+            centered_new = (good_new - RADAR_CART_CENTER) * RANGE_RESOLUTION_CART_M
+            # Initial Transform guess
+            T_wj = prev_pose @ np.block([[R,                h],
+                                         [np.zeros((2,)),   1]])
+
+            # Give Motion Distort info on two new frames
+            debug = False
+            if seqInd == wantToPlot:
+                debug = False
+            # Centered_new is in meters, p_w is in meters, T_wj is in meters, prev_pose is meters
+            MDS.update_problem(prev_pose, p_w, centered_new, T_wj, debug)
+            undistort_solution = MDS.optimize_library()
+
+            # Extract new info
+            pose_vector = undistort_solution[3:]
+            new_transform = convertPoseToTransform(pose_vector)
+            relative_transform = MDS.T_wj0_inv @ new_transform
+            '''
+            if seqInd == wantToPlot:
+                plt.figure()
+                plt.subplot(1, 3, 1)
+                plt.scatter(p_w[:,0], p_w[:,1])
+                plt.subplot(1, 3, 2)
+                plt.title("World coordinates")
+                applied = homogenize(centered_new) @ new_transform.T
+                plt.scatter(centered_new[:,0], centered_new[:,1])
+                plt.title(f"Post-Transform: {(np.max(centered_new[:,0]) - np.min(centered_new[:,0]))/(np.max(p_w[:,0]) - np.min(p_w[:,0]))}")
+                plt.subplot(1, 3, 3)
+                diff = p_w - applied[:, :2]
+                plt.scatter(diff[:,0], diff[:,1])
+                plt.show()
+            '''
+            R = relative_transform[:2, :2]
+            h = relative_transform[:2, 2:]
+            velocity = undistort_solution[:3]
+            #velocity[:2] /= RANGE_RESOLUTION_CART_M
 
             # Update trajectory
-            self.updateTrajectory(R, h, seqInd)
+            #self.updateTrajectory(R, h, seqInd)
+            self.updateTrajectoryAbsolute(pose_vector, seqInd)
+
+            latestPose = pose_vector #self.estTraj.poses[-1]
+            # Good new is given in pixels, given velocity in meters, uh oh, pose in meters
+            possible_kf.updateInfo(latestPose, centered_new, currImgPolar, velocity)
+
+            # Add a keyframe if it fulfills criteria
+            # 1) large enough translation from previous keyframe
+            # 2) large enough rotation from previous KF
+            # TODO: Not sure if this criteria is actually correct, perhaps we should be adding the previous keyframe instead
+            # 3) not enough features in current keyframe (ie about to have new features coming up)
+            # NOTE: Feature check and appending will only be checked in the next iteration,
+            #       so we can prematuraly do it here and add a keyframe first
+            nFeatures = good_new.shape[0]
+            retrack = (nFeatures <= N_FEATURES_BEFORE_RETRACK)
+            if retrack or \
+                self.map.isGoodKeyframe(possible_kf):
+
+                print("\nAdding keyframe...\n")
+
+                #old_kf.copyFromOtherKeyframe(possible_kf)
+                #self.map.addKeyframe(possible_kf)
+
+                # TODO: Aliasing above? old_kf is never assigned the object possible_kf,
+                # map ends up with a list of N pointers to the same keyframe
+                # Proposed fix: old_kf = possible_kf # switch ptr to new object
+                # Initialize new poss_kf for new ptr
+                old_kf = possible_kf
+                # TODO: Never replenished blobCoord. Proposed fix: Done here.
+                if retrack:
+                    good_new, _ = appendNewFeatures(currImgCart, good_new)
+                    centered_new = (good_new - RADAR_CART_CENTER) * RANGE_RESOLUTION_CART_M
+                    old_kf.updateInfo(latestPose, centered_new, currImgPolar, velocity)
+                possible_kf = Keyframe(latestPose, centered_new, currImgPolar, velocity)
+                # TODO: do bundle adjustment here
+                #self.map.bundleAdjustment()
 
             # Plotting and prints and stuff
-            self.plot(prevImgCart, currImgCart, good_old, good_new, R, h,
-                      seqInd)
-
+            if seqInd == endSeqInd:
+                self.plot(prevImgCart, currImgCart, good_old, good_new, R, h,
+                        seqInd, save = True, show = False)
+            else:
+                self.plot(prevImgCart, currImgCart, good_old, good_new, R, h,
+                        seqInd, save = False, show = False)
             # Update incremental variables
             blobCoord = good_new.copy()
             prevImgCart = currImgCart
+            prev_pose = convertPoseToTransform(latestPose)
 
     # TODO: Move into trajectory class?
     def updateTrajectory(self, R, h, seqInd):
@@ -166,6 +287,12 @@ def updateTrajectory(self, R, h, seqInd):
         self.estTraj.appendRelativeDeltas(timestamp, est_deltas)
         # self.estTraj.appendRelativeTransform(timestamp, R, h)
 
+    def updateTrajectoryAbsolute(self, pose_vector, seqInd):
+        imgPathArr = self.imgPathArr
+
+        timestamp = radarImgPathToTimestamp(imgPathArr[seqInd])
+        self.estTraj.appendAbsoluteTransform(timestamp, pose_vector)
+
     def plot(self,
              prevImg,
              currImg,
@@ -178,6 +305,7 @@ def plot(self,
              show=False):
 
         # Draw as subplots
+        plt.clf()
         self.fig.clear()
 
         ax1 = self.fig.add_subplot(1, 2, 1)
@@ -190,6 +318,9 @@ def plot(self,
                           show=False)
 
         ax2 = self.fig.add_subplot(1, 2, 2)
+        # TODO: Plotting for map points
+        #self.map.plot(self.fig, show=False)
+
         self.plotTraj(seqInd, R, h, save=False, show=False)
 
         trajSavePath = self.filePaths["trajSave"]
@@ -200,14 +331,14 @@ def plot(self,
         self.fig.savefig(trajSavePathInd)
 
         # # Save by subplot
-        # if save:
-        #     imgSavePath = self.filePaths["imgSave"]
-        #     imgSavePathInd = os.path.join(imgSavePath, f"{seqInd:04d}.jpg")
-        #     plt_savefig_by_axis(imgSavePathInd, self.fig, ax1)
+        if save:
+            imgSavePath = self.filePaths["imgSave"]
+            imgSavePathInd = os.path.join(imgSavePath, f"{seqInd:04d}.jpg")
+            plt_savefig_by_axis(imgSavePathInd, self.fig, ax1)
 
-        #     trajSavePath = self.filePaths["trajSave"]
-        #     trajSavePathInd = os.path.join(trajSavePath, f"{seqInd:04d}.jpg")
-        #     plt_savefig_by_axis(trajSavePathInd, self.fig, ax2)
+            trajSavePath = self.filePaths["trajSave"]
+            trajSavePathInd = os.path.join(trajSavePath, f"{seqInd:04d}.jpg")
+            plt_savefig_by_axis(trajSavePathInd, self.fig, ax2)
 
         if show:
             plt.pause(0.01)
@@ -285,7 +416,7 @@ def plotTraj(self, seqInd, R, h, save=False, show=False):
     print("Generating mp4 with script (requires bash and FFMPEG command)...")
     try:
         # Save video sequence
-        os.system(f"./img/mp4-from-folder.sh {imgSavePath} {startSeqInd + 1}")
+        os.system(f"./img/mp4-from-folder.sh {imgSavePath} {startSeqInd + 1} 20")
         print(f"mp4 saved to {imgSavePath.strip(os.path.sep)}.mp4")
 
         if REMOVE_OLD_RESULTS:
diff --git a/Tracker.py b/Tracker.py
index dfe1ddb..79bf79e 100644
--- a/Tracker.py
+++ b/Tracker.py
@@ -2,10 +2,11 @@
 import os
 from typing import Tuple
 from matplotlib import pyplot as plt
+from matplotlib.ft2font import BOLD
 
 import numpy as np
-from FMT import getRotationUsingFMT, getTranslationUsingPhaseCorrelation, rotateImg
-from getTransformKLT import calculateTransformDxDth, calculateTransformSVD, getTrackedPointsKLT, visualize_transform
+from FMT import getRotationUsingFMT
+from getTransformKLT import calculateTransformSVD, getTrackedPointsKLT, visualize_transform
 from outlierRejection import rejectOutliers
 from parseData import RANGE_RESOLUTION_CART_M
 from trajectoryPlotting import Trajectory
@@ -52,6 +53,7 @@ def track(
         @return good_old Coordinates of old good feature points (K' x 2) in [x, y] format
         @return good_new Coordinates of new good feature points (K' x 2) in [x, y] format
         @return angleRotRad Angle used to rotate image
+        @return corrStatus (K x 2) correspondence status @note Needed for mapping to track keyframe points
         '''
         # Timing
         start = tic()
@@ -80,16 +82,32 @@ def track(
         print(
             f"{seqInd} | Num good features: {nGoodFeatures} of {nFeatures} ({(nGoodFeatures / nFeatures) * 100:.2f}%) | Time: {toc(start):.2f}s"
         )
+        
+        # plt.subplot(1, 2, 1)
+        # plt.scatter(good_old[:,0], good_old[:,1])
+        # plt.subplot(1, 2, 2)
+        # plt.scatter(good_new[:,0], good_new[:,1])
+        # plt.show()
 
         # Outlier rejection
         doOutlierRejection = self.paramFlags.get("rejectOutliers", True)
         if doOutlierRejection:
-            good_old, good_new = rejectOutliers(good_old, good_new)
+            good_old, good_new, pruning_mask = rejectOutliers(good_old, good_new)
+        
+        # plt.subplot(1, 2, 1)
+        # plt.scatter(good_old[:,0], good_old[:,1])
+        # plt.subplot(1, 2, 2)
+        # plt.scatter(good_new[:,0], good_new[:,1])
+        # plt.show()
+        # Ensure correct correspondence status
+        rng = np.arange(nFeatures)
+        corrStatus[rng[corrStatus.flatten().astype(bool)]] &= pruning_mask[:, np.newaxis]
 
-        return good_old, good_new, angleRotRad
+        return good_old, good_new, angleRotRad, corrStatus
 
     def getTransform(self, srcCoord: np.ndarray,
-                     targetCoord: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+                     targetCoord: np.ndarray,
+                     pixel: bool) -> Tuple[np.ndarray, np.ndarray]:
         '''
         @brief Obtain transform from coordinate correspondnces
         
@@ -103,7 +121,8 @@ def getTransform(self, srcCoord: np.ndarray,
         # Obtain transforms
         # R, h = calculateTransformDxDth(srcCoord, targetCoord)
         R, h = calculateTransformSVD(srcCoord, targetCoord)
-        h *= RANGE_RESOLUTION_CART_M
+        if not pixel:
+            h *= RANGE_RESOLUTION_CART_M
 
         return R, h
 
diff --git a/testDeadReckoning.py b/archive/testDeadReckoning.py
similarity index 100%
rename from testDeadReckoning.py
rename to archive/testDeadReckoning.py
diff --git a/testTransformKLT2.py b/archive/testTransformKLT2.py
similarity index 100%
rename from testTransformKLT2.py
rename to archive/testTransformKLT2.py
diff --git a/data/full_seq_1/INFO.md b/data/full_seq_1/INFO.md
index f85f4b2..f73d0d9 100644
--- a/data/full_seq_1/INFO.md
+++ b/data/full_seq_1/INFO.md
@@ -11,7 +11,7 @@ Distance: 9.02 km
 Frames: 8867
 FPS: 4.00 Hz
 Resolution: 4.32 cm
-Range: 165 m
+Range: 165 m 
 
 ## Dataset Structure
 
@@ -23,7 +23,6 @@ Range: 165 m
          |-- (bunch of radar images labelled as <timestamp>.png)
       |-- info.txt (this file)
       |-- radar.timestamps (in form: <timestamp> <valid-bit>
-      |-- gt
-         |-- radar_odometry.csv (processed ground truth odometry data from GPS)
-
+      |-- radar_odometry.csv (processed ground truth odometry data from GPS)
+      
  ```
diff --git a/data/full_seq_1/gt/trajectory.png b/data/full_seq_1/gt/trajectory.png
deleted file mode 100644
index d1880a9..0000000
Binary files a/data/full_seq_1/gt/trajectory.png and /dev/null differ
diff --git a/g2o b/g2o
index cc40301..b1ba729 160000
--- a/g2o
+++ b/g2o
@@ -1 +1 @@
-Subproject commit cc403018ea05475380794f433666a442d675f0e2
+Subproject commit b1ba729aa569267e179fa2e237db0b3ad5169e2e
diff --git a/genFakeData.py b/genFakeData.py
index 1477d18..0f681c3 100644
--- a/genFakeData.py
+++ b/genFakeData.py
@@ -1,6 +1,6 @@
 import numpy as np
 import matplotlib.pyplot as plt
-from utils import getRotationMatrix
+from utils import getRotationMatrix, invert_transform
 from parseData import *
 
 from parseData import RANGE_RESOLUTION_CART_M
@@ -40,7 +40,7 @@ def plotFakeFeatures(srcCoord,
                     color='blue',
                     marker='.',
                     alpha=alpha,
-                    label=f'Features 0{title_append}')
+                    label=f'Instantaneous Radar Scan{title_append}')
 
     if targetCoord is not None:
         plt.scatter(targetCoord[:, 0],
@@ -48,7 +48,7 @@ def plotFakeFeatures(srcCoord,
                     color='red',
                     marker='+',
                     alpha=alpha,
-                    label=f'Features 1{title_append}')
+                    label=f'Scan with Distortion{title_append}')
 
     if targetCoord2 is not None:
         plt.scatter(targetCoord2[:, 0],
@@ -56,7 +56,7 @@ def plotFakeFeatures(srcCoord,
                     color='green',
                     marker='x',
                     alpha=alpha,
-                    label=f'Features 2{title_append}')
+                    label=f'Original Points{title_append}')
 
     if plotDisplace:
         for i in range(targetCoord.shape[0]):
@@ -116,44 +116,44 @@ def convertPolarPointsToCartesian(points):
     y = np.expand_dims(ranges * np.sin(angles), axis = 1)
     return np.hstack((x, y))
 
-def generateFakeCorrespondencesPolar(srcCoord=None,
+def generateFakeCorrespondencesPolar(currentFrame=None,
                                     n_points=100,
                                     theta_max_deg=20,
                                     max_translation_m=3):
     '''
     @brief Generate fake correspondences with transform, randomly generated from max range and degree
-    @param[in] srcCoord Source coordinate to transform from. If none, will randomly generate features
-    @param[in] n_points Number of points to generate, only applies if srcCoord = None
+    @param[in] currentFrame Source coordinate to transform from. If none, will randomly generate features
+    @param[in] n_points Number of points to generate, only applies if currentFrame = None
     @param[in] theta_max_deg Maximum degree of rotation
     @param[in] max_range_m Maximum range (for translation) in meters
 
-    @return srcCoord Generated or passed in srcCoord
+    @return currentFrame Generated or passed in currentFrame
     @return targetCoord Corresponding targetCoord generated using (theta_deg, h)
     @return theta_deg Theta component of transform
     @return h Translation component of transform
     '''
 
-    if srcCoord is None:
+    if currentFrame is None:
         print("Generating fake features..")
         max_range_m = max_translation_m * 3
-        srcCoord = generateFakeFeaturesPolar(n_points, max_range_m)
-        #print(srcCoord.shape)
-        srcCoord = convertPolarPointsToCartesian(srcCoord)
+        currentFrame = generateFakeFeaturesPolar(n_points, max_range_m)
+        #print(currentFrame.shape)
+        currentFrame = convertPolarPointsToCartesian(currentFrame)
     else:
-        n_points = srcCoord.shape[0]
+        n_points = currentFrame.shape[0]
 
     theta_deg = np.random.random() * theta_max_deg
     R = getRotationMatrix(theta_deg, degrees=True)
     h = generateTranslationVector(max_translation_m)
-    #print(srcCoord.shape)
-    targetCoord = transformCoords(srcCoord, R, h)
+    #h = np.array([[0], [0]])
+    groundTruth = transformCoords(currentFrame, R, h)
 
-    return srcCoord, targetCoord, theta_deg, h
+    return groundTruth, currentFrame, theta_deg, h
 
 def distort(coords, velocity, frequency, h):
     
     coords_norm = coords - h.flatten() # N x 2
-    angles = np.arctan2(coords_norm[:, 1], -coords_norm[:, 0]) # - y to follow clockwise convention
+    angles = np.arctan2(coords_norm[:, 1], -coords_norm[:, 0]) # - x to follow clockwise convention
     period = 1 / frequency
     times = angles / (2 * np.pi) * period
     #print(angles) # lesson: need to use arctan2 wisely, it wraps [-pi, pi]
@@ -166,15 +166,14 @@ def distort(coords, velocity, frequency, h):
     #print(displacement)
     #print(displacement)
     dx = displacement[0, :]
-    print(dx)
     dy = displacement[1, :]
     dtheta = displacement[2, :] / 180 * np.pi
     c = np.cos(dtheta)
     s = np.sin(dtheta)
     ones = np.ones(times.shape)
     zeros = np.zeros(times.shape)
-    distortion = np.array([[c, s, -s* dy - c*dx],
-                           [-s,  c, -c*dy + s*dx],
+    distortion = np.array([[ c, s, -s*dy - c*dx],
+                           [-s, c, -c*dy + s*dx],
                            [zeros, zeros, ones]]) # 3 x 3 x N, need to invert?
     distorted = np.transpose(distortion, axes = (2, 0, 1)) @ np.expand_dims(coords, axis = 2) # N x 3 x 1
     distorted = distorted[:, :2, 0]
diff --git a/getFeatures.py b/getFeatures.py
index 2564f16..9fe1e34 100644
--- a/getFeatures.py
+++ b/getFeatures.py
@@ -54,7 +54,7 @@ def getBlobsFromCart(cartImage: np.ndarray,
 
 # Thresholds for feature loss
 PERCENT_FEATURE_LOSS_THRESHOLD = 0.75
-N_FEATURES_BEFORE_RETRACK = -1  # TODO: Make it dynamic (find the overall loss)
+N_FEATURES_BEFORE_RETRACK = 60  # TODO: Make it dynamic (find the overall loss)
 
 
 # TODO: Make dynamic?
@@ -103,7 +103,6 @@ def appendNewFeatures(srcImg, oldFeaturesCoord):
     @param[in] srcImg Source image to obtain features on
     @param[in] oldFeaturesCoord (K x 2) array of [x, y] coordinate of features
     '''
-
     newFeatureCoord, newFeatureRadii = getFeatures(srcImg)
     print("Added", newFeatureCoord.shape[0], "new features!")
 
diff --git a/getTransformKLT.py b/getTransformKLT.py
index af73c65..ce6b7e7 100644
--- a/getTransformKLT.py
+++ b/getTransformKLT.py
@@ -15,7 +15,7 @@
 from utils import *
 
 PLOT_BAD_FEATURES = False
-N_FEATURES_BEFORE_RETRACK = 80
+N_FEATURES_BEFORE_RETRACK = 60
 
 def visualize_transform(prevImg: np.ndarray,
                         currImg: np.ndarray,
@@ -61,7 +61,7 @@ def visualize_transform(prevImg: np.ndarray,
                     marker='.',
                     color='yellow',
                     alpha=alpha,
-                    label=f'Image 0 Features{extraLabel}')
+                    label=f'Previous Features{extraLabel}')
 
     plt.legend()
     plt.axis("off")
@@ -349,6 +349,7 @@ def getTrackedPointsKLT(
     if nFeatures < N_FEATURES_BEFORE_RETRACK:
         featurePtSrc, N_FEATURES_BEFORE_RETRACK = \
             appendNewFeatures(srcImg, featurePtSrc)
+        print("WARNING: getTransformKLT added new features!")
 
     # Perform KLT to get corresponding points
     # Stupid conversions to appropriate types
@@ -468,15 +469,16 @@ def getTrackedPointsKLT(
             good_old, good_new = rejectOutliers(good_old, good_new)
 
             # Obtain transforms
-            R, h = calculateTransformDxDth(good_old, good_new)
-            # R, h = calculateTransformSVD(good_old, good_new)
+            #R, h = calculateTransformDxDth(good_old, good_new)
+            R, h = calculateTransformSVD(good_old, good_new)
             # print(h)
             # h[0] += 0
             # for i in range(good_old.shape[0]):
             #     plotFakeFeatures(good_old[i:i+1,:], (R @ good_new[i:i+1,:].T + h).T, show= True)
-            # transformed_pts = (R @ good_new.T + h).T
+            transformed_pts = (R @ good_new.T + h).T
             # print(f"RMSE = {np.sum(np.square(good_old - transformed_pts))}")
-            # plotFakeFeatures(good_old, transformed_pts, good_new, show = True)
+            #plotFakeFeatures(good_old, good_new, show = True)
+            plotFakeFeatures(good_old, transformed_pts, good_new, show = True)
             h *= RANGE_RESOLUTION_CART_M
 
             #R, h = estimateTransformUsingDelats(good_old, good_new)
diff --git a/img/render_gif.py b/img/render_gif.py
index 2d6533b..d641670 100644
--- a/img/render_gif.py
+++ b/img/render_gif.py
@@ -5,8 +5,8 @@
 # truly a windows moment
 images = []
 for fname in sorted(
-    glob.glob("track_klt_thresholding/full_seq_1_traj/*.jpg"), key=lambda x: int(x.split("\\")[1].split(".")[0])
+    glob.glob("roam/full_seq_1_traj/*.jpg"), key=lambda x: int(x.split("\\")[1].split(".")[0])
 ):
     print(fname)
     images.append(imageio.imread(fname))
-imageio.mimsave("track_klt_thresholding/full_seq_1_traj_1.gif", images, fps=10)
\ No newline at end of file
+imageio.mimsave("roam/full_seq_1_traj_1.gif", images, fps=10)
\ No newline at end of file
diff --git a/img/roam_mapping/tiny_traj.mp4 b/img/roam_mapping/tiny_traj.mp4
new file mode 100644
index 0000000..a2d3de9
Binary files /dev/null and b/img/roam_mapping/tiny_traj.mp4 differ
diff --git a/img/roam_mapping/tiny_traj/0001.jpg b/img/roam_mapping/tiny_traj/0001.jpg
new file mode 100644
index 0000000..b2a5ce8
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0001.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0002.jpg b/img/roam_mapping/tiny_traj/0002.jpg
new file mode 100644
index 0000000..37676ec
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0002.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0003.jpg b/img/roam_mapping/tiny_traj/0003.jpg
new file mode 100644
index 0000000..2237332
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0003.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0004.jpg b/img/roam_mapping/tiny_traj/0004.jpg
new file mode 100644
index 0000000..dbb22df
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0004.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0005.jpg b/img/roam_mapping/tiny_traj/0005.jpg
new file mode 100644
index 0000000..ffeffc0
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0005.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0006.jpg b/img/roam_mapping/tiny_traj/0006.jpg
new file mode 100644
index 0000000..3a758e6
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0006.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0007.jpg b/img/roam_mapping/tiny_traj/0007.jpg
new file mode 100644
index 0000000..47b0af0
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0007.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0008.jpg b/img/roam_mapping/tiny_traj/0008.jpg
new file mode 100644
index 0000000..99e54b8
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0008.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0009.jpg b/img/roam_mapping/tiny_traj/0009.jpg
new file mode 100644
index 0000000..ff59f5f
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0009.jpg differ
diff --git a/img/roam_mapping/tiny_traj/0010.jpg b/img/roam_mapping/tiny_traj/0010.jpg
new file mode 100644
index 0000000..1974303
Binary files /dev/null and b/img/roam_mapping/tiny_traj/0010.jpg differ
diff --git a/motionDistortion.py b/motionDistortion.py
index b668b07..7f1dbba 100644
--- a/motionDistortion.py
+++ b/motionDistortion.py
@@ -1,5 +1,7 @@
 import numpy as np
 import scipy as sp
+from utils import *
+import matplotlib.pyplot as plt
 #sp.linalg.sqrtm
 #sp.lin
 
@@ -21,6 +23,8 @@
 TODO: Run parse data to look at how the scanline moves. Adjust code to match
 TODO: Verify code against math
 TODO: Integration with the code. Adding mapping capability
+TODO: Check arctan2 conventions for dT stuff. the image might be flipped because
+images have the origin in the top left.
 
 Plan:
 Check out parse data
@@ -29,35 +33,79 @@
 Debug
 
 '''
+RADAR_SCAN_FREQUENCY = 4 # 4 hz data
+VERBOSE = False
 class MotionDistortionSolver():
-    def __init__(self, T_wj0, p_w, p_jt, v_j0, T_wj, lambda_p, lambda_v):
+    def __init__(self, T_wj0, p_w, p_jt, T_wj, sigma_p, sigma_v, 
+                 frequency = RADAR_SCAN_FREQUENCY):
         # e_p Parameters
         self.T_wj0 = T_wj0 # Prior transform, T_w,j0
         self.T_wj0_inv = np.linalg.inv(T_wj0)
-        self.p_w = p_w # Estimated point world positions, N x 3
-        self.p_jt = p_jt # Observed points at time t, N x 3
+        self.p_w = homogenize(p_w) # Estimated point world positions, N x 3
+        self.p_jt = homogenize(p_jt) # Observed points at time t, N x 3
+        self.T_wj_initial = T_wj
+
+        # Radar Data Params
+        self.total_scan_time = 1 / frequency
 
         # e_v Parameters
-        self.v_j_initial = v_j0 # Initial velocity guess (prior velocity/ velocity from SVD solution)
-        self.T_wj_initial = T_wj # Initial Transform guess (T from SVD solution)
+        self.v_j_initial = self.infer_velocity(self.T_wj0_inv @ T_wj)
+        # Initial velocity guess (prior velocity/ velocity from SVD solution)
 
         # Optimization parameters
-        self.lambda_p = lambda_p # Info matrix, point error, lamdba_p
-        self.lambda_v = lambda_v # Info matrix, velocity, lambda_v
-        nv = self.lambda_p.shape[0]
-        np = self.lambda_v.shape[0]
-        self.lambda_total = np.block([[lambda_v, np.zeros((nv, np))],
-                                      [np.zeros((np, nv)), lambda_p]])
-        self.info_sqrt = sp.linalg.sqrtm(np.linalg.inv(self.lambda_total)) # 5 x 5
-        self.dT = None
-        self.T_wj_best = T_wj
-        self.v_j_best = v_j0 # might not be good enough a guess, too far from optimal
+        self.sigma_p = sigma_p # Info matrix, point error, lamdba_p
+        self.sigma_v = sigma_v # Info matrix, velocity, sigma_v
+        # n_v = self.sigma_v.shape[0]
+        # n_p = self.sigma_p.shape[0]
+        # self.sigma_total = np.block([[sigma_v, np.zeros((n_v, n_p))],
+        #                               [np.zeros((n_p, n_v)), sigma_p]])
+        # self.info_sqrt = sp.linalg.sqrtm(np.linalg.inv(self.sigma_total))
+        sigma_p_vector = np.tile(np.diag(sigma_p), p_jt.shape[0])
+        sigma_v_vector = np.diag(sigma_v)
+        sigma_vector = np.concatenate((sigma_p_vector, sigma_v_vector))
+        self.info_vector = 1 / sigma_vector
+        self.dT = MotionDistortionSolver.compute_time_deltas(self.total_scan_time, p_jt)
+        pass
 
+    def __init__(self, sigma_p, sigma_v, 
+                 frequency = RADAR_SCAN_FREQUENCY):
         # Radar Data Params
-        self.total_scan_time = 1 / 4 # assuming 4 Hz
+        self.total_scan_time = 1 / frequency
+
+        # Optimization parameters
+        self.sigma_p = np.diag(sigma_p) # Info matrix, point error, lamdba_p
+        self.sigma_v = np.diag(sigma_v) # Info matrix, velocity, sigma_v
         pass
+    
+    def update_problem(self, T_wj0, p_w, p_jt, T_wj, debug= False):
+        # e_p Parameters
+        self.T_wj0 = T_wj0 # Prior transform, T_w,j0
+        self.T_wj0_inv = np.linalg.inv(T_wj0)
+        self.p_w = homogenize(p_w) # Estimated point world positions, N x 3
+        self.p_jt = homogenize(p_jt) # Observed points at time t, N x 3
+        assert(p_w.shape == p_jt.shape)
+        self.T_wj_initial = T_wj
+        self.debug = debug
 
-    def compute_time_deltas(self):
+        # e_v Parameters
+        self.v_j_initial = self.infer_velocity(self.T_wj0_inv @ T_wj)
+        # Initial velocity guess (prior velocity/ velocity from SVD solution)
+        self.dT = MotionDistortionSolver.compute_time_deltas(self.total_scan_time, p_jt)
+
+        # Info matrix scales to number of points in the optimization problem
+        sigma_p_vector = np.tile(self.sigma_p, p_jt.shape[0])
+        sigma_v_vector = self.sigma_v
+        sigma_vector = np.concatenate((sigma_p_vector, sigma_v_vector))
+        self.info_vector = 1 / sigma_vector
+
+    def infer_velocity(self, transform):
+        dx = transform[0, 2]
+        dy = transform[1, 2]
+        dtheta = np.arctan2(transform[1, 0], transform[0, 0])
+        return np.array([dx, dy, dtheta]) / self.total_scan_time
+
+    @staticmethod
+    def compute_time_deltas(period, points):
         '''
         Get the time deltas for each point. This depends solely on where the
         points are in scan angle. The further away from center, the greater the
@@ -67,33 +115,42 @@ def compute_time_deltas(self):
         idea to re-run this function once an undistorted frame is obtained for
         better estimates.
         '''
-        points = self.undistort()#self.p_jt # provide in N x 3
         x = points[:, 0]
         y = points[:, 1]
-        angles = np.arctan2(-y, x) # scanline starts at positive x axis and moves clockwise
-        dT = self.total_scan_time * angles / (2 * np.pi)
-        dT -= self.total_scan_time / 2 # offset range, as defined in [-scan_time /2 , scan_time/2]
-        self.dT = dT
-
-    def undistort(self, v_j):
+        # scanline starts at positive x axis and moves clockwise (counter-clockwise?)
+        # We take midpoint pi/2 as 0
+        angles = np.arctan2(-y, -x) 
+        dT = period * angles / (2 * np.pi)
+        return dT
+    
+    @staticmethod
+    def undistort(v_j, points, period = 1 / RADAR_SCAN_FREQUENCY, times = None):
         '''
         Computes a new set of undistorted observed points, based on the current
         best estimate of v_T, T_wj, dT
         '''
-        displacement = np.expand_dims(v_j, axis = 1) * self.dT # 3 x N
-        assert(displacement.shape = (3,points.shape[0]))
-        theta = displacement[2, :]
-        dx = displacement[0, :]
-        dy = displacement[1, :]
+        # Turn points in homogeneous form if not already
+        points = homogenize(points)
+
+        # Get the time deltas for motion distortion 
+        if times is None:
+            assert(period > 0)
+            times = MotionDistortionSolver.compute_time_deltas(period, points)
+        
+        v_j_column = np.expand_dims(v_j, axis = 1)
+        displacement = v_j_column * times # 3 x N
+
+        theta = displacement[2, :] # (N,)
+        dx = displacement[0, :] # (N,)
+        dy = displacement[1, :] # (N,)
+        shape = theta.shape
         # Correction matrix for time drift, 3 x 3 x N
         T_j_jt = np.array([[np.cos(theta), -np.sin(theta), dx],
                            [np.sin(theta), np.cos(theta),  dy],
-                           [0,             0,              1]])
-        
-        p_jt_col = np.expand_dims(self.p_jt, axis = 1).transpose(axis = (2, 0, 1)) # N x 3 x 1
-        undistorted = T_j_jt.transpose(axis = (2, 0, 1)) @ p_jt_col # N x 3 x 1
-        return undistorted
-
+                           [np.zeros(shape), np.zeros(shape), np.ones(shape)]])
+        p_jt_col = np.expand_dims(points, axis = 2) # N x 3 x 1
+        undistorted = T_j_jt.transpose((2, 0, 1)) @ p_jt_col # N x 3 x 1
+        return np.squeeze(undistorted) # N x 3
 
     def expected_observed_pts(self, T_wj):
         '''
@@ -103,13 +160,19 @@ def expected_observed_pts(self, T_wj):
         return np.linalg.inv(T_wj) @ self.p_w.T
 
     def error_vector(self, params):
-        theta = params[2]
-        x = params[0]
-        y = params[1]
+        '''
+        Because we are optimizing over rotations, we choose to keep the rotation
+        in a theta form, we have to do matrix exponential in here to convert
+        into the SO(1) form, then augment to the rotation-translation transform
+        '''
+        theta = params[5]
+        x = params[3]
+        y = params[4]
         T = np.array([[np.cos(theta), -np.sin(theta), x],
                       [np.sin(theta),  np.cos(theta), y],
                       [0            ,  0            , 1]])
-        return self.info_sqrt @ self.error(params[:3], T)
+        #return self.info_sqrt @ self.error(params[:3], T)
+        return self.info_vector * self.error(params[:3], T)
 
     def error(self, v_j, T_wj):
         '''
@@ -117,33 +180,41 @@ def error(self, v_j, T_wj):
         estimated observed positions and the velocity error.
         '''
         # Compute point error
-        undistorted = self.undistort(v_j)
-        expected = self.expected_observed_pts(self, T_wj)
-        naive_e_p = expected - np.squeeze(undistorted).T # 3 x N
+        undistorted = MotionDistortionSolver.undistort(v_j, self.p_jt, times=self.dT)
+        expected = self.expected_observed_pts(T_wj)
+        naive_e_p = expected - undistorted.T # 3 x N
+
         # Actual loss is the Cauchy robust loss, defined here:
-        e_p_i = np.log(np.square(naive_e_p[:2, :]) / 2 + 1)
-        e_p = np.sum(e_p_i, axis = 1) # 2 x 1
+        e_p_i = np.log(np.square(naive_e_p[:2, :]) / 2 + 1) # 2 x N
+        e_p = e_p_i.flatten(order='F')
+        #e_p = np.sum(e_p_i, axis = 1) # (2,)
         
         # Compute velocity error
+        # Matrix log operation
         T_j_j1 = self.T_wj0_inv @ T_wj
         dx = T_j_j1[0, 2]
         dy = T_j_j1[1, 2]
         dtheta = np.arctan2(T_j_j1[1, 0], T_j_j1[0, 0])
         v_j_prior = np.array([dx, dy, dtheta]) / self.total_scan_time
-        e_v = (v_j - v_j_prior) * e_p.shape[1] # 3 x 1
+        v_diff = (v_j - v_j_prior)
+        v_diff[2] = normalize_angles(v_diff[2])
+        e_v = v_diff * e_p_i.shape[1] # (3,)
+        # ideally should warp2pi here on theta error
+        e = np.hstack((e_p, e_v))
 
-        e = np.vstack((e_v, e_p))
         return e
 
     def jacobian_vector(self, params):
-        theta = params[2]
-        x = params[0]
-        y = params[1]
+        theta = params[5]
+        x = params[3]
+        y = params[4]
         T = np.array([[np.cos(theta), -np.sin(theta), x],
                       [np.sin(theta),  np.cos(theta), y],
                       [0            ,  0            , 1]])
-        return self.info_sqrt @ self.jacobian(params[:3], T)
-
+        velocity = params[:3]
+        #return self.info_sqrt @ self.jacobian(velocity, T)
+        return np.expand_dims(self.info_vector, axis=1) * self.jacobian(velocity, T)
+        
     def jacobian(self, v_j, T_wj):
         '''
         Compute the Jacobian. This has two parts, as defined by the RadarSLAM
@@ -153,12 +224,13 @@ def jacobian(self, v_j, T_wj):
         J_v -   gradient of point error and velocity error wrt velocity terms
                 vx, vy, vtheta
         '''
-        undistorted = self.undistort(v_j)
-        expected = self.expected_observed_pts(self, T_wj)
-        input = expected - np.squeeze(undistorted).T # 3 x N
-        cauchy_derivative = input / (np.square(input[:2, ]) / 2 + 1) # 2 x N
+        undistorted = MotionDistortionSolver.undistort(v_j, self.p_jt, times=self.dT)
+        expected = self.expected_observed_pts(T_wj)
+        input = expected - undistorted.T # 3 x N
+        naive_e_p = input[:2]
+        cauchy_derivative = naive_e_p / (np.square(naive_e_p) / 2 + 1) # 3 x N
 
-        # Compute J_p: derivative of e_p wrt
+        # Compute J_p: derivative of errors wrt the point position
         c0 = self.T_wj0[0, 0]
         s0 = self.T_wj0[1, 0]
         c1 = T_wj[0, 0]
@@ -168,14 +240,16 @@ def jacobian(self, v_j, T_wj):
         pwx = self.p_w[:, 0] # 1 x N
         pwy = self.p_w[:, 1]
         ones = np.ones(pwx.shape)
+        zeros = np.zeros(pwx.shape)
 
         # 2 x 3 x N
         J_p1 = np.array([[-c1 * ones, -s1 * ones, -pwx * s1 + pwy * c1 - c1 * Ty + s1 * Tx],
                         [s1 * ones,  -c1 * ones, -pwx * c1 - pwy * s1 + s1 * Ty + c1 * Tx]])
         J_p1 *= np.expand_dims(cauchy_derivative, axis = 1)
+        J_p1 = np.squeeze(np.vstack(np.split(J_p1, J_p1.shape[2], axis = 2)))
         J_p2 = np.array([[ c0, s0, 0],
                          [-s0, c0, 0],
-                         [0,   0,  1]]) / self.total_scan_time
+                         [0,   0,  1]]) / self.total_scan_time * pwx.shape[0]
         J_p = np.vstack((J_p1, J_p2))
 
         # Compute J_v: derivative of the errors wrt velocity
@@ -184,10 +258,11 @@ def jacobian(self, v_j, T_wj):
         y = points[:, 1]
         displacement = np.expand_dims(v_j, axis = 1) * self.dT # 3 x N
         theta = displacement[2, :]
-        J_v = np.array([[-self.dT, 0, np.sin(theta) * self.dT * x + np.cos(theta) * self.dT * y ],
-                        [0, -self.dT, -np.cos(theta) * self.dT * x + np.sin(theta) * self.dT * y]])
+        zeros = np.zeros(theta.shape)
+        J_v = np.array([[-self.dT, zeros, self.dT * (np.sin(theta) * x + np.cos(theta) * y) ],
+                        [zeros, -self.dT, self.dT * (-np.cos(theta) * x + np.sin(theta) * y)]])
         J_v *= np.expand_dims(cauchy_derivative, axis = 1)
-        J_v = np.sum(J_v, axis = -1) # 3 x 2
+        J_v = np.squeeze(np.vstack(np.split(J_v, J_v.shape[2], axis = 2))) # 2N x 3
         J_v = np.vstack((J_v, np.eye(3) * x.shape[0]))
         
         # J = [J_v, J_p]
@@ -218,13 +293,19 @@ def optimize(self, max_iters = 20):
         pass
 
     def optimize_library(self):
+        '''
+        Optimize using the LM implementation in the scipy library.
+        '''
         # Initialize v, T
         T0 = self.T_wj_initial
         T_params = np.array([T0[0, 2], T0[1, 2], np.arctan2(T0[1, 0], T0[0, 0])])
         initial_guess = np.hstack((self.v_j_initial, T_params))
+        if VERBOSE:
+            print(f"Initial v guess: {self.v_j_initial}")
+            print(f"Initial T guess: {T_params}")
 
-        result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = self.jacobian_vector, method = 'lm')
-
+        result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = '2-point', method = 'lm')
+        # result = sp.optimize.least_squares(self.error_vector, initial_guess, jac = self.jacobian_vector, method = 'lm')
         # return v, T
         best_params = result.x
         num_evals = result.nfev # number of function evaluations: measure of how quickly we converged
@@ -235,7 +316,10 @@ def optimize_library(self):
                         2 : "ftol termination condition is satisfied",
                         3 : "xtol termination condition is satisfied",
                         4 : "Both ftol and xtol termination conditions are satisfied"}
-        print(f"Final v: {best_params[:3]}, t: {best_params[3:]}")
-        print(f"Used {num_evals} evaluations")
-        print(status_dict[status])
+        if VERBOSE:
+            print(f"Final v: {best_params[:3]}")
+            print(f"Final t: {best_params[3:]}")
+            print(f"Used {num_evals} evaluations")
+            print(f"Residuals were {result.fun}")
+            print(status_dict[status])
         return best_params
diff --git a/outlierRejection.py b/outlierRejection.py
index 45722d1..e43f7e4 100644
--- a/outlierRejection.py
+++ b/outlierRejection.py
@@ -7,7 +7,7 @@
 from genFakeData import addNoise, createOutliers, generateFakeCorrespondences, plotFakeFeatures
 
 # TODO: Tune this
-DIST_THRESHOLD_M = 1  # changed from 2.5m
+DIST_THRESHOLD_M = 0.5  # changed from 2.5m
 DIST_THRESHOLD_PX = DIST_THRESHOLD_M / RANGE_RESOLUTION_CART_M  # Euclidean distance threshold
 # NOTE: this is Euclid distance squared (i.e. 25 = ~5 px of error allowed)
 DISTSQ_THRESHOLD_PX = DIST_THRESHOLD_PX * DIST_THRESHOLD_PX
@@ -15,9 +15,8 @@
 
 def rejectOutliers(
     prev_coord: np.ndarray,
-    new_coord: np.ndarray,
-    #    outlierInd=None
-) -> tuple[np.ndarray, np.ndarray]:
+    new_coord: np.ndarray
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     '''
     @brief Reject outliers by using radar geometry to find dynamic/moving features
     
@@ -37,6 +36,7 @@ def rejectOutliers(
 
     @return pruned_prev_coord   (k x 2) Pruned previous coordinates
     @return pruned_new_coord    (k x 2) Pruned current/new coordinates
+    @return pruning_mask        (K x 2) Pruning mask
     '''
     assert prev_coord.shape == new_coord.shape, "Coordinates should be the same shape"
 
@@ -56,7 +56,7 @@ def rejectOutliers(
 
     distDiff = np.abs(dist_prev - dist_new)
     distThreshMask = (distDiff <= DIST_THRESHOLD_PX).astype(np.int8)
-    
+
     print("Mean Distance Error:", distDiff.mean(), "[px] | Threshold:", DIST_THRESHOLD_PX, "[px]")
 
     # Form graph using the distThreshMask matrix
@@ -92,7 +92,7 @@ def rejectOutliers(
     pruned_prev_coord = prev_coord[pruning_mask]
     pruned_new_coord = new_coord[pruning_mask]
 
-    return pruned_prev_coord, pruned_new_coord
+    return pruned_prev_coord, pruned_new_coord, pruning_mask
 
 
 if __name__ == "__main__":
diff --git a/parseData.py b/parseData.py
index 864ffc0..984221c 100644
--- a/parseData.py
+++ b/parseData.py
@@ -1,5 +1,5 @@
 from tkinter.messagebox import NO
-from typing import Tuple
+from typing import Tuple, List
 import numpy as np
 import cv2
 import os, sys
@@ -55,7 +55,7 @@ def extractDataFromRadarImage(
 
 def drawCVPoint(img: np.ndarray,
                 point: CartCoord,
-                point_color: tuple[int, int, int] = (0, 0, 255)):
+                point_color: Tuple[int, int, int] = (0, 0, 255)):
     if isinstance(point, CartCoord):
         point = point.asTuple()
 
@@ -132,7 +132,6 @@ def convertPolarImageToCartesian(
         flags += cv2.WARP_POLAR_LOG
 
     imgCart = cv2.warpPolar(imgPolar, cartSize, center, maxRadius, flags)
-
     return imgCart
 
 
@@ -159,7 +158,7 @@ def convertPolarImgToLogPolar(imgPolar: np.ndarray):
 
 
 def getDataFromImgPathsByIndex(
-    imgPathArr: list[str], index: int
+    imgPathArr: List[str], index: int
 ) -> Tuple[np.ndarray, float, float, np.ndarray, np.ndarray, np.ndarray]:
     '''
     @brief Get information from image path array, indexing accordingly
@@ -180,7 +179,7 @@ def getDataFromImgPathsByIndex(
     return extractDataFromRadarImage(imgPolarData)
 
 
-def getPolarImageFromImgPaths(imgPathArr: list[str], index: int) -> np.ndarray:
+def getPolarImageFromImgPaths(imgPathArr: List[str], index: int) -> np.ndarray:
     '''
     @brief Get polar image from image path array, indexing accordingly
     @param[in] imgPathArr List of image path as strings
@@ -193,7 +192,7 @@ def getPolarImageFromImgPaths(imgPathArr: list[str], index: int) -> np.ndarray:
     return imgPolar
 
 
-def getCartImageFromImgPaths(imgPathArr: list[str], index: int) -> np.ndarray:
+def getCartImageFromImgPaths(imgPathArr: List[str], index: int) -> np.ndarray:
     '''
     @brief Get polar image from image path array, indexing accordingly
     @param[in] imgPathArr List of image path as strings
@@ -206,7 +205,7 @@ def getCartImageFromImgPaths(imgPathArr: list[str], index: int) -> np.ndarray:
     return convertPolarImageToCartesian(imgPolar)
 
 
-def getRadarImgPaths(dataPath: str, timestampPath: str) -> list[str]:
+def getRadarImgPaths(dataPath: str, timestampPath: str) -> List[str]:
     '''
     @brief Obtain list of radar image paths
     
diff --git a/testMotionDistortion.py b/testMotionDistortion.py
new file mode 100644
index 0000000..5840d2d
--- /dev/null
+++ b/testMotionDistortion.py
@@ -0,0 +1,100 @@
+from getTransformKLT import *
+from genFakeData import *
+from motionDistortion import *
+import matplotlib.pyplot as plt
+
+if __name__ == "__main__":
+    N = 100
+    outlier_rate = 0.4
+    noisy = False
+    useOld = False
+    frequency = 4
+    period = 1 / frequency
+
+    # Generate fake data
+    # srcCoord = frame 1
+    # currentFrame = frame 2
+    # R_theta_deg @ currentFrame + h = srcCoord
+    # pose at frame 1 (srcCoord) is I
+    groundTruth, currentFrame, theta_deg, h = generateFakeCorrespondencesPolar(n_points=N)
+    # This theta_deg reflects the radar's own motion. To distort points, the opposite must be used
+    velocity = np.array([h[0, 0], h[1, 0], theta_deg]) / period
+    distorted = distort(currentFrame, velocity, frequency, h)
+    if noisy:
+        currentFrame, outlier_ind = createOutliers(currentFrame, int(N * outlier_rate), 
+                                    noiseToAdd=10)
+
+    '''
+    Naive Fit: rotation and translation
+    '''
+    if useOld:
+        R_fit, h_fit = calculateTransform(groundTruth, distorted)
+
+        A = np.block([[R_fit, h_fit],
+                    [np.zeros((1, 2)), 1]])
+        A_inv = np.linalg.inv(A)
+        R_fit = A_inv[:2, :2]
+        h_fit = A_inv[:2, 2:]
+    else:
+        R_fit, h_fit = calculateTransformSVD(groundTruth, distorted)
+    #R_fit, h_fit = calculateTransformDxDth(srcCoord, currentFrame)
+
+    theta_fit = np.arctan2(R_fit[1, 0], R_fit[0, 0]) * 180 / np.pi
+    print(f"Actual Transform:\ntheta:\n{theta_deg}\nh:\n{h}")
+    print(f"Fitted Transform:\ntheta:\n{theta_fit}\nh:\n{h_fit}")
+
+    # Visualize
+    plt.subplot(1,2,1)
+    srcCoord2 = (R_fit @ distorted.T + h_fit).T    
+    plotFakeFeatures(groundTruth, srcCoord2, 
+                     title_append="", alpha=0.5, clear=False, show=False, 
+                     plotDisplace = True)
+
+    '''
+    Applying Motion Distortion Solving
+    '''
+    # Prior frame's pose
+    T_wj0 = np.eye(3)
+    # Point world positions
+    p_w = groundTruth
+    # Observed points at time 1
+    p_jt = distorted
+    # Initial velocity guess
+    v_j0 = np.array([h_fit[0,0], h_fit[1,0], theta_fit * np.pi / 180]) / period
+    # Initial Transform guess
+    T_wj = np.block([[R_fit,            h_fit],
+                     [np.zeros((2,)),    1]])
+    # Covariance matrix, point errors
+    cov_p = np.diag([4, 4])
+    # Covariance matrix, velocity errors
+    cov_v = np.diag([1, 1, (5 * np.pi / 180) ** 2]) # 4 ^2 since 4 Hz
+    # Information matrix, 
+    MDS = MotionDistortionSolver(T_wj0, p_w, p_jt, v_j0, T_wj, cov_p, cov_v)
+    MDS.compute_time_deltas(p_jt)
+    undistorted = MDS.undistort(v_j0)
+    print(undistorted.shape)
+    undistorted = undistorted[:, :2, 0]
+    R_fit, h_fit = calculateTransformSVD(groundTruth, undistorted)
+    srcCoord3 = (R_fit @ undistorted.T + h_fit).T 
+    plt.subplot(1,2,2)
+    plotFakeFeatures(groundTruth, srcCoord3, 
+                     title_append="", alpha=0.5, clear=False, show=True, 
+                     plotDisplace = True)
+
+
+    params = MDS.optimize_library()
+    print(f"Parameters:\nvx, vy, dx, dy, dtheta\n{params.flatten()}")
+    final_undistorted = MDS.undistort(params[:3])
+    transform = convertPoseToTransform(params[3:])
+    solution = (transform @ final_undistorted)[:, :2, 0]
+    plt.figure()
+    plotFakeFeatures(groundTruth, solution, 
+                     title_append="", alpha=0.5, clear=False, show=True, 
+                     plotDisplace = True)
+
+    #TODO: In main loop: for each iteration, track the previous frame's points.
+    # If we need to generate new points, create a new keyframe and add its points
+    # to the global map, using the undistort's velocity and transform to determine
+    # true position. We need to be able to associate local points to global points
+    # in each keyframe, and associate global points to those
+    
\ No newline at end of file
diff --git a/trajectoryPlotting.py b/trajectoryPlotting.py
index c8ced0b..4e0687e 100644
--- a/trajectoryPlotting.py
+++ b/trajectoryPlotting.py
@@ -58,6 +58,17 @@ def appendRelativeTransform(self, time, R, h):
         # new_pose = [*xy[0], *xy[1], self.poses[-1,2] + dth]
         self.poses = np.vstack((self.poses, new_pose))
 
+    def appendAbsoluteTransform(self, time, pose):
+        '''
+        @brief Append a relative transform to the trajectory
+               h should already be scaled by radar resolution
+        @param[in] time timestamp of the transform
+        @param[in] pose pose vector (3, )
+        '''
+        # Add to timestamps
+        self.timestamps = np.append(self.timestamps, time)
+        self.poses = np.vstack((self.poses, pose))
+
     def getPoseAtTimes(self, times):
         '''
         @brief Given timestamps, return the pose at that time using cubic interpolation
@@ -116,7 +127,7 @@ def plotGtAndEstTrajectory(gtTraj,
                            title='GT and EST Trajectories',
                            info=None,
                            savePath=None,
-                           arrow=True):
+                           arrow=False):
     '''
     @brief Plot ground truth trajectory and estimated trajectory
     @param[in] gtTrajectory Ground truth trajectory
diff --git a/utils.py b/utils.py
index 6b60797..1b1e1f1 100644
--- a/utils.py
+++ b/utils.py
@@ -143,3 +143,23 @@ def plt_savefig_by_axis(filePath, fig, ax, pad=0.0):
     extent = plt_full_extent(ax, pad).transformed(fig.dpi_scale_trans.inverted())
     # extent = ax.get_tightbbox(fig.canvas.get_renderer())
     fig.savefig(filePath, bbox_inches=extent)
+
+def invert_transform(T):
+    theta = np.arctan2(T[1, 0], T[0, 0])
+    x = T[0, 2]
+    y = T[1, 2]
+    c = np.cos(theta)
+    s = np.sin(theta)
+    T_inv = np.array([[ c, s, -s * y - c * x],
+                      [-s, c, -c * y + s * x],
+                      [ 0, 0, 1]])
+    return T_inv
+
+def homogenize(points):
+    result = None
+    if points.shape[1] == 2:
+        ones = np.ones((points.shape[0], 1))
+        result = np.concatenate((points, ones) , axis = 1)
+    else:
+        result = points.copy()
+    return result
\ No newline at end of file