"""
Feature Matching & Classification Module
Implements SURF, ORB, BRISK, and Hybrid matching according to
the methodology of Electronics 2023, 12(17), 3640.
"""

import os
import cv2
import numpy as np
import threading
from typing import Dict, List, Tuple, Optional, Any
from identity import decide_identity


class LocalFeatureExtractor:
    """Encapsulates individual OpenCV local feature algorithms."""

    def __init__(self, method: str = "orb"):
        self.method = method.lower()
        self.detector = self._create_detector(self.method)
        self.norm_type = self._get_norm_type(self.method)

    def _create_detector(self, method: str):
        if method == "surf":
            # Check if SURF is supported in the current OpenCV build
            if hasattr(cv2, "xfeatures2d") and hasattr(cv2.xfeatures2d, "SURF_create"):
                try:
                    return cv2.xfeatures2d.SURF_create(hessianThreshold=400, nOctaves=4, nOctaveLayers=3)
                except Exception as e:
                    print(f"[Warning] SURF initialization failed: {e}. Falling back to SIFT.")
            else:
                print("[Info] SURF not available in this OpenCV build. Using SIFT as high-grade floating point descriptor.")
            return cv2.SIFT_create(nfeatures=1000)

        elif method == "orb":
            return cv2.ORB_create(
                nfeatures=1500,
                scaleFactor=1.2,
                nlevels=8,
                edgeThreshold=31,
                firstLevel=0,
                WTA_K=2,
                scoreType=cv2.ORB_HARRIS_SCORE,
                patchSize=31,
                fastThreshold=20
            )

        elif method == "brisk":
            return cv2.BRISK_create(thresh=30, octaves=3, patternScale=1.0)

        elif method == "sift":
            return cv2.SIFT_create(nfeatures=1000)

        else:
            raise ValueError(f"Unsupported feature extractor method: {method}")

    def _get_norm_type(self, method: str):
        # Binary descriptors use Hamming distance; floating-point descriptors use L2 distance
        if method in ["orb", "brisk"]:
            return cv2.NORM_HAMMING
        return cv2.NORM_L2

    def detect_and_compute(self, image_gray: np.ndarray) -> Tuple[List[cv2.KeyPoint], Optional[np.ndarray]]:
        """Detects keypoints and computes their descriptors."""
        if image_gray is None or image_gray.size == 0:
            return [], None
        return self.detector.detectAndCompute(image_gray, None)


class TemplateFeatureDatabase:
    """Precomputes and caches local feature descriptors for all stock templates."""

    def __init__(self, stock_dir: str = "stock"):
        self.stock_dir = stock_dir
        self.templates: Dict[str, Dict[str, Any]] = {}
        self.extractors = {
            "surf": LocalFeatureExtractor("surf"),
            "orb": LocalFeatureExtractor("orb"),
            "brisk": LocalFeatureExtractor("brisk")
        }
        self.load_and_cache_templates()

    def load_and_cache_templates(self):
        """Loads all SKU images and extracts their multi-feature descriptors."""
        if not os.path.exists(self.stock_dir):
            return

        image_files = [
            f for f in os.listdir(self.stock_dir)
            if f.lower().endswith((".jpeg", ".jpg", ".png"))
        ]

        print(f"[FeatureDatabase] Pre-caching {len(image_files)} stock SKU templates...")
        for filename in image_files:
            sku_id = os.path.splitext(filename)[0]
            filepath = os.path.join(self.stock_dir, filename)
            img = cv2.imread(filepath)
            if img is None:
                continue

            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            features = {}
            for name, extractor in self.extractors.items():
                kp, des = extractor.detect_and_compute(gray)
                features[name] = {"keypoints": kp, "descriptors": des}

            self.templates[sku_id] = {
                "sku_id": sku_id,
                "filepath": filepath,
                "image": img,
                "gray": gray,
                "features": features
            }
        print(f"[FeatureDatabase] Successfully cached {len(self.templates)} templates.")


class ProductFeatureMatcher:
    """
    Main product classification engine.
    Matches a query cropped bounding box against cached SKU templates using
    SURF, ORB, BRISK, or the Paper's Hybrid combination.
    """

    def __init__(self, stock_dir: str = "stock"):
        self.db = TemplateFeatureDatabase(stock_dir)
        self.extractors = self.db.extractors
        self._thread_local = threading.local()

    def thread_extractors(self) -> Dict[str, "LocalFeatureExtractor"]:
        """One extractor set per thread: OpenCV Feature2D objects are not safe to
        call concurrently from several threads."""
        extractors = getattr(self._thread_local, "extractors", None)
        if extractors is None:
            extractors = {
                name: LocalFeatureExtractor(name) for name in ("surf", "orb", "brisk")
            }
            self._thread_local.extractors = extractors
        return extractors

    def match_descriptors(
        self,
        des_query: Optional[np.ndarray],
        des_template: Optional[np.ndarray],
        norm_type: int,
        ratio_thresh: float = 0.75
    ) -> List[cv2.DMatch]:
        """Performs k-NN matching with Lowe's ratio test."""
        if des_query is None or des_template is None:
            return []
        if len(des_query) < 2 or len(des_template) < 2:
            return []

        bf = cv2.BFMatcher(norm_type, crossCheck=False)
        try:
            raw_matches = bf.knnMatch(des_query, des_template, k=2)
        except Exception:
            return []

        good_matches = []
        for match_pair in raw_matches:
            if len(match_pair) == 2:
                m, n = match_pair
                if m.distance < ratio_thresh * n.distance:
                    good_matches.append(m)
        return good_matches

    def compute_geometric_inliers(
        self,
        kp_query: List[cv2.KeyPoint],
        kp_template: List[cv2.KeyPoint],
        good_matches: List[cv2.DMatch]
    ) -> Tuple[int, List[cv2.DMatch], Optional[np.ndarray]]:
        """Applies RANSAC Homography to verify geometric consistency."""
        if len(good_matches) < 4:
            return len(good_matches), good_matches, None

        src_pts = np.float32([kp_query[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp_template[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

        try:
            H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
            if mask is None:
                return len(good_matches), good_matches, None

            inlier_matches = [good_matches[i] for i in range(len(good_matches)) if mask[i][0] == 1]
            return len(inlier_matches), inlier_matches, H
        except Exception:
            return len(good_matches), good_matches, None

    def score_crop(
        self,
        crop_bgr: np.ndarray,
        method: str = "hybrid"
    ) -> Tuple[Dict[str, float], Dict[str, Dict[str, Any]]]:
        """
        Scores a product crop against every SKU template.
        Returns (scores, details); raises ValueError for an unsupported method and
        returns empty results for an empty crop.
        Methods supported: 'surf', 'orb', 'brisk', 'hybrid'.
        """
        if crop_bgr is None or crop_bgr.size == 0:
            return {}, {}

        crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
        method = method.lower()
        extractors = self.thread_extractors()

        scores: Dict[str, float] = {}
        details: Dict[str, Dict[str, Any]] = {}

        if method in ["surf", "orb", "brisk"]:
            extractor = extractors.get(method)
            if extractor is None:
                extractor = LocalFeatureExtractor(method)
            kp_query, des_query = extractor.detect_and_compute(crop_gray)

            for sku_id, tmpl in self.db.templates.items():
                kp_tmpl = tmpl["features"][method]["keypoints"]
                des_tmpl = tmpl["features"][method]["descriptors"]

                good_matches = self.match_descriptors(
                    des_query, des_tmpl, extractor.norm_type, ratio_thresh=0.75
                )
                inlier_count, inliers, _ = self.compute_geometric_inliers(
                    kp_query, kp_tmpl, good_matches
                )

                # Score based on geometric inliers & match ratio
                denom = max(1, min(len(kp_query), len(kp_tmpl)))
                inlier_ratio = inlier_count / denom
                score = inlier_count * 2.0 + inlier_ratio * 100.0

                scores[sku_id] = score
                details[sku_id] = {
                    "inliers": inlier_count,
                    "inlier_matches": inliers,
                    "total_good": len(good_matches),
                    "kp_query": kp_query,
                    "kp_tmpl": kp_tmpl,
                    "score": score
                }

        elif method == "hybrid":
            # Paper's hybrid strategy combining SURF, ORB, and BRISK
            sub_methods = ["surf", "orb", "brisk"]
            weights = {"surf": 0.45, "orb": 0.25, "brisk": 0.30}

            # Pre-compute features for query crop across all 3 extractors
            query_features = {}
            for sm in sub_methods:
                kp_q, des_q = extractors[sm].detect_and_compute(crop_gray)
                query_features[sm] = (kp_q, des_q)

            for sku_id, tmpl in self.db.templates.items():
                composite_score = 0.0
                sku_sub_details = {}

                for sm in sub_methods:
                    kp_q, des_q = query_features[sm]
                    kp_t = tmpl["features"][sm]["keypoints"]
                    des_t = tmpl["features"][sm]["descriptors"]
                    norm_t = extractors[sm].norm_type

                    matches = self.match_descriptors(des_q, des_t, norm_t, ratio_thresh=0.75)
                    inliers, inlier_list, _ = self.compute_geometric_inliers(kp_q, kp_t, matches)

                    denom = max(1, min(len(kp_q), len(kp_t)))
                    sub_score = inliers * 1.5 + (inliers / denom) * 80.0
                    composite_score += weights[sm] * sub_score

                    sku_sub_details[sm] = {
                        "inliers": inliers,
                        "matches": len(matches),
                        "kp_q_count": len(kp_q),
                        "kp_t_count": len(kp_t),
                        "sub_score": round(sub_score, 2)
                    }

                scores[sku_id] = composite_score
                details[sku_id] = {
                    "score": composite_score,
                    "sub_details": sku_sub_details,
                    "kp_query": query_features["surf"][0],
                    "kp_tmpl": tmpl["features"]["surf"]["keypoints"]
                }
        else:
            raise ValueError(f"Unknown method: {method}")

        return scores, details

    def score_crop_fast(
        self,
        crop_bgr: np.ndarray,
        shortlist: int = 5
    ) -> Tuple[Dict[str, float], Dict[str, Dict[str, Any]]]:
        """
        ORB-shortlist then hybrid-score the top candidates only.

        Same scoring formula as 'hybrid', but the expensive SURF/ORB/BRISK
        descriptor matching runs on `shortlist` templates instead of all of them.
        Every other template scores 0.0, so decide_identity() still ranks the
        whole catalog. This is a live-mode speed path; validate it against
        'hybrid' with bench_engine.py --compare before trusting it.
        """
        if crop_bgr is None or crop_bgr.size == 0:
            return {}, {}
        crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
        extractors = self.thread_extractors()

        orb = extractors["orb"]
        kp_orb, des_orb = orb.detect_and_compute(crop_gray)
        orb_scores: Dict[str, float] = {}
        for sku_id, tmpl in self.db.templates.items():
            kp_t = tmpl["features"]["orb"]["keypoints"]
            des_t = tmpl["features"]["orb"]["descriptors"]
            matches = self.match_descriptors(des_orb, des_t, orb.norm_type, ratio_thresh=0.75)
            inliers, _, _ = self.compute_geometric_inliers(kp_orb, kp_t, matches)
            denom = max(1, min(len(kp_orb), len(kp_t)))
            orb_scores[sku_id] = inliers * 1.5 + (inliers / denom) * 80.0

        candidates = sorted(orb_scores.items(), key=lambda kv: kv[1], reverse=True)[:shortlist]

        scores: Dict[str, float] = {sku_id: 0.0 for sku_id in self.db.templates}
        details: Dict[str, Dict[str, Any]] = {}
        sub_methods = ["surf", "orb", "brisk"]
        weights = {"surf": 0.45, "orb": 0.25, "brisk": 0.30}
        query_features = {}
        for sm in sub_methods:
            kp_q, des_q = extractors[sm].detect_and_compute(crop_gray)
            query_features[sm] = (kp_q, des_q)

        for sku_id, _ in candidates:
            tmpl = self.db.templates[sku_id]
            composite = 0.0
            sku_sub_details = {}
            for sm in sub_methods:
                kp_q, des_q = query_features[sm]
                kp_t = tmpl["features"][sm]["keypoints"]
                des_t = tmpl["features"][sm]["descriptors"]
                matches = self.match_descriptors(des_q, des_t, extractors[sm].norm_type, ratio_thresh=0.75)
                inliers, _, _ = self.compute_geometric_inliers(kp_q, kp_t, matches)
                denom = max(1, min(len(kp_q), len(kp_t)))
                sub_score = inliers * 1.5 + (inliers / denom) * 80.0
                composite += weights[sm] * sub_score
                sku_sub_details[sm] = {
                    "inliers": inliers,
                    "matches": len(matches),
                    "kp_q_count": len(kp_q),
                    "kp_t_count": len(kp_t),
                    "sub_score": round(sub_score, 2),
                }
            scores[sku_id] = composite
            details[sku_id] = {
                "score": composite,
                "sub_details": sku_sub_details,
                "kp_query": query_features["surf"][0],
                "kp_tmpl": tmpl["features"]["surf"]["keypoints"],
            }
        return scores, details

    def classify_crop(
        self,
        crop_bgr: np.ndarray,
        method: str = "hybrid"
    ) -> Dict[str, Any]:
        """
        Classifies a product crop against all SKU templates.
        Methods supported: 'surf', 'orb', 'brisk', 'hybrid', 'hybrid_fast'.
        """
        if crop_bgr is None or crop_bgr.size == 0:
            return {"status": "error", "message": "Empty crop image"}

        method = method.lower()
        if method == "hybrid_fast":
            scores, details = self.score_crop_fast(crop_bgr)
        elif method in ["surf", "orb", "brisk", "hybrid"]:
            scores, details = self.score_crop(crop_bgr, method=method)
        else:
            return {"status": "error", "message": f"Unknown method: {method}"}

        ranked = decide_identity(scores, details=details, method=method)
        return ranked

    def generate_match_visualization(
        self,
        crop_bgr: np.ndarray,
        matched_sku: str,
        method: str = "surf",
        max_matches: int = 30
    ) -> Optional[np.ndarray]:
        """Draws visual matching lines between crop and matched template."""
        if matched_sku not in self.db.templates:
            return None

        tmpl = self.db.templates[matched_sku]
        crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
        eval_method = method if method in ["surf", "orb", "brisk"] else "surf"

        extractor = self.thread_extractors()[eval_method]
        kp_q, des_q = extractor.detect_and_compute(crop_gray)
        kp_t = tmpl["features"][eval_method]["keypoints"]
        des_t = tmpl["features"][eval_method]["descriptors"]

        matches = self.match_descriptors(des_q, des_t, extractor.norm_type, ratio_thresh=0.75)
        _, inliers, _ = self.compute_geometric_inliers(kp_q, kp_t, matches)

        matches_to_draw = inliers[:max_matches] if inliers else matches[:max_matches]

        vis = cv2.drawMatches(
            crop_bgr, kp_q,
            tmpl["image"], kp_t,
            matches_to_draw, None,
            matchColor=(0, 255, 0),
            singlePointColor=(0, 0, 255),
            flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS
        )
        return vis