import cv2
import math
import numpy as np
from numpy.linalg import inv
from typing import Optional
import constants
from scipy.spatial.transform import Rotation
import utility

class Position:
    """Класс позиции с полной ориентацией БПЛА в 3D пространстве"""

    x: float  # Координата X
    y: float  # Координата Y
    z: float  # Масштаб
    yaw: float    # Рыскание (rotation around Z-axis)
    pitch: float  # Тангаж (rotation around Y-axis)
    roll: float   # Крен (rotation around X-axis)

    def __init__(
        self, 
        x: float = 0,
        y: float = 0,
        z: float = 1,
        yaw: float = 0,
        pitch: float = 0,
        roll: float = 0
    ):
        self.x = x
        self.y = y
        self.z = z
        self.yaw = yaw
        self.pitch = pitch
        self.roll = roll

    def __str__(self) -> str:
        return (
            f"Position(x={self.x:.2f}, y={self.y:.2f}, z={self.z:.2f}, "
            f"yaw={math.degrees(self.yaw):.1f}°, "
            f"pitch={math.degrees(self.pitch):.1f}°, "
            f"roll={math.degrees(self.roll):.1f}°)"
        )

    def __imul__(self, scalar: float):
        self.x *= scalar
        self.y *= scalar
        self.z *= scalar
        return self

    def __mul__(self, scalar: float) -> 'Position':
        pos = self.copy()
        pos *= scalar
        return pos

    def __itruediv__(self, scalar: float):
        self.x /= scalar
        self.y /= scalar
        self.z /= scalar
        return self

    def __truediv__(self, scalar: float) -> 'Position':
        pos = self.copy()
        pos /= scalar
        return pos

    def get_homography_matrix(self, K_in: np.ndarray = constants.K, K_out: np.ndarray | None = None, sliding: bool = True) -> np.ndarray:
        """ Возвращает матрицу гомографии """
        R = self.get_rotation_matrix()
        T = self.get_translation_matrix(K_in)
        if not sliding:
            T[0, 2] = T[1, 2] = 0
        if K_out is None: K_out = K_in
        return K_out @ R @ T @ np.linalg.inv(K_in)

    def copy(self) -> 'Position':
        """Создает полную копию объекта"""
        return Position(self.x, self.y, self.z, self.yaw, self.pitch, self.roll)
    
    def get_translation_matrix(self, K: np.ndarray = constants.K) -> np.ndarray:
        return np.array([
            [1, 0, self.x / K[0][0]],
            [0, 1, self.y / K[0][0]],
            [0, 0, self.z]
        ])

    def get_rotation_matrix(self) -> np.ndarray:
        """
        Матрица вращения с порядком применения: yaw → pitch → roll
        Умножение: R = Rx(roll) * Ry(pitch) * Rz(yaw)
        """
        cy, sy = math.cos(self.yaw), math.sin(self.yaw)
        cp, sp = math.cos(self.pitch), math.sin(self.pitch)
        cr, sr = math.cos(self.roll), math.sin(self.roll)

        Rz = np.array([
            [cy, -sy, 0],
            [sy, cy, 0],
            [0, 0, 1],
        ])

        Ry = np.array([
            [cp, 0, sp],
            [0, 1, 0],
            [-sp, 0, cp],
        ])

        Rx = np.array([
            [1, 0, 0],
            [0, cr, -sr],
            [0, sr, cr],
        ])
        
        return Rx @ Ry @ Rz

    def iapply(self, homography_matrix: np.ndarray, K = constants.K) -> 'Position':
        """Применяет матрицу трансформации для вычисления новой позиции и ориентации."""

        np.set_printoptions(suppress=True)
        H = homography_matrix @ self.get_homography_matrix(sliding=False)

        # Decompose homography
        _, R, t, _ = cv2.decomposeHomographyMat(H, K)
        R = np.array(R)
        t = np.array(t)

        # print(cv2.decomposeHomographyMat(inv(H), K))
        # _, _, z, _ = cv2.decomposeHomographyMat(inv(H), K)
        # print(z)
        # z = z.copy()
        # z *= constants._K_FOCUS_DISTANCE
        # print(z)
        
        T = inv(R) @ inv(K) @ H @ K
        ind = np.array([A[2][0] ** 2 + A[2][1] ** 2 for A in T])
        top_k = max(1, len(T) // 2)
        if (len(T) == 3): raise "len(T) == 3"
        ind = np.argpartition(ind, top_k - 1)[:top_k]
        T = T[ind[0]]
        T = T @ np.array([0, 0, 1]) / np.mean((T[0][0], T[1][1]))
        T[2] -= 1
        R = R[ind]
        t = t[ind]

        best_id = ((t - T) ** 2).sum((1, 2)).argmin()

        R = R[best_id]
        rot = Rotation.from_matrix(R).as_euler('XYZ').flatten()
        self.roll = min(np.radians(5), max(np.radians(-5), rot[0]))
        self.pitch = min(np.radians(5), max(np.radians(-5), rot[1]))
        self.yaw = rot[2]

        t = t[best_id].flatten()
        self.x -= T[0] * constants._K_FOCUS_DISTANCE
        self.y += T[1] * constants._K_FOCUS_DISTANCE
        self.z = max(0.7, min(1.3, 1 + T[2]))
        T[0] *= constants._K_FOCUS_DISTANCE
        T[1] *= constants._K_FOCUS_DISTANCE

    def apply(self, homography_matrix: np.ndarray, K = constants.K) -> 'Position':
        """Применяет матрицу трансформации для вычисления новой позиции и ориентации."""
        pos = self.copy()
        pos.iapply(homography_matrix, K)
        return pos