import numpy as np
from typing import Dict, Tuple
from scipy.optimize import minimize
from .pose_math import rvec_tvec_to_matrix, matrix_to_rvec_tvec
from .depth_verify import compute_depth_residual


def extrinsics_to_params(T: np.ndarray) -> np.ndarray:
    rvec, tvec = matrix_to_rvec_tvec(T)
    return np.concatenate([rvec, tvec])


def params_to_extrinsics(params: np.ndarray) -> np.ndarray:
    rvec = params[:3]
    tvec = params[3:]
    return rvec_tvec_to_matrix(rvec, tvec)


def depth_residual_objective(
    params: np.ndarray,
    marker_corners_world: Dict[int, np.ndarray],
    depth_map: np.ndarray,
    K: np.ndarray,
    initial_params: np.ndarray,
    regularization_weight: float = 0.1,
) -> float:
    T = params_to_extrinsics(params)
    residuals = []

    for marker_id, corners in marker_corners_world.items():
        for corner in corners:
            residual = compute_depth_residual(corner, T, depth_map, K, window_size=5)
            if residual is not None:
                residuals.append(residual)

    if len(residuals) == 0:
        return 1e6

    residuals_array = np.array(residuals)
    data_term = np.mean(residuals_array**2)

    param_diff = params - initial_params
    rotation_diff = np.linalg.norm(param_diff[:3])
    translation_diff = np.linalg.norm(param_diff[3:])
    regularization = regularization_weight * (rotation_diff + translation_diff)

    return float(np.real(data_term + regularization))


def refine_extrinsics_with_depth(
    T_initial: np.ndarray,
    marker_corners_world: Dict[int, np.ndarray],
    depth_map: np.ndarray,
    K: np.ndarray,
    max_translation_m: float = 0.05,
    max_rotation_deg: float = 5.0,
    regularization_weight: float = 0.1,
) -> Tuple[np.ndarray, dict[str, float]]:
    initial_params = extrinsics_to_params(T_initial)

    max_rotation_rad = np.deg2rad(max_rotation_deg)

    bounds = [
        (initial_params[0] - max_rotation_rad, initial_params[0] + max_rotation_rad),
        (initial_params[1] - max_rotation_rad, initial_params[1] + max_rotation_rad),
        (initial_params[2] - max_rotation_rad, initial_params[2] + max_rotation_rad),
        (initial_params[3] - max_translation_m, initial_params[3] + max_translation_m),
        (initial_params[4] - max_translation_m, initial_params[4] + max_translation_m),
        (initial_params[5] - max_translation_m, initial_params[5] + max_translation_m),
    ]

    result = minimize(
        depth_residual_objective,
        initial_params,
        args=(
            marker_corners_world,
            depth_map,
            K,
            initial_params,
            regularization_weight,
        ),
        method="L-BFGS-B",
        bounds=bounds,
        options={"maxiter": 100, "disp": False},
    )

    T_refined = params_to_extrinsics(result.x)

    delta_params = result.x - initial_params
    delta_rotation_rad = np.linalg.norm(delta_params[:3])
    delta_rotation_deg = np.rad2deg(delta_rotation_rad)
    delta_translation = np.linalg.norm(delta_params[3:])

    initial_cost = depth_residual_objective(
        initial_params,
        marker_corners_world,
        depth_map,
        K,
        initial_params,
        regularization_weight,
    )

    stats = {
        "iterations": result.nit,
        "success": result.success,
        "initial_cost": float(initial_cost),
        "final_cost": float(result.fun),
        "delta_rotation_deg": float(delta_rotation_deg),
        "delta_translation_norm_m": float(delta_translation),
    }

    return T_refined, stats