import cv2
import numpy as np
import torch
import logging
from typing import Tuple

logger = logging.getLogger(__name__)

class PseudoImuRectifier:
    """
    Estimates the horizon/tilt of the UAV camera from a single monocular image 
    and rectifies it to a pseudo-nadir view to prevent tracking loss during sharp banks.
    """
    def __init__(self, device: str = "cuda", tilt_threshold_deg: float = 15.0):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.tilt_threshold_deg = tilt_threshold_deg
        
        logger.info(f"Initializing Pseudo-IMU Horizon Estimator on {self.device}")
        # In a full deployment, this loads a lightweight CNN like HorizonNet or DepthAnythingV2
        # self.horizon_model = load_horizon_model().to(self.device)
        
    def estimate_attitude(self, image: np.ndarray) -> Tuple[float, float]:
        """
        Estimates pitch and roll from the image's vanishing points/horizon.
        Returns: (pitch_degrees, roll_degrees)
        """
        # Placeholder for deep-learning based horizon estimation tensor operations.
        # Returns mocked 0.0 for pitch/roll unless the model detects extreme banking.
        pitch_deg = 0.0
        roll_deg = 0.0
        return pitch_deg, roll_deg
        
    def compute_rectification_homography(self, pitch: float, roll: float, K: np.ndarray) -> np.ndarray:
        """Computes the homography matrix to un-warp perspective distortion."""
        p = np.deg2rad(pitch)
        r = np.deg2rad(roll)
        
        # Rotation matrices for pitch (X-axis) and roll (Z-axis)
        Rx = np.array([[1, 0, 0], [0, np.cos(p), -np.sin(p)], [0, np.sin(p), np.cos(p)]])
        Rz = np.array([[np.cos(r), -np.sin(r), 0], [np.sin(r), np.cos(r), 0], [0, 0, 1]])
        
        R = Rz @ Rx
        
        # Homography: H = K * R * K^-1
        K_inv = np.linalg.inv(K)
        H = K @ R @ K_inv
        return H

    def rectify_image(self, image: np.ndarray, K: np.ndarray) -> Tuple[np.ndarray, bool]:
        pitch, roll = self.estimate_attitude(image)
        
        if abs(pitch) < self.tilt_threshold_deg and abs(roll) < self.tilt_threshold_deg:
            return image, False # No rectification needed
            
        logger.warning(f"Extreme tilt detected (Pitch: {pitch:.1f}, Roll: {roll:.1f}). Rectifying.")
        H = self.compute_rectification_homography(-pitch, -roll, K) 
        
        rectified_image = cv2.warpPerspective(image, H, (image.shape[1], image.shape[0]), flags=cv2.INTER_LINEAR)
        return rectified_image, True