# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# Modified from https://github.com/facebookresearch/vggt

from typing import Union

import numpy as np
import torch

ArrayLike = Union[np.ndarray, torch.Tensor]


def _is_numpy(x: ArrayLike) -> bool:
    return isinstance(x, np.ndarray)


def _is_torch(x: ArrayLike) -> bool:
    return isinstance(x, torch.Tensor)


def _ensure_torch(x: ArrayLike) -> torch.Tensor:
    """Convert input to torch tensor if it's not already one."""
    if _is_numpy(x):
        return torch.from_numpy(x)
    elif _is_torch(x):
        return x
    else:
        return torch.tensor(x)


def single_undistortion(params, tracks_normalized):
    """
    Apply undistortion to the normalized tracks using the given distortion parameters once.

    Args:
        params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN.
        tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks, 2].

    Returns:
        torch.Tensor: Undistorted normalized tracks tensor.
    """
    params = _ensure_torch(params)
    tracks_normalized = _ensure_torch(tracks_normalized)

    u, v = tracks_normalized[..., 0].clone(), tracks_normalized[..., 1].clone()
    u_undist, v_undist = apply_distortion(params, u, v)
    return torch.stack([u_undist, v_undist], dim=-1)


def iterative_undistortion(
    params,
    tracks_normalized,
    max_iterations=100,
    max_step_norm=1e-10,
    rel_step_size=1e-6,
):
    """
    Iteratively undistort the normalized tracks using the given distortion parameters.

    Args:
        params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN.
        tracks_normalized (torch.Tensor or numpy.ndarray): Normalized tracks tensor of shape [batch_size, num_tracks, 2].
        max_iterations (int): Maximum number of iterations for the undistortion process.
        max_step_norm (float): Maximum step norm for convergence.
        rel_step_size (float): Relative step size for numerical differentiation.

    Returns:
        torch.Tensor: Undistorted normalized tracks tensor.
    """
    params = _ensure_torch(params)
    tracks_normalized = _ensure_torch(tracks_normalized)

    B, N, _ = tracks_normalized.shape
    u, v = tracks_normalized[..., 0].clone(), tracks_normalized[..., 1].clone()
    original_u, original_v = u.clone(), v.clone()

    eps = torch.finfo(u.dtype).eps
    for idx in range(max_iterations):
        u_undist, v_undist = apply_distortion(params, u, v)
        dx = original_u - u_undist
        dy = original_v - v_undist

        step_u = torch.clamp(torch.abs(u) * rel_step_size, min=eps)
        step_v = torch.clamp(torch.abs(v) * rel_step_size, min=eps)

        J_00 = (
            apply_distortion(params, u + step_u, v)[0]
            - apply_distortion(params, u - step_u, v)[0]
        ) / (2 * step_u)
        J_01 = (
            apply_distortion(params, u, v + step_v)[0]
            - apply_distortion(params, u, v - step_v)[0]
        ) / (2 * step_v)
        J_10 = (
            apply_distortion(params, u + step_u, v)[1]
            - apply_distortion(params, u - step_u, v)[1]
        ) / (2 * step_u)
        J_11 = (
            apply_distortion(params, u, v + step_v)[1]
            - apply_distortion(params, u, v - step_v)[1]
        ) / (2 * step_v)

        J = torch.stack(
            [
                torch.stack([J_00 + 1, J_01], dim=-1),
                torch.stack([J_10, J_11 + 1], dim=-1),
            ],
            dim=-2,
        )

        delta = torch.linalg.solve(J, torch.stack([dx, dy], dim=-1))

        u += delta[..., 0]
        v += delta[..., 1]

        if torch.max((delta**2).sum(dim=-1)) < max_step_norm:
            break

    return torch.stack([u, v], dim=-1)


def apply_distortion(extra_params, u, v):
    """
    Applies radial or OpenCV distortion to the given 2D points.

    Args:
        extra_params (torch.Tensor or numpy.ndarray): Distortion parameters of shape BxN, where N can be 1, 2, or 4.
        u (torch.Tensor or numpy.ndarray): Normalized x coordinates of shape Bxnum_tracks.
        v (torch.Tensor or numpy.ndarray): Normalized y coordinates of shape Bxnum_tracks.

    Returns:
        points2D (torch.Tensor): Distorted 2D points of shape BxNx2.
    """
    extra_params = _ensure_torch(extra_params)
    u = _ensure_torch(u)
    v = _ensure_torch(v)

    num_params = extra_params.shape[1]

    if num_params == 1:
        # Simple radial distortion
        k = extra_params[:, 0]
        u2 = u * u
        v2 = v * v
        r2 = u2 + v2
        radial = k[:, None] * r2
        du = u * radial
        dv = v * radial

    elif num_params == 2:
        # RadialCameraModel distortion
        k1, k2 = extra_params[:, 0], extra_params[:, 1]
        u2 = u * u
        v2 = v * v
        r2 = u2 + v2
        radial = k1[:, None] * r2 + k2[:, None] * r2 * r2
        du = u * radial
        dv = v * radial

    elif num_params == 4:
        # OpenCVCameraModel distortion
        k1, k2, p1, p2 = (
            extra_params[:, 0],
            extra_params[:, 1],
            extra_params[:, 2],
            extra_params[:, 3],
        )
        u2 = u * u
        v2 = v * v
        uv = u * v
        r2 = u2 + v2
        radial = k1[:, None] * r2 + k2[:, None] * r2 * r2
        du = u * radial + 2 * p1[:, None] * uv + p2[:, None] * (r2 + 2 * u2)
        dv = v * radial + 2 * p2[:, None] * uv + p1[:, None] * (r2 + 2 * v2)
    else:
        raise ValueError("Unsupported number of distortion parameters")

    u = u.clone() + du
    v = v.clone() + dv

    return u, v


if __name__ == "__main__":
    import random

    import pycolmap

    max_diff = 0
    for i in range(1000):
        # Define distortion parameters (assuming 1 parameter for simplicity)
        B = random.randint(1, 500)
        track_num = random.randint(100, 1000)
        params = torch.rand((B, 1), dtype=torch.float32)  # Batch size 1, 4 parameters
        tracks_normalized = torch.rand(
            (B, track_num, 2), dtype=torch.float32
        )  # Batch size 1, 5 points

        # Undistort the tracks
        undistorted_tracks = iterative_undistortion(params, tracks_normalized)

        for b in range(B):
            pycolmap_intri = np.array([1, 0, 0, params[b].item()])
            pycam = pycolmap.Camera(
                model="SIMPLE_RADIAL",
                width=1,
                height=1,
                params=pycolmap_intri,
                camera_id=0,
            )

            undistorted_tracks_pycolmap = pycam.cam_from_img(
                tracks_normalized[b].numpy()
            )
            diff = (undistorted_tracks[b] - undistorted_tracks_pycolmap).abs().median()
            max_diff = max(max_diff, diff)
            print(f"diff: {diff}, max_diff: {max_diff}")

    import pdb

    pdb.set_trace()