import logging
import math
import typing

import cv2 as cv
import numpy as np
from wpimath.geometry import Rotation2d, Rotation3d, Translation3d
from wpimath.units import hertz, seconds

from ..estimation import RotTrlTransform3d


class SimCameraProperties:
    def __init__(self):
        self.resWidth: int = -1
        self.resHeight: int = -1
        self.camIntrinsics: np.ndarray = np.zeros((3, 3))  # [3,3]
        self.distCoeffs: np.ndarray = np.zeros((8, 1))  # [8,1]
        self.avgErrorPx: float = 0.0
        self.errorStdDevPx: float = 0.0
        self.frameSpeed: seconds = 0.0
        self.exposureTime: seconds = 0.0
        self.avgLatency: seconds = 0.0
        self.latencyStdDev: seconds = 0.0
        self.viewplanes: list[np.ndarray] = []  # [3,1]

        self.setCalibrationFromFOV(960, 720, fovDiag=Rotation2d(math.radians(90.0)))

    def setCalibrationFromFOV(
        self, width: int, height: int, fovDiag: Rotation2d
    ) -> None:
        if fovDiag.degrees() < 1.0 or fovDiag.degrees() > 179.0:
            fovDiag = Rotation2d.fromDegrees(max(min(fovDiag.degrees(), 179.0), 1.0))
            logging.error("Requested invalid FOV! Clamping between (1, 179) degrees...")

        resDiag = math.sqrt(width * width + height * height)
        diagRatio = math.tan(fovDiag.radians() / 2.0)
        fovWidth = Rotation2d(math.atan((diagRatio * (width / resDiag)) * 2))
        fovHeight = Rotation2d(math.atan(diagRatio * (height / resDiag)) * 2)

        newDistCoeffs = np.zeros((8, 1))

        cx = width / 2.0 - 0.5
        cy = height / 2.0 - 0.5

        fx = cx / math.tan(fovWidth.radians() / 2.0)
        fy = cy / math.tan(fovHeight.radians() / 2.0)

        newCamIntrinsics = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]])

        self.setCalibrationFromIntrinsics(
            width, height, newCamIntrinsics, newDistCoeffs
        )

    def setCalibrationFromIntrinsics(
        self,
        width: int,
        height: int,
        newCamIntrinsics: np.ndarray,
        newDistCoeffs: np.ndarray,
    ) -> None:

        self.resWidth = width
        self.resHeight = height
        self.camIntrinsics = newCamIntrinsics
        self.distCoeffs = newDistCoeffs

        p = [
            Translation3d(
                1.0,
                Rotation3d(
                    0.0,
                    0.0,
                    (self.getPixelYaw(0) + Rotation2d(math.pi / 2.0)).radians(),
                ),
            ),
            Translation3d(
                1.0,
                Rotation3d(
                    0.0,
                    0.0,
                    (self.getPixelYaw(width) + Rotation2d(math.pi / 2.0)).radians(),
                ),
            ),
            Translation3d(
                1.0,
                Rotation3d(
                    0.0,
                    0.0,
                    (self.getPixelPitch(0) + Rotation2d(math.pi / 2.0)).radians(),
                ),
            ),
            Translation3d(
                1.0,
                Rotation3d(
                    0.0,
                    0.0,
                    (self.getPixelPitch(height) + Rotation2d(math.pi / 2.0)).radians(),
                ),
            ),
        ]

        self.viewplanes = []

        for i in p:
            self.viewplanes.append(np.array([i.X(), i.Y(), i.Z()]))

    def setCalibError(self, newAvgErrorPx: float, newErrorStdDevPx: float):
        self.avgErrorPx = newAvgErrorPx
        self.errorStdDevPx = newErrorStdDevPx

    def setFPS(self, fps: hertz):
        self.frameSpeed = max(1.0 / fps, self.exposureTime)

    def setExposureTime(self, newExposureTime: seconds):
        self.exposureTime = newExposureTime
        self.frameSpeed = max(self.frameSpeed, self.exposureTime)

    def setAvgLatency(self, newAvgLatency: seconds):
        self.vgLatency = newAvgLatency

    def setLatencyStdDev(self, newLatencyStdDev: seconds):
        self.latencyStdDev = newLatencyStdDev

    def getResWidth(self) -> int:
        return self.resWidth

    def getResHeight(self) -> int:
        return self.resHeight

    def getResArea(self) -> int:
        return self.resWidth * self.resHeight

    def getAspectRatio(self) -> float:
        return 1.0 * self.resWidth / self.resHeight

    def getIntrinsics(self) -> np.ndarray:
        return self.camIntrinsics

    def getDistCoeffs(self) -> np.ndarray:
        return self.distCoeffs

    def getFPS(self) -> hertz:
        return 1.0 / self.frameSpeed

    def getFrameSpeed(self) -> seconds:
        return self.frameSpeed

    def getExposureTime(self) -> seconds:
        return self.exposureTime

    def getAverageLatency(self) -> seconds:
        return self.avgLatency

    def getLatencyStdDev(self) -> seconds:
        return self.latencyStdDev

    def getContourAreaPercent(self, points: np.ndarray) -> float:
        return cv.contourArea(cv.convexHull(points)) / self.getResArea() * 100.0

    def getPixelYaw(self, pixelX: float) -> Rotation2d:
        fx = self.camIntrinsics[0, 0]
        cx = self.camIntrinsics[0, 2]
        xOffset = cx - pixelX
        return Rotation2d(fx, xOffset)

    def getPixelPitch(self, pixelY: float) -> Rotation2d:
        fy = self.camIntrinsics[1, 1]
        cy = self.camIntrinsics[1, 2]
        yOffset = cy - pixelY
        return Rotation2d(fy, -yOffset)

    def getPixelRot(self, point: cv.typing.Point2f) -> Rotation3d:
        return Rotation3d(
            0.0,
            self.getPixelPitch(point[1]).radians(),
            self.getPixelYaw(point[0]).radians(),
        )

    def getCorrectedPixelRot(self, point: cv.typing.Point2f) -> Rotation3d:
        fx = self.camIntrinsics[0, 0]
        cx = self.camIntrinsics[0, 2]
        xOffset = cx - point[0]

        fy = self.camIntrinsics[1, 1]
        cy = self.camIntrinsics[1, 2]
        yOffset = cy - point[1]

        yaw = Rotation2d(fx, xOffset)
        pitch = Rotation2d(fy / math.cos(math.atan(xOffset / fx)), -yOffset)
        return Rotation3d(0.0, pitch.radians(), yaw.radians())

    def getHorizFOV(self) -> Rotation2d:
        left = self.getPixelYaw(0)
        right = self.getPixelYaw(self.resWidth)
        return left - right

    def getVertFOV(self) -> Rotation2d:
        above = self.getPixelPitch(0)
        below = self.getPixelPitch(self.resHeight)
        return below - above

    def getDiagFOV(self) -> Rotation2d:
        return Rotation2d(
            math.hypot(self.getHorizFOV().radians(), self.getVertFOV().radians())
        )

    def getVisibleLine(
        self, camRt: RotTrlTransform3d, a: Translation3d, b: Translation3d
    ) -> typing.Tuple[float | None, float | None]:
        relA = camRt.applyTranslation(a)
        relB = camRt.applyTranslation(b)

        if relA.X() <= 0.0 and relB.X() <= 0.0:
            return (None, None)

        av = np.array([relA.X(), relA.Y(), relA.Z()])
        bv = np.array([relB.X(), relB.Y(), relB.Z()])
        abv = bv - av

        aVisible = True
        bVisible = True

        for normal in self.viewplanes:
            aVisibility = av.dot(normal)
            if aVisibility < 0:
                aVisible = False

            bVisibility = bv.dot(normal)
            if bVisibility < 0:
                bVisible = False
            if aVisibility <= 0 and bVisibility <= 0:
                return (None, None)

        if aVisible and bVisible:
            return (0.0, 1.0)

        intersections = [float("nan"), float("nan"), float("nan"), float("nan")]

        # Optionally 3x1 vector
        ipts: typing.List[np.ndarray | None] = [None, None, None, None]

        for i, normal in enumerate(self.viewplanes):
            a_projn = (av.dot(normal) / normal.dot(normal)) * normal

            if abs(abv.dot(normal)) < 1.0e-5:
                continue
            intersections[i] = a_projn.dot(a_projn) / -(abv.dot(a_projn))

            apv = intersections[i] * abv
            intersectpt = av + apv
            ipts[i] = intersectpt

            for j in range(1, len(self.viewplanes)):
                if j == 0:
                    continue
                oi = (i + j) % len(self.viewplanes)
                onormal = self.viewplanes[oi]
                if intersectpt.dot(onormal) < 0:
                    intersections[i] = float("nan")
                    ipts[i] = None
                    break

            if ipts[i] is None:
                continue

            for j in range(i - 1, 0 - 1):
                oipt = ipts[j]
                if not oipt:
                    continue

                diff = oipt - intersectpt
                if abs(diff).max() < 1e-4:
                    intersections[i] = float("nan")
                    ipts[i] = None
                    break

        inter1 = float("nan")
        inter2 = float("nan")
        for inter in intersections:
            if not math.isnan(inter):
                if math.isnan(inter1):
                    inter1 = inter
                else:
                    inter2 = inter

        if not math.isnan(inter2):
            max_ = max(inter1, inter2)
            min_ = min(inter1, inter2)
            if aVisible:
                min_ = 0
            if bVisible:
                max_ = 1
            return (min_, max_)
        elif not math.isnan(inter1):
            if aVisible:
                return (0, inter1)
            if bVisible:
                return (inter1, 1)
            return (inter1, None)
        else:
            return (None, None)

    def estPixelNoise(self, points: np.ndarray) -> np.ndarray:
        assert points.shape[1] == 1, points.shape
        assert points.shape[2] == 2, points.shape
        if self.avgErrorPx == 0 and self.errorStdDevPx == 0:
            return points

        noisyPts: list[list] = []
        for p in points:
            error = np.random.normal(self.avgErrorPx, self.errorStdDevPx, 1)[0]
            errorAngle = np.random.uniform(-math.pi, math.pi)
            noisyPts.append(
                [
                    [
                        float(p[0, 0] + error * math.cos(errorAngle)),
                        float(p[0, 1] + error * math.sin(errorAngle)),
                    ]
                ]
            )
        retval = np.array(noisyPts, dtype=np.float32)
        assert points.shape == retval.shape, retval
        return retval

    def estLatency(self) -> seconds:
        return max(
            float(np.random.normal(self.avgLatency, self.latencyStdDev, 1)[0]),
            0.0,
        )

    def estSecUntilNextFrame(self) -> seconds:
        return self.frameSpeed + max(0.0, self.estLatency() - self.frameSpeed)

    @classmethod
    def PERFECT_90DEG(cls) -> typing.Self:
        return cls()

    @classmethod
    def PI4_LIFECAM_320_240(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            320,
            240,
            newCamIntrinsics=np.array(
                [
                    [328.2733242048587, 0.0, 164.8190261141906],
                    [0.0, 318.0609794305216, 123.8633838438093],
                    [0.0, 0.0, 1.0],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.09957946553445934,
                        -0.9166265114485799,
                        0.0019519890627236526,
                        -0.0036071725380870333,
                        1.5627234622420942,
                        0,
                        0,
                        0,
                    ]
                ]
            ),
        )
        prop.setCalibError(0.21, 0.0124)
        prop.setFPS(30.0)
        prop.setAvgLatency(30.0e-3)
        prop.setLatencyStdDev(10.0e-3)
        return prop

    @classmethod
    def PI4_LIFECAM_640_480(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            640,
            480,
            newCamIntrinsics=np.array(
                [
                    [669.1428078983059, 0.0, 322.53377249329213],
                    [0.0, 646.9843137061716, 241.26567383784163],
                    [0.0, 0.0, 1.0],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.12788470750464645,
                        -1.2350335805796528,
                        0.0024990767286192732,
                        -0.0026958287600230705,
                        2.2951386729115537,
                        0,
                        0,
                        0,
                    ]
                ]
            ),
        )
        prop.setCalibError(0.26, 0.046)
        prop.setFPS(15.0)
        prop.setAvgLatency(65.0e-3)
        prop.setLatencyStdDev(15.0e-3)
        return prop

    @classmethod
    def LL2_640_480(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            640,
            480,
            newCamIntrinsics=np.array(
                [
                    [511.22843367007755, 0.0, 323.62049380211096],
                    [0.0, 514.5452336723849, 261.8827920543568],
                    [0.0, 0.0, 1.0],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.1917469998873756,
                        -0.5142936883324216,
                        0.012461562046896614,
                        0.0014084973492408186,
                        0.35160648971214437,
                        0,
                        0,
                        0,
                    ]
                ]
            ),
        )
        prop.setCalibError(0.25, 0.05)
        prop.setFPS(15.0)
        prop.setAvgLatency(35.0e-3)
        prop.setLatencyStdDev(8.0e-3)
        return prop

    @classmethod
    def LL2_960_720(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            960,
            720,
            newCamIntrinsics=np.array(
                [
                    [769.6873145148892, 0.0, 486.1096609458122],
                    [0.0, 773.8164483705323, 384.66071662358354],
                    [0.0, 0.0, 1.0],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.189462064814501,
                        -0.49903003669627627,
                        0.007468423590519429,
                        0.002496885298683693,
                        0.3443122090208624,
                        0,
                        0,
                        0,
                    ]
                ]
            ),
        )
        prop.setCalibError(0.35, 0.10)
        prop.setFPS(10.0)
        prop.setAvgLatency(50.0e-3)
        prop.setLatencyStdDev(15.0e-3)
        return prop

    @classmethod
    def LL2_1280_720(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            1280,
            720,
            newCamIntrinsics=np.array(
                [
                    [1011.3749416937393, 0.0, 645.4955139388737],
                    [0.0, 1008.5391755084075, 508.32877656020196],
                    [0.0, 0.0, 1.0],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.13730101577061535,
                        -0.2904345656989261,
                        8.32475714507539e-4,
                        -3.694397782014239e-4,
                        0.09487962227027584,
                        0,
                        0,
                        0,
                    ]
                ]
            ),
        )
        prop.setCalibError(0.37, 0.06)
        prop.setFPS(7.0)
        prop.setAvgLatency(60.0e-3)
        prop.setLatencyStdDev(20.0e-3)
        return prop

    @classmethod
    def OV9281_640_480(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            640,
            480,
            newCamIntrinsics=np.array(
                [
                    [627.1573807284262, 0, 307.79423851611824],
                    [0, 626.6621595938243, 219.02625533911998],
                    [0, 0, 1],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.054834081023049625,
                        -0.15994111706817074,
                        -0.0017587106009926158,
                        -0.0014671022483263552,
                        0.049742166267499596,
                        0,
                        0,
                        0,
                    ],
                ]
            ),
        )
        prop.setCalibError(0.25, 0.05)
        prop.setFPS(30.0)
        prop.setAvgLatency(60.0e-3)
        prop.setLatencyStdDev(20.0e-3)
        return prop

    @classmethod
    def OV9281_800_600(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            800,
            600,
            newCamIntrinsics=np.array(
                [
                    [783.9467259105329, 0, 384.7427981451478],
                    [0, 783.3276994922804, 273.7828191739],
                    [0, 0, 1],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.054834081023049625,
                        -0.15994111706817074,
                        -0.0017587106009926158,
                        -0.0014671022483263552,
                        0.049742166267499596,
                        0,
                        0,
                        0,
                    ],
                ]
            ),
        )
        prop.setCalibError(0.25, 0.05)
        prop.setFPS(25.0)
        prop.setAvgLatency(60.0e-3)
        prop.setLatencyStdDev(20.0e-3)
        return prop

    @classmethod
    def OV9281_1280_720(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            1280,
            720,
            newCamIntrinsics=np.array(
                [
                    [940.7360710926395, 0, 615.5884770322365],
                    [0, 939.9932393907364, 328.53938300868],
                    [0, 0, 1],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.054834081023049625,
                        -0.15994111706817074,
                        -0.0017587106009926158,
                        -0.0014671022483263552,
                        0.049742166267499596,
                        0,
                        0,
                        0,
                    ],
                ]
            ),
        )
        prop.setCalibError(0.25, 0.05)
        prop.setFPS(15.0)
        prop.setAvgLatency(60.0e-3)
        prop.setLatencyStdDev(20.0e-3)
        return prop

    @classmethod
    def OV9281_1920_1080(cls) -> typing.Self:
        prop = cls()
        prop.setCalibrationFromIntrinsics(
            1920,
            1080,
            newCamIntrinsics=np.array(
                [
                    [1411.1041066389591, 0, 923.3827155483548],
                    [0, 1409.9898590861046, 492.80907451301994],
                    [0, 0, 1],
                ]
            ),
            newDistCoeffs=np.array(
                [
                    [
                        0.054834081023049625,
                        -0.15994111706817074,
                        -0.0017587106009926158,
                        -0.0014671022483263552,
                        0.049742166267499596,
                        0,
                        0,
                        0,
                    ],
                ]
            ),
        )
        prop.setCalibError(0.25, 0.05)
        prop.setFPS(10.0)
        prop.setAvgLatency(60.0e-3)
        prop.setLatencyStdDev(20.0e-3)
        return prop