import numpy as np
import math
import logging
from typing import Tuple, Optional, Any, Dict
from pydantic import BaseModel

logger = logging.getLogger(__name__)

# --- Data Models ---

class GPSPoint(BaseModel):
    lat: float
    lon: float
    altitude_m: Optional[float] = 0.0

class Sim3Transform(BaseModel):
    translation: np.ndarray
    rotation: np.ndarray
    scale: float
    
    model_config = {"arbitrary_types_allowed": True}

class ObjectGPSResponse(BaseModel):
    gps: GPSPoint
    accuracy_meters: float

class OriginNotSetError(Exception):
    pass

# --- Implementation ---

class CoordinateTransformer:
    """
    F14 (also referenced as F13 in some architectural diagrams): Coordinate Transformer.
    Maps precise pixel object coordinates to Ray-Cloud geospatial intersections.
    Handles transformations between 2D pixels, 3D local maps, 3D global ENU, and GPS.
    """
    def __init__(self, camera_model: Any):
        self.camera_model = camera_model
        self.origins: Dict[str, GPSPoint] = {}
        self.conversion_factors: Dict[str, Tuple[float, float]] = {}

    def _compute_meters_per_degree(self, latitude: float) -> Tuple[float, float]:
        lat_rad = math.radians(latitude)
        meters_per_degree_lat = 111319.5
        meters_per_degree_lon = 111319.5 * math.cos(lat_rad)
        return (meters_per_degree_lon, meters_per_degree_lat)

    def set_enu_origin(self, flight_id: str, origin_gps: GPSPoint) -> None:
        """Sets the global [Lat, Lon, Alt] origin for ENU conversions per flight."""
        self.origins[flight_id] = origin_gps
        self.conversion_factors[flight_id] = self._compute_meters_per_degree(origin_gps.lat)
        logger.info(f"Coordinate Transformer ENU origin set for flight {flight_id}.")

    def get_enu_origin(self, flight_id: str) -> GPSPoint:
        if flight_id not in self.origins:
            raise OriginNotSetError(f"Origin not set for flight {flight_id}")
        return self.origins[flight_id]

    def gps_to_enu(self, flight_id: str, gps: GPSPoint) -> Tuple[float, float, float]:
        """Converts global GPS Geodetic coordinates to local Metric ENU."""
        origin = self.get_enu_origin(flight_id)
        factors = self.conversion_factors[flight_id]
        delta_lat = gps.lat - origin.lat
        delta_lon = gps.lon - origin.lon
        east = delta_lon * factors[0]
        north = delta_lat * factors[1]
        return (east, north, 0.0)

    def enu_to_gps(self, flight_id: str, enu: Tuple[float, float, float]) -> GPSPoint:
        """Converts local metric ENU coordinates to global GPS Geodetic coordinates."""
        origin = self.get_enu_origin(flight_id)
        factors = self.conversion_factors[flight_id]
        east, north, up = enu
        delta_lon = east / factors[0]
        delta_lat = north / factors[1]
        alt = (origin.altitude_m or 0.0) + up
        return GPSPoint(lat=origin.lat + delta_lat, lon=origin.lon + delta_lon, altitude_m=alt)

    def _ray_cloud_intersection(self, ray_origin: np.ndarray, ray_dir: np.ndarray, point_cloud: np.ndarray, max_dist: float = 2.0) -> Optional[np.ndarray]:
        """
        Finds the 3D point in the local point cloud that intersects with (or is closest to) the ray.
        """
        if point_cloud is None or len(point_cloud) == 0:
            return None
            
        # Vectors from the ray origin to all points in the cloud
        w = point_cloud - ray_origin
        
        # Projection scalar of w onto the ray direction
        proj = np.dot(w, ray_dir)
        
        # We only care about points that are in front of the camera (positive projection)
        valid_idx = proj > 0
        if not np.any(valid_idx):
            return None
            
        w_valid = w[valid_idx]
        proj_valid = proj[valid_idx]
        pc_valid = point_cloud[valid_idx]
        
        # Perpendicular distance squared from valid points to the ray (Pythagorean theorem)
        w_sq_norm = np.sum(w_valid**2, axis=1)
        dist_sq = w_sq_norm - (proj_valid**2)
        
        min_idx = np.argmin(dist_sq)
        min_dist = math.sqrt(max(0.0, dist_sq[min_idx]))
        
        if min_dist > max_dist:
            logger.warning(f"No point cloud feature found near the object ray (closest was {min_dist:.2f}m away).")
            return None
            
        # Return the actual 3D feature point from the map
        return pc_valid[min_idx]

    def pixel_to_gps(self, flight_id: str, u: float, v: float, local_pose_matrix: np.ndarray, local_point_cloud: np.ndarray, sim3: Sim3Transform) -> Optional[ObjectGPSResponse]:
        """
        Executes the Ray-Cloud intersection algorithm to geolocate an object in an image.
        Decouples external DEM errors to meet AC-2 and AC-10.
        """
        d_cam = self.camera_model.pixel_to_ray(u, v)
        
        R_local = local_pose_matrix[:3, :3]
        T_local = local_pose_matrix[:3, 3]
        ray_dir_local = R_local @ d_cam
        ray_dir_local = ray_dir_local / np.linalg.norm(ray_dir_local)
        
        p_local = self._ray_cloud_intersection(T_local, ray_dir_local, local_point_cloud)
        if p_local is None: return None
            
        p_metric = sim3.scale * (sim3.rotation @ p_local) + sim3.translation
        gps_coord = self.enu_to_gps(flight_id, (p_metric[0], p_metric[1], p_metric[2]))
        
        return ObjectGPSResponse(gps=gps_coord, accuracy_meters=(5.0 * sim3.scale))