import cv2 as cv
import numpy as np

class Detector:
    def __init__(self, target_image, min_matches=4, angle_tolerance=30, length_ratio_tolerance=2.5, nfeatures=30000):
        self.orb = cv.ORB_create(nfeatures=nfeatures)
        self.target_features, self.target_des = self.orb.detectAndCompute(target_image, None)
        self.target_image_h, self.target_image_w = target_image.shape[:2]
        self.min_matches = min_matches
        self.angle_tolerance = angle_tolerance
        self.length_ratio_tolerance = length_ratio_tolerance

    def valid_rectangle(self, points):
        def angle_between(p1, p2, p3):
            v1 = p1 - p2
            v2 = p3 - p2
            angle = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
            return np.degrees(angle)

        if len(points) != 4:
            return False

        angles = [
            angle_between(points[i], points[(i + 1) % 4], points[(i + 2) % 4])
            for i in range(4)
        ]
        if not all(abs(angle - 90) <= self.angle_tolerance for angle in angles):
            print(f'Not a valid rectangle, angle criterion {[angle for angle in angles]}')
            return False

        side_lengths = [
            np.linalg.norm(points[i] - points[(i + 1) % 4]) for i in range(4)
        ]
        max_length = max(side_lengths)
        min_length = min(side_lengths)
        if max_length / min_length > self.length_ratio_tolerance:
            print(f'Not a valid rectangle, length criterion')
            return False
        
        return True

    def match_feature_find_object(self, train_img):
        features2, des2 = self.orb.detectAndCompute(train_img, None)

        bf = cv.BFMatcher(cv.NORM_HAMMING)
        matches = bf.knnMatch(self.target_des, des2, k=2)

        good = []
        good_without_lists = []
        matches = [match for match in matches if len(match) == 2]
        for m, n in matches:
            if m.distance < 0.8 * n.distance:
                good.append([m])
                good_without_lists.append(m)

        if len(good) >= self.min_matches:
            src_pts = np.float32([self.target_features[m.queryIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
            dst_pts = np.float32([features2[m.trainIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)

            M, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)

            h, w = self.target_image_h, self.target_image_w
            pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)

            try:
                dst = cv.perspectiveTransform(pts, M)
            except cv.error as e:
                print(f"Perspective transform failed: {e}")
                return None, train_img

            rectangle_points = np.squeeze(dst)
            if not self.valid_rectangle(rectangle_points):
                return None, train_img

            train_img = cv.polylines(train_img, [np.int32(dst)], True, (0, 255, 0), 2, cv.LINE_AA)
            x, y, w, h = cv.boundingRect(np.int32(dst))
            return (x, y, w, h), train_img

        else:
            print(f"Not enough good matches are found - {len(good)}/{self.min_matches}")
            return None, train_img