# Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license

import time

import cv2
import numpy as np

from boxmot.motion.cmc.base_cmc import BaseCMC
from boxmot.utils import BOXMOT
from boxmot.utils import logger as LOGGER


class ECC(BaseCMC):
    def __init__(
        self,
        warp_mode: int = cv2.MOTION_TRANSLATION,
        eps: float = 1e-5,
        max_iter: int = 100,
        scale: float = 0.15,
        align: bool = False,
        grayscale: bool = True,
    ) -> None:
        """Compute the warp matrix from src to dst.

        Parameters
        ----------
        warp_mode: opencv flag
            translation: cv2.MOTION_TRANSLATION
            rotated and shifted: cv2.MOTION_EUCLIDEAN
            affine(shift,rotated,shear): cv2.MOTION_AFFINE
            homography(3d): cv2.MOTION_HOMOGRAPHY
        eps: float
            the threshold of the increment in the correlation coefficient between two iterations
        max_iter: int
            the number of iterations.
        scale: float or [int, int]
            scale_ratio: float
            scale_size: [W, H]
        align: bool
            whether to warp affine or perspective transforms to the source image
        grayscale: bool
            whether to transform 3 channel RGB to single channel grayscale for faster computations

        Returns
        -------
        warp matrix : ndarray
            Returns the warp matrix from src to dst.
            if motion models is homography, the warp matrix will be 3x3, otherwise 2x3
        src_aligned: ndarray
            aligned source image of gray
        """
        self.align = align
        self.grayscale = grayscale
        self.scale = scale
        self.warp_mode = warp_mode
        self.termination_criteria = (
            cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
            max_iter,
            eps,
        )
        self.prev_img = None

    def apply(self, img: np.ndarray, dets: np.ndarray = None) -> np.ndarray:
        """Apply sparse optical flow to compute the warp matrix.

        Parameters:
            img (ndarray): The input image.
            dets: Description of dets parameter.

        Returns:
            ndarray: The warp matrix from the source to the destination.
                If the motion model is homography, the warp matrix will be 3x3; otherwise, it will be 2x3.
        """

        if self.warp_mode == cv2.MOTION_HOMOGRAPHY:
            warp_matrix = np.eye(3, 3, dtype=np.float32)
        else:
            warp_matrix = np.eye(2, 3, dtype=np.float32)

        if self.prev_img is None:
            self.prev_img = self.preprocess(img)
            return warp_matrix

        img = self.preprocess(img)

        try:
            (ret_val, warp_matrix) = cv2.findTransformECC(
                self.prev_img,  # already processed
                img,
                warp_matrix,
                self.warp_mode,
                self.termination_criteria,
                None,
                1,
            )
        except cv2.error as e:
            # error 7 is StsNoConv, according to https://docs.opencv.org/3.4/d1/d0d/namespacecv_1_1Error.html
            if e.code == cv2.Error.StsNoConv:
                LOGGER.warning(
                    f"Affine matrix could not be generated: {e}. Returning identity"
                )
                return warp_matrix
            else:  # other error codes
                raise

        # upscale warp matrix to original images size
        if self.scale < 1:
            warp_matrix[0, 2] /= self.scale
            warp_matrix[1, 2] /= self.scale

        if self.align:
            h, w = self.prev_img.shape
            if self.warp_mode == cv2.MOTION_HOMOGRAPHY:
                # Use warpPerspective for Homography
                self.prev_img_aligned = cv2.warpPerspective(
                    self.prev_img, warp_matrix, (w, h), flags=cv2.INTER_LINEAR
                )
            else:
                # Use warpAffine for Translation, Euclidean and Affine
                self.prev_img_aligned = cv2.warpAffine(
                    self.prev_img, warp_matrix, (w, h), flags=cv2.INTER_LINEAR
                )
        else:
            self.prev_img_aligned = None

        self.prev_img = img

        return warp_matrix  # , prev_img_aligned


def main():
    ecc = ECC(scale=0.5, align=True, grayscale=True)
    curr_img = cv2.imread("assets/MOT17-mini/train/MOT17-2-FRCNN/img1/000005.jpg")
    prev_img = cv2.imread("assets/MOT17-mini/train/MOT17-2-FRCNN/img1/000001.jpg")

    warp_matrix = ecc.apply(prev_img, None)
    warp_matrix = ecc.apply(curr_img, None)

    start = time.process_time()
    for i in range(0, 100):
        warp_matrix = ecc.apply(prev_img, None)
        warp_matrix = ecc.apply(curr_img, None)
    end = time.process_time()
    print("Total time", end - start)
    print(warp_matrix)

    if ecc.prev_img_aligned is not None:
        curr_img = ecc.preprocess(curr_img)
        prev_img = ecc.preprocess(prev_img)
        weighted_img = cv2.addWeighted(curr_img, 0.5, ecc.prev_img_aligned, 0.5, 0)
        cv2.imshow("prev_img_aligned", weighted_img)
        cv2.waitKey(0)
        cv2.imwrite(str(BOXMOT / "motion/cmc/ecc_aligned.jpg"), weighted_img)


if __name__ == "__main__":
    main()
