from typing import List, Tuple, Optional

class EKFSLAM:
    def __init__(self, state_dim: int = 6):
        self.state_dim = state_dim  # Robot state dimensions
        self.state = np.zeros(state_dim)
        self.P = np.eye(state_dim) * 0.1  # Covariance
        self.landmarks = []  # List of landmark positions
        self.P_landmarks = []  # Landmark covariances
    
    def motion_model(self, state, control, dt):
        """Unicycle motion model for drone."""
        x, y, z, roll, pitch, yaw = state
        v_x, v_y, v_z, omega_x, omega_y, omega_z = control
        
        new_state = np.array([
            x + v_x * np.cos(yaw) * dt - v_y * np.sin(yaw) * dt,
            y + v_x * np.sin(yaw) * dt + v_y * np.cos(yaw) * dt,
            z + v_z * dt,
            roll + omega_x * dt,
            pitch + omega_y * dt,
            yaw + omega_z * dt
        ])
        return new_state
    
    def Jacobian_motion(self, state, control, dt):
        """Jacobian of motion model w.r.t. state."""
        x, y, z, roll, pitch, yaw = state
        v_x = control[0]
        
        F = np.eye(self.state_dim)
        F[0, 5] = -v_x * np.sin(yaw) * dt
        F[1, 5] = v_x * np.cos(yaw) * dt
        return F
    
    def predict(self, control, dt, Q):
        """EKF Predict step."""
        F = self.Jacobian_motion(self.state, control, dt)
        self.state = self.motion_model(self.state, control, dt)
        self.P = F @ self.P @ F.T + Q
    
    def observation_model(self, robot_state, landmark_pos):
        """Compute expected observation from robot to landmark."""
        dx = landmark_pos[0] - robot_state[0]
        dy = landmark_pos[1] - robot_state[1]
        dz = landmark_pos[2] - robot_state[2]
        dist = np.sqrt(dx**2 + dy**2 + dz**2)
        angle_y = np.arctan2(dx, dy) - robot_state[5]
        angle_p = np.arcsin(dz / dist) if dist > 0.01 else 0
        
        return np.array([dist, angle_y, angle_p])
    
    def update(self, observation, landmark_idx, R):
        """EKF Update step with landmark observation."""
        if landmark_idx >= len(self.landmarks):
            self.add_landmark(observation)
            return
        
        landmark = self.landmarks[landmark_idx]
        expected = self.observation_model(self.state, landmark)
        
        innovation = observation - expected
        innovation[1] = np.arctan2(np.sin(innovation[1]), np.cos(innovation[1]))
        
        H = self.jacobian_observation(self.state, landmark)
        
        S = H @ self.P @ H.T + R
        K = self.P @ H.T @ np.linalg.inv(S)
        
        self.state = self.state + K @ innovation
        self.state[5] = np.arctan2(np.sin(self.state[5]), np.cos(self.state[5]))
        self.P = (np.eye(self.state_dim) - K @ H) @ self.P
    
    def add_landmark(self, observation):
        """Initialize a new landmark from observation."""
        dist, angle_y, angle_p = observation
        yaw = self.state[5]
        
        lx = self.state[0] + dist * np.cos(angle_p) * np.sin(angle_y + yaw)
        ly = self.state[1] + dist * np.cos(angle_p) * np.cos(angle_y + yaw)
        lz = self.state[2] + dist * np.sin(angle_p)
        
        self.landmarks.append(np.array([lx, ly, lz]))
        self.P_landmarks.append(np.eye(3) * 0.5)
    
    def jacobian_observation(self, robot_state, landmark_pos):
        """Jacobian of observation model."""
        dx = landmark_pos[0] - robot_state[0]
        dy = landmark_pos[1] - robot_state[1]
        dz = landmark_pos[2] - robot_state[2]
        dist = np.sqrt(dx**2 + dy**2 + dz**2)
        
        H = np.zeros((3, self.state_dim))
        H[0, 0] = dx / dist
        H[0, 1] = dy / dist
        H[0, 2] = dz / dist
        H[1, 5] = -1
        H[2, 0] = -dx * dz / (dist**2 * np.sqrt(1 - (dz/dist)**2))
        H[2, 1] = -dy * dz / (dist**2 * np.sqrt(1 - (dz/dist)**2))
        H[2, 2] = (dx**2 + dy**2) / (dist**2 * np.sqrt(1 - (dz/dist)**2))
        
        return H
    
    def detect_loop_closure(self, current_obs, threshold=0.1):
        """Simple loop closure detection based on landmark re-observation."""
        for i, landmark in enumerate(self.landmarks):
            expected = self.observation_model(self.state, landmark)
            if np.allclose(current_obs, expected, atol=threshold):
                return i
        return None

class SimpleVisualSLAM:
    def __init__(self):
        self.orb = cv2.ORB_create(nfeatures=2000)
        self.bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        self.keyframes = []
        self.map_points = []
        self.camera_matrix = None
        self.dist_coeffs = None
    
    def extract_features(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        keypoints, descriptors = self.orb.detectAndCompute(gray, None)
        return keypoints, descriptors
    
    def track_pose(self, prev_desc, curr_desc, prev_kp, curr_kp, K):
        matches = self.bf.match(prev_desc, curr_desc)
        matches = sorted(matches, key=lambda x: x.distance)
        
        if len(matches) < 10:
            return None, None
        
        pts_prev = np.float32([prev_kp[m.queryIdx].pt for m in matches])
        pts_curr = np.float32([curr_kp[m.trainIdx].pt for m in matches])
        
        E, mask = cv2.findEssentialMat(pts_curr, pts_prev, K, 
                                        method=cv2.RANSAC, threshold=1.0)
        _, R, t, mask = cv2.recoverPose(E, pts_curr, pts_prev, K)
        
        return R, t
    
    def triangulate_points(self, R1, t1, R2, t2, pts1, pts2, K):
        P1 = K @ np.hstack([np.eye(3), np.zeros((3, 1))])
        P2 = K @ np.hstack([R2, t2.reshape(3, 1)])
        
        pts4d = cv2.triangulatePoints(P1, P2, pts1.T, pts2.T)
        pts3d = pts4d[:3] / pts4d[3]
        
        return pts3d.T
    
    def optimize_bundle_adjustment(self, poses, points, observations):
        from scipy.optimize import least_squares
        
        def residual(params):
            n_poses = len(poses)
            residuals = []
            for i, (pose_idx, pt_idx, obs) in enumerate(observations):
                pose = params[pose_idx * 6:pose_idx * 6 + 6]
                point = params[n_poses * 6 + pt_idx * 3:n_poses * 6 + pt_idx * 3 + 3]
                projected = project_point(point, pose, K)
                residuals.extend(obs - projected)
            return np.array(residuals)
        
        x0 = np.concatenate([poses.flatten(), points.flatten()])
        result = least_squares(residual, x0, method='lm')
        return result.x

from sklearn.neighbors import KDTree

class ICP:
    def __init__(self, max_iterations=50, tolerance=1e-6):
        self.max_iterations = max_iterations
        self.tolerance = tolerance
    
    def fit(self, source, target):
        tree = KDTree(target)
        R_total = np.eye(3)
        t_total = np.zeros(3)
        
        for iteration in range(self.max_iterations):
            # Find closest points
            distances, indices = tree.query(source)
            
            # Compute centroids
            src_centroid = np.mean(source, axis=0)
            tgt_centroid = np.mean(target[indices], axis=0)
            
            # Center the points
            src_centered = source - src_centroid
            tgt_centered = target[indices] - tgt_centroid
            
            # SVD for optimal rotation
            H = src_centered.T @ tgt_centered
            U, S, Vt = np.linalg.svd(H)
            R = Vt.T @ U.T
            
            # Ensure proper rotation (det = 1)
            if np.linalg.det(R) < 0:
                Vt[-1, :] *= -1
                R = Vt.T @ U.T
            
            # Translation
            t = tgt_centroid - R @ src_centroid
            
            # Apply transformation
            source = (R @ source.T).T + t
            R_total = R @ R_total
            t_total = R @ t_total + t
            
            # Check convergence
            if np.linalg.norm(t) < self.tolerance:
                break
        
        return R_total, t_total

class LoopClosureDetector:
    def __init__(self, vocabulary_size=1000):
        self.vocabulary_size = vocabulary_size
        self.db_descriptors = []
        self.db_frames = []
        self.bow_model = None
    
    def compute_bow_descriptor(self, descriptors):
        """Compute Bag of Words descriptor from ORB features."""
        if descriptors is None:
            return np.zeros(self.vocabulary_size)
        
        bow = np.zeros(self.vocabulary_size)
        for desc in descriptors:
            # Simplified: assign to nearest visual word
            word_id = hash(desc.tobytes()) % self.vocabulary_size
            bow[word_id] += 1
        
        # L1 normalize
        if np.sum(bow) > 0:
            bow = bow / np.sum(bow)
        return bow
    
    def detect_loop(self, current_bow, threshold=0.7):
        """Detect if current frame matches a previous keyframe."""
        best_score = 0
        best_frame = -1
        
        for i, (bow, frame_id) in enumerate(zip(self.db_descriptors, self.db_frames)):
            score = self.bow_similarity(current_bow, bow)
            if score > best_score and score > threshold:
                best_score = score
                best_frame = frame_id
        
        return best_frame, best_score
    
    def bow_similarity(self, bow1, bow2):
        """Compute cosine similarity between BoW descriptors."""
        dot = np.dot(bow1, bow2)
        norm1 = np.linalg.norm(bow1)
        norm2 = np.linalg.norm(bow2)
        if norm1 > 0 and norm2 > 0:
            return dot / (norm1 * norm2)
        return 0
    
    def add_to_database(self, descriptors, frame_id):
        bow = self.compute_bow_descriptor(descriptors)
        self.db_descriptors.append(bow)
        self.db_frames.append(frame_id)

from collections import deque

class VisualInertialSLAM:
    def __init__(self, camera_params, imu_params):
        self.camera = camera_params
        self.imu = imu_params
        
        # State: [position, velocity, orientation, gyro_bias, accel_bias]
        self.state = np.zeros(15)
        self.covariance = np.eye(15) * 0.01
        
        # Map
        self.landmarks = {}  # id -> (position, descriptor, covariance)
        self.keyframes = deque(maxlen=20)
        
        # Parameters
        self.gravity = np.array([0, 0, -9.81])
        self.feature_detector = cv2.ORB_create(nfeatures=500)
    
    def imu_callback(self, accel, gyro, dt):
        """Process IMU measurement - prediction step."""
        # Remove gravity from acceleration
        R = self.rotation_matrix(self.state[6:9])
        accel_world = R @ accel - self.gravity
        
        # Update velocity and position
        self.state[3:6] += accel_world * dt
        self.state[0:3] += self.state[3:6] * dt + 0.5 * accel_world * dt**2
        
        # Update orientation
        omega = gyro - self.state[12:15]  # Remove bias
        self.state[6:9] += omega * dt
        
        # Update covariance
        F = self.state_transition_matrix(dt)
        Q = self.process_noise_covariance(dt)
        self.covariance = F @ self.covariance @ F.T + Q
    
    def vision_callback(self, frame):
        """Process camera frame - update step."""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        keypoints, descriptors = self.feature_detector.detectAndCompute(gray, None)
        
        if len(keypoints) < 10:
            return
        
        observations = []
        for kp, desc in zip(keypoints, descriptors):
            # Try to match with existing landmarks
            match_id = self.match_landmark(desc)
            
            if match_id is not None:
                # Update existing landmark
                obs = self.project_to_3d(kp.pt)
                self.update_landmark(match_id, obs)
                observations.append((match_id, obs))
            else:
                # Create new landmark
                pos_3d = self.triangulate_new_point(kp.pt)
                if pos_3d is not None:
                    self.add_landmark(pos_3d, desc)
        
        # Loop closure check
        if len(self.keyframes) > 10:
            loop_frame = self.detect_loop_closure(descriptors)
            if loop_frame is not None:
                self.optimize_with_loop_closure(loop_frame)
        
        self.keyframes.append((self.state.copy(), keypoints, descriptors))
    
    def rotation_matrix(self, euler_angles):
        roll, pitch, yaw = euler_angles
        Rx = np.array([[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]])
        Ry = np.array([[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]])
        Rz = np.array([[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]])
        return Rz @ Ry @ Rx
    
    def state_transition_matrix(self, dt):
        F = np.eye(15)
        F[0:3, 3:6] = np.eye(3) * dt
        return F
    
    def process_noise_covariance(self, dt):
        Q = np.zeros((15, 15))
        Q[0:3, 0:3] = np.eye(3) * 0.01 * dt**2
        Q[3:6, 3:6] = np.eye(3) * 0.1 * dt**2
        Q[6:9, 6:9] = np.eye(3) * 0.01 * dt**2
        Q[9:12, 9:12] = np.eye(3) * 0.001 * dt
        Q[12:15, 12:15] = np.eye(3) * 0.0001 * dt
        return Q
    
    def match_landmark(self, descriptor):
        """Find best matching landmark by descriptor distance."""
        best_id = None
        best_dist = 100
        
        for lid, (_, ldesc, _) in self.landmarks.items():
            dist = np.sum(np.abs(descriptor.astype(int) - ldesc.astype(int)))
            if dist < best_dist:
                best_dist = dist
                best_id = lid
        
        return best_id if best_dist < 50 else None
    
    def run_slam(self, imu_data, image_data, dt=0.01):
        """Main SLAM loop."""
        trajectory = []
        
        for i, (accel, gyro, frame) in enumerate(zip(imu_data, image_data)):
            # IMU prediction at higher rate
            self.imu_callback(accel, gyro, dt)
            
            # Vision update at camera rate
            if i % 5 == 0:  # 20Hz vision, 100Hz IMU
                self.vision_callback(frame)
            
            trajectory.append(self.state[0:3].copy())
        
        return np.array(trajectory), self.landmarks