# -*- encoding: utf-8 -*-

"""
@File    :   camera.py
@Time    :   2025/9/3 15:25:00
@Author  :   lh9171338
@Version :   1.0
@Contact :   2909171338@qq.com
"""

import numpy as np
import cv2


class Camera:
    """
    Base Camera

    Args:
        coeff (dict | None): camera coefficients
        **kwargs: keyword arguments
    """

    def __init__(self, coeff=None, **kwargs):
        self.coeff = coeff
        self.format_coeff()

    def format_coeff(self):
        """
        Format coeff

        Args:
            None

        Returns:
            None
        """
        if self.coeff:
            self.coeff = {k: np.array(v) for k, v in self.coeff.items()}

    def load_coeff(self, filename):
        """
        Load coeff

        Args:
            filename (str): filename

        Returns:
            None
        """
        fs = cv2.FileStorage(filename, cv2.FileStorage_READ)
        K = fs.getNode("K").mat()
        D = fs.getNode("D").mat()
        fs.release()
        self.coeff = {"K": K, "D": D}

    def save_coeff(self, filename):
        """
        Save coeff

        Args:
            filename (str): filename

        Returns:
            None
        """
        fs = cv2.FileStorage(filename, cv2.FileStorage_WRITE)
        fs.write("K", self.coeff["K"])
        fs.write("D", self.coeff["D"])
        fs.release()

    def distort_point(self, undistorted):
        """
        Distort point

        Args:
            undistorted (np.ndarray): undistorted points, shape [N, 2]

        Returns:
            distorted (np.ndarray): distorted points, shape [N, 2]
        """
        raise NotImplementedError

    def undistort_point(self, distorted):
        """
        Undistort point

        Args:
            distorted (np.ndarray): distorted points, shape [N, 2]

        Returns:
            undistorted (np.ndarray): undistorted points, shape [N, 2]
        """
        raise NotImplementedError

    def distort_image(self, image, transform=None):
        """
        Distort image

        Args:
            image (np.ndarray): image
            transform (list): transform, [tx, ty, sx, sy]

        Returns:
            image (np.ndarray): distorted image
        """
        if transform is None:
            transform = [0.0, 0.0, 1.0, 1.0]
        tx, ty, sx, sy = transform[0], transform[1], transform[2], transform[3]

        height, width = image.shape[0], image.shape[1]

        distorted = np.mgrid[0:width, 0:height].T.reshape(-1, 2).astype(np.float64)
        undistorted = self.undistort_point(distorted)
        undistorted = undistorted.reshape(height, width, 2)
        map1 = (undistorted[:, :, 0].astype(np.float32) - tx) / sx
        map2 = (undistorted[:, :, 1].astype(np.float32) - ty) / sy

        image = cv2.remap(image, map1, map2, cv2.INTER_CUBIC)

        return image

    def undistort_image(self, image, transform=None):
        """
        Undistort image

        Args:
            image (np.ndarray): image
            transform (list): transform, [tx, ty, sx, sy]

        Returns:
            image (np.ndarray): undistorted image
        """
        if transform is None:
            transform = [0.0, 0.0, 1.0, 1.0]
        tx, ty, sx, sy = transform[0], transform[1], transform[2], transform[3]

        height, width = image.shape[0], image.shape[1]

        undistorted = np.mgrid[0:width, 0:height].T.reshape(-1, 2).astype(np.float64)
        undistorted[:, 0] = undistorted[:, 0] * sx + tx
        undistorted[:, 1] = undistorted[:, 1] * sy + ty
        distorted = self.distort_point(undistorted)
        distorted = distorted.reshape(height, width, 2)
        map1 = distorted[:, :, 0].astype(np.float32)
        map2 = distorted[:, :, 1].astype(np.float32)
        image = cv2.remap(image, map1, map2, cv2.INTER_CUBIC)

        return image

    def interp_line(self, lines, num=None, resolution=1.0):
        """
        Interpolate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            num (int | None): number of interpolated points per line
            resolution (float): resolution of interpolation

        Returns:
            pts_list (list): list of interpolated points
        """
        raise NotImplementedError

    def interp_arc(self, arcs, num=None, resolution=0.01):
        """
        Interpolate arc

        Args:
            arcs (np.ndarray): arcs, shape [N, 2, 2]
            num (int | None): number of interpolated points per line
            resolution (float): resolution of interpolation

        Returns:
            pts_list (list): list of interpolated points
        """
        resolution *= np.pi / 180.0

        pts_list = []
        for arc in arcs:
            pt1, pt2 = arc[0], arc[1]
            normal = np.cross(pt1, pt2)
            normal /= np.linalg.norm(normal)
            angle = np.arccos(normal[2])
            axes = np.array([-normal[1], normal[0], 0])
            axes /= max(np.linalg.norm(axes), np.finfo(np.float64).eps)
            rotation_vector = angle * axes
            rotation_matrix, _ = cv2.Rodrigues(rotation_vector)
            pt1 = np.matmul(rotation_matrix.T, pt1[:, None]).flatten()
            pt2 = np.matmul(rotation_matrix.T, pt2[:, None]).flatten()
            min_angle = np.arctan2(pt1[1], pt1[0])
            max_angle = np.arctan2(pt2[1], pt2[0])
            if max_angle < min_angle:
                max_angle += 2 * np.pi

            K = int(round((max_angle - min_angle) / resolution) + 1) if num is None else num
            angles = np.linspace(min_angle, max_angle, K)
            pts = np.hstack((np.cos(angles)[:, None], np.sin(angles)[:, None], np.zeros((K, 1))))
            pts = np.matmul(rotation_matrix, pts.T).T
            pts_list.append(pts)

        return pts_list

    def insert_line(self, image, pts_list, color, thickness=1):
        """
        Insert line

        Args:
            image (np.ndarray): image
            pts_list (list): list of points
            color (tuple): color
            thickness (int): thickness

        Returns:
            image (np.ndarray): image
        """
        for pts in pts_list:
            pts = np.round(pts).astype(np.int32)
            cv2.polylines(image, [pts], isClosed=False, color=color, thickness=thickness)

        return image

    def truncate_line(self, lines):
        """
        Truncate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            image_size (tuple): image size, [width, height]

        Returns:
            lines (np.ndarray): truncated lines, shape [M, 2, 2]
        """
        return lines


class Pinhole(Camera):
    """
    Pinhole camera
    """

    def distort_point(self, undistorted):
        """
        Distort point

        Args:
            undistorted (np.ndarray): undistorted points, shape [N, 2]

        Returns:
            distorted (np.ndarray): distorted points, shape [N, 2]
        """
        if self.coeff is not None:
            K, D = self.coeff["K"], self.coeff["D"]
            fx, fy = K[0, 0], K[1, 1]
            cx, cy = K[0, 2], K[1, 2]

            undistorted = undistorted.copy().astype(np.float64)
            undistorted[:, 0] = (undistorted[:, 0] - cx) / fx
            undistorted[:, 1] = (undistorted[:, 1] - cy) / fy
            undistorted = np.hstack((undistorted, np.ones((undistorted.shape[0], 1), np.float64)))
            distorted = cv2.projectPoints(undistorted.reshape(1, -1, 3), (0, 0, 0), (0, 0, 0), K, D)[0].reshape(-1, 2)
        else:
            distorted = undistorted

        return distorted

    def undistort_point(self, distorted):
        """
        Undistort point

        Args:
            distorted (np.ndarray): distorted points, shape [N, 2]

        Returns:
            undistorted (np.ndarray): undistorted points, shape [N, 2]
        """
        if self.coeff is not None:
            K, D = self.coeff["K"], self.coeff["D"]
            distorted = distorted.copy().astype(np.float64)
            undistorted = cv2.undistortPoints(distorted.reshape(1, -1, 2), K, D, R=None, P=K).reshape(-1, 2)
        else:
            undistorted = distorted

        return undistorted

    def interp_line(self, lines, num=None, resolution=0.1):
        """
        Interpolate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            num (int | None): number of interpolated points per line
            resolution (float): resolution of interpolation

        Returns:
            pts_list (list): list of interpolated points
        """
        distorted = lines.reshape(-1, 2)
        undistorted = self.undistort_point(distorted)
        lines = undistorted.reshape(-1, 2, 2)

        pts_list = []
        for line in lines:
            K = num or int(round(max(abs(line[1] - line[0])) / resolution)) + 1
            lambda_ = np.linspace(0, 1, K)[:, None]
            pts = line[1] * lambda_ + line[0] * (1 - lambda_)
            pts = self.distort_point(pts)
            pts_list.append(pts)

        return pts_list

    def insert_line(self, image, lines, color, thickness=1):
        """
        Insert line

        Args:
            image (np.ndarray): image
            lines (np.ndarray): lines, shape [N, 2, 2]
            color (tuple): color
            thickness (int): thickness

        Returns:
            image (np.ndarray): image
        """
        pts_list = self.interp_line(lines)
        super().insert_line(image, pts_list, color, thickness)

        return image


class Fisheye(Camera):
    """
    Fisheye camera
    """

    def distort_point(self, undistorted):
        """
        Distort point

        Args:
            undistorted (np.ndarray): undistorted points, shape [N, 2]

        Returns:
            distorted (np.ndarray): distorted points, shape [N, 2]
        """
        undistorted = undistorted.copy().astype(np.float64)

        K, D = self.coeff["K"], self.coeff["D"]
        fx, fy = K[0, 0], K[1, 1]
        cx, cy = K[0, 2], K[1, 2]

        undistorted[:, 0] = (undistorted[:, 0] - cx) / fx
        undistorted[:, 1] = (undistorted[:, 1] - cy) / fy
        distorted = cv2.fisheye.distortPoints(undistorted.reshape(1, -1, 2), K, D).reshape(-1, 2)

        return distorted

    def undistort_point(self, distorted):
        """
        Undistort point

        Args:
            distorted (np.ndarray): distorted points, shape [N, 2]

        Returns:
            undistorted (np.ndarray): undistorted points, shape [N, 2]
        """
        distorted = distorted.copy().astype(np.float64)

        K, D = self.coeff["K"], self.coeff["D"]
        undistorted = cv2.fisheye.undistortPoints(distorted.reshape(1, -1, 2), K, D, P=K).reshape(-1, 2)

        return undistorted

    def interp_line(self, lines, num=None, resolution=0.1):
        """
        Interpolate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            num (int | None): number of interpolated points per line
            resolution (float): resolution of interpolation

        Returns:
            pts_list (list): list of interpolated points
        """
        distorted = lines.reshape(-1, 2)
        undistorted = self.undistort_point(distorted)
        undistorted = np.hstack((undistorted, np.ones((undistorted.shape[0], 1), np.float64)))
        undistorted = undistorted / np.linalg.norm(undistorted, axis=1, keepdims=True)

        arcs = undistorted.reshape(-1, 2, 3)
        undistorted_list = self.interp_arc(arcs, num, resolution)
        distorted_list = []
        for undistorted in undistorted_list:
            undistorted = undistorted / (undistorted[:, 2:] + np.finfo(np.float64).eps)
            undistorted = undistorted[:, :2]
            distorted = self.distort_point(undistorted)
            distorted_list.append(distorted)

        return distorted_list

    def insert_line(self, image, lines, color, thickness=1):
        """
        Insert line

        Args:
            image (np.ndarray): image
            lines (np.ndarray): lines, shape [N, 2, 2]
            color (tuple): color
            thickness (int): thickness

        Returns:
            image (np.ndarray): image
        """
        pts_list = self.interp_line(lines)
        super().insert_line(image, pts_list, color, thickness)

        return image


class Spherical(Camera):
    """
    Spherical camera

    Args:
        image_size (tuple): image size, [width, height]
        **kwargs: keyword arguments
    """

    def __init__(self, image_size, **kwargs):
        super().__init__(**kwargs)

        self.image_size = image_size

    def distort_point(self, undistorted):
        """
        Distort point

        Args:
            undistorted (np.ndarray): undistorted points, shape [N, 3]

        Returns:
            distorted (np.ndarray): distorted points, shape [N, 2]
        """
        undistorted = undistorted.copy().astype(np.float64)
        width, height = self.image_size

        if self.coeff is not None:
            K, D = self.coeff["K"], self.coeff["D"]
            cx = cy = (height - 1.0) / 2.0

            mask = undistorted[:, 2] < 0
            undistorted[mask, 0] = -undistorted[mask, 0]
            undistorted[mask, 2] = -undistorted[mask, 2]
            undistorted = undistorted / (undistorted[:, 2:] + np.finfo(np.float64).eps)
            undistorted = undistorted[:, :2]
            distorted = cv2.fisheye.distortPoints(undistorted.reshape(1, -1, 2), K, D).reshape(-1, 2)
            x = (distorted[:, 0] - cx) / cx
            y = (distorted[:, 1] - cy) / cy
            theta = np.arctan2(y, x)
            phi = np.sqrt(x**2 + y**2) * np.pi / 2.0
            x = np.sin(phi) * np.cos(theta)
            y = np.sin(phi) * np.sin(theta)
            z = np.cos(phi)
            undistorted = np.hstack((x[:, None], y[:, None], z[:, None]))
            undistorted[mask, 0] = -undistorted[mask, 0]
            undistorted[mask, 2] = -undistorted[mask, 2]

        x, y, z = undistorted[:, 0], undistorted[:, 1], undistorted[:, 2]
        lat = np.pi - np.arccos(y)
        lon = np.pi - np.arctan2(z, x)
        u = width * lon / (2 * np.pi)
        v = height * lat / np.pi
        u = np.mod(u, width)
        v = np.mod(v, height)
        distorted = np.stack([u, v], axis=-1)

        return distorted

    def undistort_point(self, distorted):
        """
        Undistort point

        Args:
            distorted (np.ndarray): distorted points, shape [N, 2]

        Returns:
            undistorted (np.ndarray): undistorted points, shape [N, 3]
        """
        distorted = distorted.copy().astype(np.float64)
        width, height = self.image_size

        u, v = distorted[:, 0], distorted[:, 1]
        lon = np.pi - u / width * 2 * np.pi
        lat = np.pi - v / height * np.pi
        y = np.cos(lat)
        x = np.sin(lat) * np.cos(lon)
        z = np.sin(lat) * np.sin(lon)
        undistorted = np.stack([x, y, z], axis=-1)

        if self.coeff is not None:
            K, D = self.coeff["K"], self.coeff["D"]
            cx = cy = (height - 1.0) / 2.0

            mask = undistorted[:, 2] < 0
            undistorted[mask, 0] = -undistorted[mask, 0]
            undistorted[mask, 2] = -undistorted[mask, 2]
            x, y, z = undistorted[:, 0], undistorted[:, 1], undistorted[:, 2]
            theta = np.arctan2(y, x)
            phi = np.arccos(z)
            r = phi * 2.0 / np.pi
            x = r * np.cos(theta) * cx + cx
            y = r * np.sin(theta) * cy + cy
            distorted = np.hstack((x[:, None], y[:, None]))
            undistorted = cv2.fisheye.undistortPoints(distorted.reshape(1, -1, 2), K, D).reshape(-1, 2)
            undistorted = np.hstack((undistorted, np.ones((undistorted.shape[0], 1), np.float64)))
            undistorted = undistorted / np.linalg.norm(undistorted, axis=1, keepdims=True)
            undistorted[mask, 0] = -undistorted[mask, 0]
            undistorted[mask, 2] = -undistorted[mask, 2]

        return undistorted

    def interp_line(self, lines, num=None, resolution=0.01):
        """
        Interpolate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            num (int | None): number of interpolated points per line
            resolution (float): resolution of interpolation

        Returns:
            pts_list (list): list of interpolated points
        """
        distorted = lines.reshape(-1, 2)
        undistorted = self.undistort_point(distorted)
        arcs = undistorted.reshape(-1, 2, 3)
        undistorted_list = self.interp_arc(arcs, num, resolution)
        distorted_list = []
        for undistorted in undistorted_list:
            distorted = self.distort_point(undistorted)
            distorted_list.append(distorted)

        return distorted_list

    def insert_line(self, image, lines, color, thickness=1):
        """
        Insert line

        Args:
            image (np.ndarray): image
            lines (np.ndarray): lines, shape [N, 2, 2]
            color (tuple): color
            thickness (int): thickness

        Returns:
            image (np.ndarray): image
        """
        pts_list = self.interp_line(lines)
        super().insert_line(image, pts_list, color, thickness)

        return image

    def truncate_line(self, lines):
        """
        Truncate line

        Args:
            lines (np.ndarray): lines, shape [N, 2, 2]
            image_size (tuple): image size, [width, height]

        Returns:
            lines (np.ndarray): truncated lines, shape [M, 2, 2]
        """
        width = self.image_size[0]
        pts_list = self.interp_line(lines)
        lines = []
        for pts in pts_list:
            dx = abs(pts[:-1, 0] - pts[1:, 0])
            mask = dx > width / 2.0
            s = sum(mask)
            assert s <= 1
            if s == 0:
                lines.append([pts[0], pts[-1]])
            else:
                ind = np.where(mask)[0][0]
                lines.append([pts[0], pts[ind]])
                lines.append([pts[ind + 1], pts[-1]])
        lines = np.asarray(lines)

        return lines