"""
Position Estimator Module - Multi-lateration positioning
"""

import numpy as np
from typing import List, Tuple, Dict, Optional
from dataclasses import dataclass


@dataclass
class PositionEstimate:
    """Position estimate with uncertainty"""
    position: np.ndarray  # [x, y, z]
    covariance: np.ndarray  # 3x3 covariance matrix
    timestamp: float
    num_measurements: int
    dilution_of_precision: float


class PositionEstimator:
    """
    Multi-lateration position estimator with outlier rejection
    
    Solves: min Σⱼ [||x - xⱼ|| - rⱼ]²
    """
    
    def __init__(self, max_iterations: int = 10, tolerance: float = 1e-6):
        """
        Initialize position estimator
        
        Args:
            max_iterations: Maximum iterations for solver
            tolerance: Convergence tolerance
        """
        self.max_iterations = max_iterations
        self.tolerance = tolerance
    
    def solve(self, 
             ranges: List[float],
             anchor_positions: List[np.ndarray],
             initial_guess: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Solve multi-lateration problem
        
        Args:
            ranges: List of measured ranges to anchors
            anchor_positions: List of anchor positions [x, y, z]
            initial_guess: Initial position guess
            
        Returns:
            Estimated position [x, y, z]
        """
        if len(ranges) < 4:
            raise ValueError(f"Need at least 4 measurements, got {len(ranges)}")
        
        n = len(ranges)
        anchors = np.array(anchor_positions)
        
        # Initial guess
        if initial_guess is None:
            x0 = np.mean(anchors, axis=0)
        else:
            x0 = initial_guess
        
        # Gauss-Newton iteration
        x = x0.copy()
        
        for iteration in range(self.max_iterations):
            # Compute residuals and Jacobian
            r = np.zeros(n)
            J = np.zeros((n, 3))
            
            for i in range(n):
                d = np.linalg.norm(x - anchors[i])
                r[i] = d - ranges[i]
                if d > 1e-6:
                    J[i] = (x - anchors[i]) / d
            
            # Solve normal equations
            try:
                delta = np.linalg.solve(J.T @ J, -J.T @ r)
            except np.linalg.LinAlgError:
                # Singular matrix - use pseudoinverse
                delta = -np.linalg.pinv(J.T @ J) @ J.T @ r
            
            x += delta
            
            if np.linalg.norm(delta) < self.tolerance:
                break
        
        return x
    
    def solve_weighted(self,
                      ranges: List[float],
                      anchor_positions: List[np.ndarray],
                      weights: List[float],
                      initial_guess: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Weighted least squares solution
        
        Args:
            ranges: List of measured ranges
            anchor_positions: List of anchor positions
            weights: Measurement weights (inverse variance)
            initial_guess: Initial position guess
            
        Returns:
            Weighted position estimate
        """
        if len(ranges) < 4:
            raise ValueError(f"Need at least 4 measurements, got {len(ranges)}")
        
        n = len(ranges)
        anchors = np.array(anchor_positions)
        W = np.diag(weights)
        
        # Initial guess
        if initial_guess is None:
            x0 = np.mean(anchors, axis=0)
        else:
            x0 = initial_guess
        
        # Weighted Gauss-Newton
        x = x0.copy()
        
        for iteration in range(self.max_iterations):
            # Compute residuals and Jacobian
            r = np.zeros(n)
            J = np.zeros((n, 3))
            
            for i in range(n):
                d = np.linalg.norm(x - anchors[i])
                r[i] = d - ranges[i]
                if d > 1e-6:
                    J[i] = (x - anchors[i]) / d
            
            # Weighted normal equations
            JT_W = J.T @ W
            try:
                delta = np.linalg.solve(JT_W @ J, -JT_W @ r)
            except np.linalg.LinAlgError:
                delta = -np.linalg.pinv(JT_W @ J) @ JT_W @ r
            
            x += delta
            
            if np.linalg.norm(delta) < self.tolerance:
                break
        
        return x
    
    def kalman_update(self,
                     measurement: np.ndarray,
                     uncertainty: float,
                     prior_position: np.ndarray,
                     prior_covariance: np.ndarray,
                     dt: float) -> Tuple[np.ndarray, np.ndarray]:
        """
        Kalman filter update step
        
        Args:
            measurement: Position measurement
            uncertainty: Measurement uncertainty (std dev)
            prior_position: Prior position estimate
            prior_covariance: Prior covariance matrix
            dt: Time step
            
        Returns:
            Updated position and covariance
        """
        # State transition matrix (constant velocity model)
        F = np.eye(3)
        
        # Prediction
        predicted_pos = prior_position
        predicted_cov = prior_covariance + np.eye(3) * 0.01 * dt  # Process noise
        
        # Measurement matrix
        H = np.eye(3)
        
        # Measurement noise
        R = np.eye(3) * uncertainty**2
        
        # Innovation
        y = measurement - H @ predicted_pos
        
        # Innovation covariance
        S = H @ predicted_cov @ H.T + R
        
        # Kalman gain
        K = predicted_cov @ H.T @ np.linalg.inv(S)
        
        # Update
        updated_pos = predicted_pos + K @ y
        updated_cov = (np.eye(3) - K @ H) @ predicted_cov
        
        return updated_pos, updated_cov
    
    def reject_outliers(self,
                       measurements: List[float],
                       threshold: float = 3.0) -> Tuple[List[float], List[int]]:
        """
        Reject outlier measurements
        
        Args:
            measurements: List of measurements
            threshold: Rejection threshold in standard deviations
            
        Returns:
            Filtered measurements and indices of kept measurements
        """
        if len(measurements) < 4:
            return measurements, list(range(len(measurements)))
        
        # Robust statistics
        median = np.median(measurements)
        mad = np.median(np.abs(measurements - median))  # Median absolute deviation
        std = mad * 1.4826  # Convert to standard deviation
        
        # Identify inliers
        inliers = []
        kept_indices = []
        
        for i, m in enumerate(measurements):
            if np.abs(m - median) < threshold * std:
                inliers.append(m)
                kept_indices.append(i)
        
        return inliers, kept_indices
    
    def dilution_of_precision(self,
                            anchor_positions: List[np.ndarray],
                            position: np.ndarray) -> Dict[str, float]:
        """
        Calculate dilution of precision metrics
        
        Args:
            anchor_positions: Positions of anchors
            position: Estimated position
            
        Returns:
            DOP metrics (GDOP, PDOP, HDOP, VDOP)
        """
        n = len(anchor_positions)
        anchors = np.array(anchor_positions)
        
        # Build geometry matrix
        A = np.zeros((n, 3))
        for i in range(n):
            d = np.linalg.norm(position - anchors[i])
            if d > 1e-6:
                A[i] = (position - anchors[i]) / d
        
        # Covariance matrix
        try:
            Q = np.linalg.inv(A.T @ A)
        except np.linalg.LinAlgError:
            Q = np.linalg.pinv(A.T @ A)
        
        # DOP metrics
        gdop = np.sqrt(np.trace(Q))
        pdop = np.sqrt(Q[0,0] + Q[1,1] + Q[2,2])
        hdop = np.sqrt(Q[0,0] + Q[1,1])
        vdop = np.sqrt(Q[2,2])
        
        return {
            'GDOP': gdop,
            'PDOP': pdop,
            'HDOP': hdop,
            'VDOP': vdop
        }
    
    def estimate_with_uncertainty(self,
                                 ranges: List[float],
                                 anchor_positions: List[np.ndarray],
                                 range_uncertainties: List[float]) -> PositionEstimate:
        """
        Estimate position with uncertainty quantification
        
        Args:
            ranges: Measured ranges
            anchor_positions: Anchor positions
            range_uncertainties: Uncertainty for each range
            
        Returns:
            PositionEstimate with covariance
        """
        # Initial estimate
        position = self.solve(ranges, anchor_positions)
        
        # Compute covariance
        n = len(ranges)
        anchors = np.array(anchor_positions)
        
        # Build Jacobian at solution
        J = np.zeros((n, 3))
        for i in range(n):
            d = np.linalg.norm(position - anchors[i])
            if d > 1e-6:
                J[i] = (position - anchors[i]) / d
        
        # Measurement covariance
        R = np.diag(np.square(range_uncertainties))
        
        # Position covariance
        try:
            cov = np.linalg.inv(J.T @ np.linalg.inv(R) @ J)
        except np.linalg.LinAlgError:
            cov = np.linalg.pinv(J.T @ np.linalg.inv(R) @ J)
        
        # DOP
        dop = self.dilution_of_precision(anchor_positions, position)
        
        return PositionEstimate(
            position=position,
            covariance=cov,
            timestamp=0.0,
            num_measurements=n,
            dilution_of_precision=dop['GDOP']
        )
    
    def __repr__(self) -> str:
        return f"PositionEstimator(max_iter={self.max_iterations})"
