#!/usr/bin/env python3
# Standard library
import os
import sys
import json
from pathlib import Path

# External library
import cv2
import numpy as np
import torch

# Internal library
from ecos_core.base_model import BaseModel
from ecos_core.match_tool.match_tool_algo import MatchTool

FILE = Path(__file__).resolve()
FILE_DIR = 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 = "match_tool"

    # =====================================================================
    # 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

        if not hasattr(self.opt, "class_names"):
            self.opt.class_names = self._load_class_names()
        if not hasattr(self.opt, "visualization"):
            self.opt.visualization = True

        # Legacy field names (threshold/angle_range/min_area/subpixel) are
        # what the current production opt.json files use - kept as the
        # primary lookup, with the newer names as fallback aliases.
        self._max_pos = getattr(opt, "max_pos", 1)
        self._max_overlap = getattr(opt, "max_overlap", 0.2)
        self._score = getattr(opt, "threshold", getattr(opt, "score", 0.5))
        tolerance_angle = getattr(opt, "angle_range", getattr(opt, "tolerance_angle", 0.0))
        if isinstance(tolerance_angle, (list, tuple)):
            tolerance_angle = max(abs(tolerance_angle[0]), abs(tolerance_angle[1]))
        self._tolerance_angle = tolerance_angle
        self._min_reduce_area = getattr(opt, "min_area", getattr(opt, "min_reduce_area", 256))
        self._use_subpixel = getattr(opt, "subpixel", getattr(opt, "use_subpixel", False))
        self._stop_layer1 = getattr(opt, "stop_layer1", False)

        self._scale_range = tuple(getattr(opt, "scale_range", (1.0, 1.0)))
        self._scale_steps = getattr(opt, "scale_steps", 1)
        self._scale_prefilter_top_k = getattr(opt, "scale_prefilter_top_k", 2)

        self._match_tool = None
        self._pattern = None
        self._template_paths = []
        self._rois = [None]
        self._template_images = []
        self._refresh_pattern_state()

    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
        """
        self._match_tool = MatchTool(
            max_pos=self._max_pos,
            max_overlap=self._max_overlap,
            score=self._score,
            tolerance_angle=self._tolerance_angle,
            min_reduce_area=self._min_reduce_area,
            use_subpixel=self._use_subpixel,
            stop_layer1=self._stop_layer1,
        )
        return self

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

        def transform(image_path):
            return image_path, torch.tensor(0)

        return transform

    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".
        """
        items = option.items() if isinstance(option, dict) else vars(option).items()
        for key, value in 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 predict(self, image) -> list:
        """Run matching against every configured (template, ROI) combination.
        Args:
            image (str | np.ndarray): Image file path, or an already-loaded
                image as a NumPy array (grayscale or BGR).
        Returns:
            list[dict]: All matches found, sorted by score descending. Each
            entry has keys x, y, w, h, score, angle, scale, center, corners
            (bbox and center/corners are in full-image pixel coordinates).
        """
        self._refresh_pattern_state()
        if not self._template_images:
            raise ValueError("No template image(s) configured.")
        if self._match_tool is None:
            self.build_model()

        gray = self._to_gray(image)

        all_matches = []
        for template_image in self._template_images:
            for roi in self._rois:
                all_matches.extend(self._search_one(template_image, roi, gray))

        all_matches.sort(key=lambda m: m["score"], reverse=True)
        return all_matches

    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.
        """
        raw_image_mat = cv2.imread(raw_image_path)
        matches = self.predict(raw_image_mat)

        best = matches[0] if matches else None
        predict_index = 0 if best is not None else 1
        predict_class = self.opt.class_names[predict_index]
        boxes = [[best["x"], best["y"], best["w"], best["h"]]] if best else []
        confs = [best["score"]] if best else []

        if self.opt.visualization:
            self._draw_and_save(raw_image_mat, best, save_image_path)
        else:
            save_dir = os.path.dirname(os.path.abspath(save_image_path))
            os.makedirs(save_dir, exist_ok=True)
            cv2.imwrite(save_image_path, raw_image_mat)

        extra_info = {}
        if best is not None:
            extra_info = {
                "center_x": best["center"][0],
                "center_y": best["center"][1],
                "angle": best["angle"],
                "scale": best["scale"],
            }

        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["score"] 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

    # =====================================================================
    # PRIVATE METHODS
    # =====================================================================
    @staticmethod
    def _to_gray(image) -> np.ndarray:
        """Accept an image path or an already-loaded array (grayscale or
        BGR) and return a grayscale np.ndarray.
        """
        if isinstance(image, str):
            path = image
            image = cv2.imread(path)
            if image is None:
                raise FileNotFoundError(f"Cannot read image: {path}")
        if image.ndim == 3:
            return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        return image

    def _load_class_names(self) -> list:
        path = os.path.join(FILE_DIR, "class.names")
        if os.path.exists(path):
            with open(path, "r") as f:
                names = [line.strip() for line in f if line.strip()]
            if names:
                return names
        return ["OK", "NG"]

    def _refresh_pattern_state(self) -> None:
        """Re-resolve template paths/images and ROIs for the current
        opt.pattern. A no-op when the pattern hasn't changed since the last
        resolution - called before every predict() so a pattern swapped in
        after construction (directly on self.opt, or via update_opt) is
        picked up instead of staying stuck on the pattern seen at __init__.
        """
        pattern = str(getattr(self.opt, "pattern", "default"))
        print(f"Current pattern: {pattern}")
        if pattern == self._pattern:
            return
        self._pattern = pattern
        self._template_paths = self._resolve_template_paths(self.opt)
        self._rois = self._resolve_rois(self.opt)
        self._template_images = [
            img for img in (
                cv2.imread(p, cv2.IMREAD_GRAYSCALE) for p in self._template_paths
            ) if img is not None
        ]

    def _resolve_template_paths(self, opt) -> list:
        template_path = getattr(opt, "template_path", None)
        if template_path:
            if isinstance(template_path, (list, tuple)):
                return list(template_path)
            return [template_path]

        template_name = getattr(opt, "template_name", None)
        folder = getattr(opt, "template_folder", os.path.join(FILE_DIR, "template_folder"))
        if isinstance(template_name, dict):
            # pattern-keyed: value can itself be a filename or a list of
            # filenames, so it still needs normalizing below.
            template_name = template_name.get(self._pattern, template_name.get("default"))

        if isinstance(template_name, (list, tuple)):
            names = list(template_name)
        elif isinstance(template_name, str):
            names = [template_name]
        else:
            names = []
        return [n if os.path.isabs(n) else os.path.join(folder, n) for n in names]

    def _resolve_rois(self, opt) -> list:
        raw = getattr(opt, "roi", None)
        roi_path = getattr(opt, "roi_path", None)
        if raw is None and roi_path and os.path.exists(roi_path):
            with open(roi_path, "r") as f:
                raw = json.load(f)
        if raw is None:
            return [None]

        if isinstance(raw, dict) and "x" in raw and "y" in raw:
            return [_normalize_roi(raw)]
        if isinstance(raw, dict):
            # pattern-keyed: {"1": [...], "2": [...], ...}
            raw = raw.get(self._pattern, raw.get("default"))
            if raw is None:
                return [None]

        if isinstance(raw, (list, tuple)) and len(raw) > 0 and isinstance(raw[0], (list, tuple)):
            # multiple ROI regions for this pattern - search every one
            return [_normalize_roi(r) for r in raw]
        return [_normalize_roi(raw)]

    @staticmethod
    def _apply_roi(image: np.ndarray, roi) -> tuple:
        if roi is None:
            return image, (0, 0)
        x, y, w, h = roi["x"], roi["y"], roi["w"], roi["h"]
        return image[y: y + h, x: x + w], (x, y)

    def _scale_candidates(self) -> np.ndarray:
        if self._scale_steps <= 1:
            return np.array([1.0])
        return np.linspace(self._scale_range[0], self._scale_range[1], self._scale_steps)

    @staticmethod
    def _scale_template(template: np.ndarray, scale: float) -> np.ndarray:
        if abs(scale - 1.0) < 1e-6:
            return template
        h, w = template.shape[:2]
        new_w = max(1, int(round(w * scale)))
        new_h = max(1, int(round(h * scale)))
        return cv2.resize(template, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

    def _rank_scales_coarse(self, template_image: np.ndarray, image: np.ndarray, candidates: np.ndarray) -> np.ndarray:
        """Cheaply rank scale candidates before running the full MatchTool
        pipeline on each one, by reusing MatchTool's own algorithm with
        `stop_layer1=True` (skips only the final, most expensive refine
        level) - not an ad-hoc heuristic, so its ranking stays consistent
        with the real search.
        Args:
            template_image (np.ndarray): The reference (unscaled) template.
            image (np.ndarray): The ROI-cropped source image.
            candidates (np.ndarray): Scale factors to rank.
        Returns:
            np.ndarray: The top `scale_prefilter_top_k` candidates, best first.
        """
        original_stop_layer1 = self._match_tool.stop_layer1
        self._match_tool.stop_layer1 = True
        try:
            scored = []
            for scale in candidates:
                scaled_template = self._scale_template(template_image, scale)
                self._match_tool.crop_xyxy = None
                self._match_tool.load_source(image)
                self._match_tool.load_template(scaled_template)
                results = self._match_tool.match()
                best_score = max((r.d_match_score for r in results), default=-1.0)
                scored.append((scale, best_score))
        finally:
            self._match_tool.stop_layer1 = original_stop_layer1

        scored.sort(key=lambda t: t[1], reverse=True)
        top_k = scored[: max(1, self._scale_prefilter_top_k)]
        return np.array([s for s, _ in top_k])

    def _search_one(self, template_image: np.ndarray, roi, gray_image: np.ndarray) -> list:
        """Run the (optionally scale-searched) MatchTool pipeline for one
        (template, ROI) combination.
        Returns:
            list[dict]: Matches in full-image coordinates.
        """
        image, offset = self._apply_roi(gray_image, roi)

        candidates = self._scale_candidates()
        if len(candidates) > 1:
            candidates = self._rank_scales_coarse(template_image, image, candidates)

        matches = []
        for scale in candidates:
            scaled_template = self._scale_template(template_image, scale)
            self._match_tool.crop_xyxy = None
            self._match_tool.load_source(image)
            self._match_tool.load_template(scaled_template)
            for r in self._match_tool.match():
                matches.append(self._to_match_dict(r, offset, scale))
        return matches

    @staticmethod
    def _to_match_dict(result, offset: tuple, scale: float) -> dict:
        ox, oy = offset
        corners = np.array(
            [result.pt_lt, result.pt_rt, result.pt_rb, result.pt_lb], dtype=np.float32
        )
        corners[:, 0] += ox
        corners[:, 1] += oy
        x, y, w, h = cv2.boundingRect(corners.astype(np.int32))
        center = (float(result.pt_center[0]) + ox, float(result.pt_center[1]) + oy)
        return {
            "x": x, "y": y, "w": w, "h": h,
            "score": float(result.d_match_score),
            "angle": float(result.d_matched_angle),
            "scale": float(scale),
            "center": center,
            "corners": corners,
        }

    def _draw_and_save(self, image: np.ndarray, best: dict | None, save_path: str) -> None:
        output = image.copy()
        if best is not None:
            pts = best["corners"].astype(int)
            cv2.polylines(output, [pts], True, (0, 255, 0), 2, cv2.LINE_AA)
            cx, cy = int(best["center"][0]), int(best["center"][1])
            cv2.drawMarker(output, (cx, cy), (0, 0, 255), cv2.MARKER_CROSS, 10, 2)
            cv2.putText(
                output,
                f"{best['score']:.3f}, {best['angle']:.1f}deg, x{best['scale']:.2f}",
                (best["x"], best["y"] - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, cv2.LINE_AA,
            )

        for roi in self._rois:
            if roi is not None:
                cv2.rectangle(
                    output, (roi["x"], roi["y"]),
                    (roi["x"] + roi["w"], roi["y"] + roi["h"]), (255, 0, 0), 2,
                )

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