import cv2
import numpy as np
import pyzed.sl as sl
from typing import Optional, Tuple, Dict
from .pose_math import compute_reprojection_error


def create_detector(
    dictionary_id: int = cv2.aruco.DICT_4X4_50,
    params: Optional[cv2.aruco.DetectorParameters] = None,
) -> cv2.aruco.ArucoDetector:
    """
    Creates an ArUco detector with the specified dictionary and parameters.

    Args:
        dictionary_id: The ArUco dictionary ID.
        params: Optional detector parameters. If None, default parameters are used.

    Returns:
        cv2.aruco.ArucoDetector: The configured ArUco detector.
    """
    dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
    if params is None:
        params = cv2.aruco.DetectorParameters()
    return cv2.aruco.ArucoDetector(dictionary, params)


def detect_markers(
    image: np.ndarray, detector: cv2.aruco.ArucoDetector
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    """
    Detects ArUco markers in an image.

    Args:
        image: Input image (BGR or grayscale).
        detector: The ArUco detector to use.

    Returns:
        Tuple[np.ndarray, Optional[np.ndarray]]:
            - corners: Detected corners as float32 shape (N, 4, 2).
            - ids: Detected marker IDs as int shape (N,) or None if no markers detected.
    """
    if image is None:
        raise ValueError("Input image is None")

    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image

    corners, ids, _ = detector.detectMarkers(gray)

    if ids is not None:
        # corners is a list of (1, 4, 2) arrays, reshape to (N, 4, 2)
        corners_array = np.array(corners, dtype=np.float32).reshape(-1, 4, 2)
        ids_array = ids.flatten().astype(np.int32)
        return corners_array, ids_array
    else:
        return np.empty((0, 4, 2), dtype=np.float32), None


def build_camera_matrix_from_zed(
    cam: sl.Camera, resizer: Optional[sl.Resolution] = None
) -> np.ndarray:
    """
    Builds a 3x3 camera intrinsic matrix from a ZED camera's calibration parameters.

    Args:
        cam: The ZED camera object.
        resizer: Optional resolution to get parameters for.

    Returns:
        np.ndarray: 3x3 float64 camera matrix K.
    """
    cam_params = cam.get_camera_information(
        resizer
    ).camera_configuration.calibration_parameters.left_cam
    fx = cam_params.fx
    fy = cam_params.fy
    cx = cam_params.cx
    cy = cam_params.cy

    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)

    return K


def estimate_pose_from_detections(
    corners: np.ndarray,
    ids: np.ndarray,
    marker_geometry: Dict[int, np.ndarray],
    K: np.ndarray,
    dist: Optional[np.ndarray] = None,
    min_markers: int = 4,
    pnp_flag: int = cv2.SOLVEPNP_SQPNP,
) -> Optional[Tuple[np.ndarray, np.ndarray, float, int]]:
    """
    Estimates the camera pose from detected ArUco markers.

    Args:
        corners: Detected corners (N, 4, 2).
        ids: Detected marker IDs (N,).
        marker_geometry: Dictionary mapping marker ID to its 3D corners (4, 3).
        K: 3x3 camera matrix.
        dist: Optional distortion coefficients.
        min_markers: Minimum number of markers required for pose estimation.
        pnp_flag: OpenCV solvePnP flag.

    Returns:
        Optional[Tuple[np.ndarray, np.ndarray, float, int]]:
            (rvec, tvec, reproj_err, n_markers) or None if estimation fails or not enough markers.
    """
    if ids is None or len(ids) == 0:
        return None

    obj_pts = []
    img_pts = []
    n_markers = 0

    for i, marker_id in enumerate(ids):
        if marker_id in marker_geometry:
            obj_pts.append(marker_geometry[marker_id])
            img_pts.append(corners[i])
            n_markers += 1

    if n_markers < min_markers:
        return None

    # Stack points into (M*4, 3) and (M*4, 2)
    obj_pts_array = np.vstack(obj_pts).astype(np.float32)
    img_pts_array = np.vstack(img_pts).astype(np.float32)

    success, rvec, tvec = cv2.solvePnP(
        obj_pts_array, img_pts_array, K, dist, flags=pnp_flag
    )

    if not success:
        return None

    reproj_err = compute_reprojection_error(
        obj_pts_array, img_pts_array, rvec, tvec, K, dist
    )

    return rvec, tvec, reproj_err, n_markers