import cv2
import numpy as np

class Postprocessor:
    @staticmethod
    def calculate_iou(boxA, boxB):
        # Determine the (x, y)-coordinates of the intersection rectangle
        xA = max(boxA['x'], boxB['x'])
        yA = max(boxA['y'], boxB['y'])
        xB = min(boxA['x'] + boxA['w'], boxB['x'] + boxB['w'])
        yB = min(boxA['y'] + boxA['h'], boxB['y'] + boxB['h'])
        
        # Compute the area of intersection rectangle
        interArea = max(0, xB - xA) * max(0, yB - yA)
        
        # Compute the area of both bounding boxes
        boxAArea = boxA['w'] * boxA['h']
        boxBArea = boxB['w'] * boxB['h']
        
        # Compute the Intersection over Union (IoU)
        unionArea = boxAArea + boxBArea - interArea
        if unionArea == 0:
            return 0
        return interArea / unionArea

    def nms(self, candidates, iou_threshold=0.3):
        """
        Filters overlapping candidate boxes using Intersection-over-Union (IoU)
        and confidence score. Keeps the highest scoring box when overlaps occur.
        """
        if not candidates:
            return []
            
        # Sort candidates by score in descending order
        sorted_cands = sorted(candidates, key=lambda x: x['score'], reverse=True)
        keep = []
        
        for cand in sorted_cands:
            should_keep = True
            for selected in keep:
                iou = self.calculate_iou(cand, selected)
                if iou > iou_threshold:
                    should_keep = False
                    break
            if should_keep:
                keep.append(cand)
                
        return keep

    @staticmethod
    def cosine_similarity(v1, v2):
        """
        Computes cosine similarity between two 1D numerical vectors.
        """
        dot = np.dot(v1, v2)
        norm_v1 = np.linalg.norm(v1)
        norm_v2 = np.linalg.norm(v2)
        if norm_v1 == 0 or norm_v2 == 0:
            return 0.0
        return float(dot / (norm_v1 * norm_v2))

    @staticmethod
    def compute_geometric_features(img):
        """
        Extracts a generalized 5D Geometric Primitive Feature Vector representing the 
        amount of straight lines, loops/circles, sharp corners, and pixel density.
        Fully scale-invariant and rotation-invariant.
        """
        if img is None or img.size == 0:
            return np.zeros(5, dtype=np.float32)
            
        h, w = img.shape[:2]
        area = float(h * w)
        
        # 1. Active Pixel Density
        density = np.count_nonzero(img) / area
        
        # 2. Corner/Junction Count (Shi-Tomasi corner detection)
        # Binarize to ensure goodFeaturesToTrack works perfectly
        img_bin = (img > 127).astype(np.uint8)
        corners = cv2.goodFeaturesToTrack(img_bin, maxCorners=100, qualityLevel=0.05, minDistance=3)
        num_corners = len(corners) if corners is not None else 0
        norm_corners = num_corners / (area / 1000.0) if area > 0 else 0
        
        # 3. Straight Lines Count & Total Length (Hough Line Transform)
        lines = cv2.HoughLinesP(img_bin, rho=1, theta=np.pi/180, threshold=8, minLineLength=6, maxLineGap=3)
        num_lines = len(lines) if lines is not None else 0
        
        total_line_len = 0.0
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                total_line_len += np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
                
        norm_lines = num_lines / (area / 1000.0) if area > 0 else 0
        norm_line_len = total_line_len / area if area > 0 else 0
        
        # 4. Circular Loops Count (Contour circularity analysis)
        contours, _ = cv2.findContours(img_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        num_circles = 0
        for c in contours:
            perimeter = cv2.arcLength(c, True)
            c_area = cv2.contourArea(c)
            if perimeter > 8:
                circularity = (4 * np.pi * c_area) / (perimeter ** 2)
                if circularity > 0.60:  # highly circular closed loop or circle
                    num_circles += 1
                    
        norm_circles = num_circles / (area / 1000.0) if area > 0 else 0
        
        return np.array([density, norm_corners, norm_lines, norm_line_len, norm_circles], dtype=np.float32)

    def verify_generalized_topology(self, cand_proc_crop, tpl_proc_crop):
        """
        GENERAL-PURPOSE ZERO-SHOT SHAPE VALIDATOR (Anti-Confusion Engine)
        Verifies that the candidate image region matches the topological structure
        of the query template using 2D spatial correlation, 1D projection profiles,
        and rotation-invariant Geometric Primitive Vectors (lines, loops, corners)
        on the CORE REGION (middle 76%).
        
        Fully robust to any arbitrary, unseen symbol during blind tests.
        """
        if cand_proc_crop is None or cand_proc_crop.size == 0:
            return 0.0
            
        # 1. Resize candidate crop to match template crop dimensions exactly
        th, tw = tpl_proc_crop.shape[:2]
        cand_resized = cv2.resize(cand_proc_crop, (tw, th), interpolation=cv2.INTER_AREA)
        
        # 2. Extract Core Region (middle 76%)
        # Discards the left, right, top, and bottom borders where leads merge into wires.
        x_pad = int(tw * 0.12)
        y_pad = int(th * 0.12)
        
        cand_core = cand_resized[y_pad : th - y_pad, x_pad : tw - x_pad]
        tpl_core = tpl_proc_crop[y_pad : th - y_pad, x_pad : tw - x_pad]
        
        if cand_core.size == 0 or tpl_core.size == 0:
            cand_core = cand_resized
            tpl_core = tpl_proc_crop
            
        # 3. 2D Cosine Similarity (Pixel-by-pixel structural overlap) on the core region
        v_cand = cand_core.flatten().astype(np.float32)
        v_tpl = tpl_core.flatten().astype(np.float32)
        sim_2d = self.cosine_similarity(v_cand, v_tpl)
        
        # 4. 1D Projection Profiles on the core region
        cand_proj_x = np.sum(cand_core, axis=0).astype(np.float32)
        tpl_proj_x = np.sum(tpl_core, axis=0).astype(np.float32)
        
        cand_proj_y = np.sum(cand_core, axis=1).astype(np.float32)
        tpl_proj_y = np.sum(tpl_core, axis=1).astype(np.float32)
        
        sim_x = self.cosine_similarity(cand_proj_x, tpl_proj_x)
        sim_y = self.cosine_similarity(cand_proj_y, tpl_proj_y)
        sim_1d = (sim_x + sim_y) / 2.0
        
        # 5. Extract and Match Rotation-Invariant Geometric Primitives
        geom_cand = self.compute_geometric_features(cand_core)
        geom_tpl = self.compute_geometric_features(tpl_core)
        sim_geom = self.cosine_similarity(geom_cand, geom_tpl)
        
        # Fused overall similarity (50% 2D structural, 20% 1D profiles, 30% Geometric primitives)
        overall_sim = (sim_2d * 0.50) + (sim_1d * 0.20) + (sim_geom * 0.30)
        
        # 6. Active Pixel Density Check on the core region (relaxed slightly for noise tolerance)
        density_cand = np.count_nonzero(cand_core) / cand_core.size
        density_tpl = np.count_nonzero(tpl_core) / tpl_core.size
        
        if density_tpl > 0:
            density_ratio = density_cand / density_tpl
            if density_ratio < 0.35 or density_ratio > 2.8:
                overall_sim *= 0.30  # penalize severely!
                
        # 7. Strict Valley/Gap Check on the core region (mean-based for discretization tolerance)
        cw = tpl_core.shape[1]
        ch = tpl_core.shape[0]
        mid_x = cw // 2
        tpl_mid_x_area = tpl_proj_x[max(0, mid_x - 3) : min(cw, mid_x + 4)]
        cand_mid_x_area = cand_proj_x[max(0, mid_x - 3) : min(cw, mid_x + 4)]
        
        tpl_has_gap = np.mean(tpl_mid_x_area) < np.mean(tpl_proj_x) * 0.15
        if tpl_has_gap:
            cand_has_gap = np.mean(cand_mid_x_area) < np.mean(cand_proj_x) * 0.25
            if not cand_has_gap:
                overall_sim *= 0.2
                
        # Same check for Y axis
        mid_y = ch // 2
        tpl_mid_y_area = tpl_proj_y[max(0, mid_y - 3) : min(ch, mid_y + 4)]
        cand_mid_y_area = cand_proj_y[max(0, mid_y - 3) : min(ch, mid_y + 4)]
        
        tpl_has_gap_y = np.mean(tpl_mid_y_area) < np.mean(tpl_proj_y) * 0.15
        if tpl_has_gap_y:
            cand_has_gap_y = np.mean(cand_mid_y_area) < np.mean(cand_proj_y) * 0.25
            if not cand_has_gap_y:
                overall_sim *= 0.2
                
        return overall_sim