#!/usr/bin/env python3
##
# @file feature_matching.py
#
# @brief Model to classify PEM images using feature-based matching.
#
# @section author_doxygen_example Author(s)
# - Created by Tran Viet Thanh on 2026-03-19.
#
# Copyright (c) 2026 HACHIX, Inc. All rights reserved.

# Standard library
import os
import sys
import json
from pathlib import Path

# External library
import torch
import cv2
import numpy as np

# Internal library
from ecos_core.base_model import BaseModel

FILE = Path(__file__).resolve()
FILE_DIR = os.path.join(os.path.dirname(__file__))
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))


def _normalize_roi(roi):
    """Accept a {"x","y","w","h"} dict or an [x1,y1,x2,y2] list/tuple.
    Args:
        roi: ROI in either supported format, or None.
    Returns:
        dict | None: {"x","y","w","h"}, or None if roi is None.
    """
    if roi is None:
        return None
    if isinstance(roi, dict):
        return {"x": int(roi["x"]), "y": int(roi["y"]), "w": int(roi["w"]), "h": int(roi["h"])}
    x1, y1, x2, y2 = roi
    return {"x": int(x1), "y": int(y1), "w": int(x2) - int(x1), "h": int(y2) - int(y1)}


class Model(BaseModel):
    NAME = "feature_matching"

    # =====================================================================
    # PUBLIC METHODS
    # =====================================================================
    def __init__(self, opt):
        """Initialize the model with options
        Args:
            opt (Namespace/Dict): options for the model
        """
        super().__init__(opt)
        self.opt = opt



        self._template_path = getattr(opt, "template_path", None)
        self._roi_path = getattr(opt, "roi_path", None)
        self._roi = self._resolve_roi(opt)
        self._method = getattr(opt, "method", "AKAZE")
        self._max_distance = getattr(opt, "max_distance", 0.7)
        self._best_matched_points = getattr(opt, "best_matched_points", 10)
        self._fast_threshold = getattr(opt, "fast_threshold", 5)
        self._edge_threshold = getattr(opt, "edge_threshold", 10)
        self._max_matches = getattr(opt, "max_matches", 1)
        self._last_angle = None

        self._detector, self._matcher = self._build_detector(
            self._method, self._fast_threshold, self._edge_threshold
        )

        self._template_image = None
        if self._template_path and os.path.exists(self._template_path):
            self._template_image = cv2.imread(self._template_path)

    def __call__(self, x):
        """Inference function
        Args:
            x: input data
        Returns:
             Inference results
        """
        return None

    def build_model(self):
        """Build the model
        Returns:
             The built model
        """
        return self

    def get_transform(self):
        """Get the data transformation function
        Returns:
             The data transformation function
        """

        def transform(image_path):
            """Data transformation function
            Args:
                image_path: path to the input image
            Returns:
                The transformed image and the raw image
            """
            return image_path, torch.tensor(0)

        return transform

    def get_selected_pattern(self):
        """Get the selected pattern from the options
        Returns:
            The selected pattern
        """
        selected_pattern_id = self.opt.selected_pattern
        for pattern in self.opt.patterns:
            if pattern.get("id") == selected_pattern_id:
                return pattern

    def update_opt(self, option, output_path="./opt.json"):
        """Update config
        Args:
            option (Namespace/Dict): option to update
            output_path (str, optional): path to save opt json. Defaults to "./opt.json".
        """
        for key, value in vars(option).items():
            if hasattr(self.opt, key):
                setattr(self.opt, key, value)

        with open(output_path, "w") as f:
            json.dump(vars(self.opt), f)

        print("Configuration updated successfully!")

    def crop_and_save_template(
        self,
        image_path: str,
        save_path: str,
    ) -> bool:
        """Select a template region and save it.

        If OpenCV GUI is available, opens an interactive ROI selector:
          - Click and drag to draw a selection box.
          - Press ENTER or SPACE to confirm.
          - Press C to redraw. Press ESC to abort.

        If GUI is not available (e.g. headless OpenCV), falls back to a
        CLI prompt asking for x, y, w, h coordinates.

        Args:
            image_path (str): Path to the source image.
            save_path (str): File path to save the cropped template image.
        Returns:
            bool: True if the template was saved, False if the user aborted.
        """
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not read image from: {image_path}")

        roi = self._select_roi(image)
        if roi is None:
            return False

        x, y, w, h = roi
        cropped = image[y: y + h, x: x + w]

        save_dir = os.path.dirname(save_path)
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)

        cv2.imwrite(save_path, cropped)

        self._template_image = cropped
        self._template_path = save_path

        print(f"Template saved to: {save_path}")
        return True

    def process_predictions(
        self,
        net_output,
        raw_image_path,
        raw_image_mat,
        image_transform,
        save_image_path,
    ):
        """Process the predictions and save the results.
        Args:
            net_output: the output from the model inference (unused)
            raw_image_path: path to the raw input image
            raw_image_mat: the raw input image in matrix form (unused, re-read from disk)
            image_transform: the transformed input image (unused)
            save_image_path: path to save the processed image
        Returns:
            str: save_image_path, once the annotated image and its JSON
            sidecar ({save_image_path}.json) have been written.
        """
        if self._template_image is None:
            raise ValueError("Template image is not set.")
        raw_image_mat = cv2.imread(raw_image_path)
        image, roi_offset = self._apply_roi(raw_image_mat)
        template_kp, template_desc = self._compute_keypoints(
            self._template_image
        )
        matches = self._find_matches(image, template_kp, template_desc)
        matched_boxes = self._estimate_poses(matches, image, roi_offset)
        matched_boxes = self._finalize_matches(matched_boxes)
        self._draw_and_save(raw_image_mat, matched_boxes, save_image_path)

        best = matched_boxes[0] if matched_boxes else None
        class_names = getattr(self.opt, "class_names", ["OK", "NG"])
        predict_index = 0 if best is not None else 1
        predict_class = class_names[predict_index]

        boxes = []
        confs = []
        extra_info = {}
        if best is not None:
            x, y, pts, score, angle = best
            bx, by, bw, bh = cv2.boundingRect(pts)
            boxes = [[bx, by, bw, bh]]
            confs = [score]
            extra_info = {"center_x": x, "center_y": y, "angle": angle}

        save_txt_path = "{}.json".format(os.path.splitext(save_image_path)[0])
        with open(save_txt_path, "w") as result_file:
            json.dump({
                "save_image_path": save_image_path,
                "score": {"image_score": best[3] if best else 0.0},
                "predict_index": predict_index,
                "class_name": predict_class,
                "bbox": boxes,
                "confidence": confs,
                "extra_info": extra_info,
            }, result_file)

        return save_image_path

    def reset_state(self) -> None:
        """Clear the angle EMA state so the next prediction is not biased by
        a previous, unrelated image. Call this between independent test
        cases (e.g. in a benchmark harness) since `_smooth_angle` otherwise
        carries state across calls.
        """
        self._last_angle = None

    def load_template(self, template_path: str) -> None:
        """Load the template image from the specified path.
        Args:
            template_path (str): Path to the template image file.
        """
        if not os.path.exists(template_path):
            raise FileNotFoundError(f"Template file not found: {template_path}")
        self._template_image = cv2.imread(template_path)
        self._template_path = template_path

    def load_roi(self, roi_path: str) -> None:
        """Load the ROI from a *_roi.json file.
        Args:
            roi_path (str): Path to the ROI json file.
        """
        if not os.path.exists(roi_path):
            raise FileNotFoundError(f"ROI file not found: {roi_path}")
        with open(roi_path, "r") as f:
            self._roi = _normalize_roi(json.load(f))
        self._roi_path = roi_path

    def select_and_save_roi(
        self, image_path: str, roi_path: str = None
    ) -> bool:
        """Interactively select a ROI from the image and save it to a
        dedicated *_roi.json file.

        If OpenCV GUI is available, opens an interactive ROI selector.
        Falls back to CLI input if GUI is unavailable.

        Args:
            image_path (str): Path to the source image to select ROI from.
            roi_path (str): Path to save the ROI json file. Defaults to
                self._roi_path if set, otherwise derives from image_path
                as <image_stem>_roi.json beside the image.
        Returns:
            bool: True if the ROI was saved, False if the user aborted.
        """
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not read image from: {image_path}")

        roi = self._select_roi(image)
        if roi is None:
            return False

        x, y, w, h = roi
        self._roi = {"x": x, "y": y, "w": w, "h": h}

        if roi_path is None:
            if self._roi_path:
                roi_path = self._roi_path
            else:
                stem = Path(image_path).stem
                roi_path = os.path.join(
                    os.path.dirname(image_path), f"{stem}_roi.json"
                )

        self._roi_path = roi_path
        save_dir = os.path.dirname(roi_path)
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)

        with open(roi_path, "w") as f:
            json.dump(self._roi, f, indent=4)

        print(f"ROI saved to {roi_path}: x={x}, y={y}, w={w}, h={h}")
        return True

    # =====================================================================
    # PRIVATE METHODS
    # =====================================================================
    def _select_roi(self, image: cv2.Mat) -> tuple | None:
        """Select a ROI interactively via GUI, or fall back to CLI input.
        Args:
            image (cv2.Mat): The source image to select from.
        Returns:
            tuple[int, int, int, int] | None: (x, y, w, h) or None if aborted.
        """
        try:
            window_name = "Select Template Region"
            cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
            display = image.copy()
            for i, line in enumerate([
                "Drag to select ROI",
                "ENTER/SPACE: confirm | C: redraw | ESC: abort",
            ]):
                cv2.putText(
                    display, line, (10, 25 + i * 22),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA,
                )
            roi = cv2.selectROI(
                window_name, display, showCrosshair=True, fromCenter=False
            )
            cv2.destroyWindow(window_name)
            x, y, w, h = roi
            if w == 0 or h == 0:
                print("Selection aborted — no template saved.")
                return None
            return x, y, w, h
        except cv2.error:
            return self._select_roi_cli(image)

    def _select_roi_cli(self, image: cv2.Mat) -> tuple | None:
        """Fallback CLI prompt for ROI selection when GUI is unavailable.
        Args:
            image (cv2.Mat): The source image (used only to show dimensions).
        Returns:
            tuple[int, int, int, int] | None: (x, y, w, h) or None if aborted.
        """
        h_img, w_img = image.shape[:2]
        print(f"Image size: {w_img}x{h_img} (width x height)")
        print("Enter ROI coordinates (or press ENTER to abort):")
        try:
            x = int(input("  x (left edge): "))
            y = int(input("  y (top edge): "))
            w = int(input("  w (width): "))
            h = int(input("  h (height): "))
        except (ValueError, EOFError):
            print("Invalid input — no template saved.")
            return None
        if w <= 0 or h <= 0:
            print("Width and height must be positive — no template saved.")
            return None
        return x, y, w, h

    def _draw_and_save(
        self,
        image: cv2.Mat,
        matches: list,
        save_path: str,
    ) -> None:
        """Draw detection results on the image and save to disk.
        Args:
            image (cv2.Mat): The original full image to draw on.
            matches (list): Matches as (x, y, pts, score, angle).
            save_path (str): Path to save the annotated image.
        """
        output = image.copy()
        for i, (x, y, pts, score, angle) in enumerate(matches):
            # Draw the true detected quadrilateral (pts) directly. Do NOT
            # reconstruct a box from cv2.boundingRect(pts) + angle: the
            # bounding rect of a rotated quad is larger than the quad
            # itself, so re-rotating it produces a box that no longer
            # coincides with what was actually matched.
            quad = pts.reshape(-1, 1, 2).astype(int)
            cv2.polylines(output, [quad], True, (0, 255, 0), 2, cv2.LINE_AA)
            center_x, center_y = int(round(x)), int(round(y))
            cv2.circle(output, (center_x, center_y), 5, (0, 0, 255), -1)
            label_x, label_y = pts[:, 0].min(), pts[:, 1].min()
            cv2.putText(
                output,
                f"#{i + 1}: {score:.3f}, {angle:.1f}deg",
                (int(label_x), int(label_y) - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 255, 0),
                1,
                cv2.LINE_AA,
            )

        if self._roi is not None:
            rx, ry = self._roi["x"], self._roi["y"]
            rw, rh = self._roi["w"], self._roi["h"]
            cv2.rectangle(output, (rx, ry), (rx + rw, ry + rh), (255, 0, 0), 2)

        save_dir = os.path.dirname(save_path)
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)
        cv2.imwrite(save_path, output)

    def _resolve_roi(self, opt) -> dict | None:
        """Resolve the search ROI, preferring an inline `opt.roi` over the
        `opt.roi_path` file (same precedence as MatchTool).
        Args:
            opt (Namespace/Dict): options for the model.
        Returns:
            dict | None: {"x","y","w","h"}, or None for full-frame search.
        """
        raw = getattr(opt, "roi", None)
        if raw is None and self._roi_path and os.path.exists(self._roi_path):
            with open(self._roi_path, "r") as f:
                raw = json.load(f)
        return _normalize_roi(raw)

    def _apply_roi(self, image: cv2.Mat) -> tuple:
        """Crop the image to the ROI and return the offset.
        Args:
            image (cv2.Mat): The input image.
        Returns:
            tuple[cv2.Mat, tuple[int, int]]: Cropped image and (x, y) offset.
        """
        if self._roi is None:
            return image, (0, 0)

        x, y, w, h = (
            self._roi["x"],
            self._roi["y"],
            self._roi["w"],
            self._roi["h"],
        )
        return image[y: y + h, x: x + w], (x, y)

    def _compute_keypoints(
        self, image: cv2.Mat
    ) -> tuple[list, np.ndarray]:
        """Detect keypoints and compute descriptors for an image.
        Args:
            image (cv2.Mat): Input image.
        Returns:
            tuple[list, np.ndarray]: Keypoints and descriptors.
        """
        return self._detector.detectAndCompute(image, None)

    def _find_matches(
        self,
        image: cv2.Mat,
        template_kp: list,
        template_desc: np.ndarray,
    ) -> list:
        """Match keypoints between the target image and the template.
        Args:
            image (cv2.Mat): Target image to search in.
            template_kp (list): Keypoints from the template.
            template_desc (np.ndarray): Descriptors from the template.
        Returns:
            list[tuple]: Each entry is (image_kp, good_matches) for one
            detection pass.
        """
        results = []
        for _ in range(self._max_matches):
            image_kp, image_desc = self._compute_keypoints(image)
            good_matches = self._filter_matches(image_desc, template_desc)
            results.append((image_kp, good_matches))
        return results

    def _filter_matches(
        self, desc1: np.ndarray, desc2: np.ndarray
    ) -> list:
        """Filter raw descriptor matches using method-specific strategy.
        Args:
            desc1 (np.ndarray): Descriptors from the target image.
            desc2 (np.ndarray): Descriptors from the template.
        Returns:
            list: Filtered list of good matches. Empty when AKAZE doesn't
            clear `best_matched_points` - too few points to trust a pose
            estimate, so this is treated as "no match" rather than an error.
        Raises:
            ValueError: If an unknown method is specified.
        """
        if self._method == "SIFT":
            raw = self._matcher.knnMatch(desc1, desc2, k=2)
            return [
                m for m, n in raw
                if m.distance < self._max_distance * n.distance
            ]

        if self._method == "ORB":
            raw = self._matcher.match(desc1, desc2)
            return sorted(raw, key=lambda m: m.distance)[
                : self._best_matched_points
            ]

        if self._method == "AKAZE":
            raw = self._matcher.knnMatch(desc1, desc2, k=2)
            good = [
                m for m, n in raw
                if m.distance < self._max_distance * n.distance
            ]
            if len(good) < self._best_matched_points:
                return []
            return good

        raise ValueError(f"Unknown method: {self._method}")

    def _estimate_poses(
        self,
        match_results: list,
        image: cv2.Mat,
        roi_offset: tuple,
    ) -> list:
        """Estimate affine pose for each set of keypoint matches.
        Args:
            match_results (list): Output of _find_matches — list of
                (image_kp, good_matches) pairs.
            image (cv2.Mat): Target image (used only for template shape ref).
            roi_offset (tuple[int, int]): (x, y) to shift corners back to
                full-image coordinates.
        Returns:
            list[tuple[float, float, np.ndarray, float, float]]:
                Each entry is (x, y, pts, score, angle).
        """
        h_t, w_t = self._template_image.shape[:2]
        template_kp, _ = self._compute_keypoints(self._template_image)
        corners = np.float32([
            [0,   0],
            [w_t, 0],
            [w_t, h_t],
            [0,   h_t],
        ]).reshape(-1, 1, 2)

        boxes = []
        for image_kp, good_matches in match_results:
            box = self._estimate_single_pose(
                image_kp, template_kp, good_matches, corners, roi_offset
            )
            if box is not None:
                boxes.append(box)

        return boxes

    def _estimate_single_pose(
        self,
        image_kp: list,
        template_kp: list,
        good_matches: list,
        corners: np.ndarray,
        roi_offset: tuple,
    ) -> tuple | None:
        """Estimate affine transform for a single match set and return a box.
        Args:
            image_kp (list): Keypoints from the target image.
            template_kp (list): Keypoints from the template.
            good_matches (list): Filtered matches between image and template.
            corners (np.ndarray): Template corner points shaped (-1, 1, 2).
            roi_offset (tuple[int, int]): (x, y) offset from ROI cropping.
        Returns:
            tuple[float, float, np.ndarray, float, float] | None:
                (x, y, pts, score, angle) or None if estimation fails.
        """
        pts_image = np.float32(
            [image_kp[m.queryIdx].pt for m in good_matches]
        )
        pts_template = np.float32(
            [template_kp[m.trainIdx].pt for m in good_matches]
        )

        try:
            M, inliers = cv2.estimateAffinePartial2D(
                pts_template,
                pts_image,
                cv2.RANSAC,
                ransacReprojThreshold=2.0,
                maxIters=5000,
                confidence=0.9950,
            )
        except cv2.error:
            print("Affine estimation failed — skipping.")
            return None

        transformed = cv2.transform(corners, M)
        offset = np.array(roi_offset, dtype=np.float32)
        transformed[:, 0, :] += offset

        score = len(inliers) / len(good_matches)
        angle = self._smooth_angle(np.degrees(np.arctan2(M[1, 0], M[0, 0])))
        pts = np.round(transformed.reshape(-1, 2)).astype(int)
        x, y = transformed.reshape(-1, 2).mean(axis=0)

        return (float(x), float(y), pts, float(score), float(angle))

    def _finalize_matches(self, boxes: list) -> list:
        """Trim to max_matches, sorted by score descending.
        Args:
            boxes (list): List of (x, y, pts, score, angle) tuples.
        Returns:
            list: Top-k matches sorted by score.
        """
        boxes = sorted(boxes, key=lambda b: b[3], reverse=True)
        if self._max_matches:
            boxes = boxes[: self._max_matches]
        return boxes

    def _smooth_angle(self, angle: float, alpha: float = 0.2) -> float:
        """Reduce rotation noise using exponential moving average.
        Args:
            angle (float): Current angle in degrees.
            alpha (float): Smoothing factor (0 = no update, 1 = no smoothing).
        Returns:
            float: Smoothed angle.
        """
        if self._last_angle is None:
            self._last_angle = angle
        else:
            self._last_angle = (1 - alpha) * self._last_angle + alpha * angle
        return self._last_angle

    @staticmethod
    def _build_detector(
        method: str, fast_threshold: int, edge_threshold: int
    ) -> tuple:
        """Instantiate the feature detector and matcher for the given method.
        Args:
            method (str): One of "SIFT", "ORB", or "AKAZE".
            fast_threshold (int): FAST threshold (ORB only).
            edge_threshold (int): Edge threshold (ORB only).
        Returns:
            tuple[cv2.Feature2D, cv2.DescriptorMatcher]: Detector and matcher.
        Raises:
            ValueError: If an unknown method name is provided.
        """
        if method == "SIFT":
            return cv2.SIFT_create(), cv2.BFMatcher()

        if method == "ORB":
            detector = cv2.ORB_create(
                fastThreshold=fast_threshold,
                edgeThreshold=edge_threshold,
            )
            return detector, cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

        if method == "AKAZE":
            return (
                cv2.AKAZE_create(),
                cv2.DescriptorMatcher_create(
                    cv2.DescriptorMatcher_BRUTEFORCE
                ),
            )

        raise ValueError(f"Unknown method: {method}")