#!/usr/bin/env python3
"""Circle Finder - reusable ECOS Core module for detecting circles.

A thin ``BaseModel`` wrapper (same contract as ``match_tool`` / ``paddle_ocr``)
around the pure ``geometry`` / ``detectors`` library. It:

1. reads config from ``opt.json`` (method(s), radius band, ROI, ...),
2. preprocesses (channel select + brightness/contrast + blur),
3. runs one or more registered detector strategies (with automatic
    coarse-to-fine seeding for the ``radial`` method),
4. re-scores every candidate with the common ``edge_support`` metric,
5. filters (radius band, min score), de-duplicates (NMS), keeps the top N,
6. optionally refines the winners to sub-pixel accuracy,
7. and, via ``process_predictions``, writes an annotated image + JSON sidecar.

Use it as a library too::

    from ecos_core.circle_finder.circle_finder import Model
    model = Model(opt); model.build_model()
    circles = model.predict(image_bgr)        # -> list[Circle], best first
"""
import os
import sys
import json
from pathlib import Path

import cv2
import numpy as np
import torch

from ecos_core.base_model import BaseModel
from ecos_core.circle_finder import preprocess as pp
from ecos_core.circle_finder.geometry import Circle, edge_support_score, refine_center_subpixel_radial
from ecos_core.circle_finder.detectors import DetectContext, build_detector

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 {"x","y","w","h"} or [x1,y1,x2,y2]; return {"x","y","w","h"} or 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 = "circle_finder"

    # =====================================================================
    # PUBLIC API
    # =====================================================================
    def __init__(self, opt):
        """Args:
            opt (Namespace/EasyDict/Dict): configuration (see README / opt.json).
        """
        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
        self._detectors = None

    def __call__(self, x):
        return None

    def get_transform(self):
        """BaseModel hook; this module reads images directly, so it is a no-op."""
        def transform(image_path):
            return image_path, torch.tensor(0)
        return transform

    def build_model(self):
        """Instantiate the configured detector strategies. Returns self."""
        methods = self._method_list()
        self._detectors = []
        for name in methods:
            params = self._get(f"methods", {}) or {}
            self._detectors.append((name, build_detector(name, params.get(name, {}))))
        return self

    def predict(self, image, roi=None, seed=None) -> list:
        """Detect circles in ``image``, across one or more ROIs.

        Args:
            image (str | np.ndarray): path, or a loaded gray/BGR array.
            roi (dict | list | None, optional): override ROI(s); defaults to
                ``opt.roi``. Accepts:
                - ``None`` → full frame;
                - a single ``{"x","y","w","h"}`` or ``[x1,y1,x2,y2]`` → one ROI;
                - a **list** of either → several ROIs, each searched
                  independently (up to ``max_circles`` circles per ROI).
                Results are returned in full-image coordinates.
            seed (Circle, optional): explicit seed for the ``radial`` method.
                Honoured only when there is a single ROI; with several ROIs each
                crop is auto-seeded.

        Returns:
            list[Circle]: results grouped by ROI (in config order); within a
            ROI, best (highest score) first. Each circle's ``meta["roi_index"]``
            records which ROI it came from.
        """
        if self._detectors is None:
            self.build_model()

        bgr_full = pp.to_bgr(image)
        gray_full = pp.to_gray(image)

        rois = self._resolve_rois(roi if roi is not None else self._get("roi"))

        single = len(rois) == 1
        results = []
        for ri, roi_n in enumerate(rois):
            ox, oy = 0, 0
            gray, bgr = gray_full, bgr_full
            if roi_n is not None:
                ox, oy = roi_n["x"], roi_n["y"]
                gray = gray_full[oy: oy + roi_n["h"], ox: ox + roi_n["w"]]
                bgr = bgr_full[oy: oy + roi_n["h"], ox: ox + roi_n["w"]] if bgr_full is not None else None
                if gray.size == 0:
                    continue
            # Radius band, resolved per ROI: diameter_mm / diameter_tol / r_min /
            # r_max may be a scalar (same band for every ROI) or a list (entry i
            # sizes ROI i) - so ROIs holding differently-sized holes each get
            # their own band. calibration / scale stay image-wide.
            r_min_i, r_max_i = pp.radius_band(
                diameter_mm=self._per_roi(self._get("diameter_mm"), ri),
                calibration=self._get("calibration", 0.0),
                tol=self._per_roi(self._get("diameter_tol", 0.3), ri),
                r_min=self._per_roi(self._get("r_min"), ri),
                r_max=self._per_roi(self._get("r_max"), ri),
                scale=self._get("scale", 1.0),
            )
            roi_seed = seed if single else None
            for c in self._detect_in_crop(gray, bgr, r_min_i, r_max_i, roi_seed):
                c.meta["roi_index"] = ri
                results.append(c.shifted(ox, oy))
        return results

    def _detect_in_crop(self, gray_crop, bgr_crop, r_min, r_max, seed) -> list:
        """Run the full detect pipeline on one already-cropped region. Returns
        crop-local circles, best first, capped at ``max_circles``. Shared by
        every ROI so the single-ROI and multi-ROI paths behave identically."""
        # Preprocess the working channel once and share across detectors.
        work = pp.select_channel(bgr_crop, self._get("channel", "gray")) if bgr_crop is not None else gray_crop
        work = pp.adjust_brightness_contrast(work, self._get("bc_min", 0), self._get("bc_max", 0))
        work = pp.blur(work, self._get("blur_ksize", 0))

        ctx = DetectContext(gray=work, bgr=bgr_crop, r_min=r_min, r_max=r_max, seed=seed)

        candidates = []
        for name, det in self._detectors:
            if name == "radial" and ctx.seed is None:
                ctx.seed = self._make_seed(det, ctx)
            candidates.extend(det.detect(ctx))

        # Common re-scoring so a union of methods is ranked on one scale.
        candidates = self._rescore(work, candidates)
        candidates = [c for c in candidates if c.score >= self._get("min_score", 0.0)]
        candidates = self._nms(candidates, self._get("nms_dist_frac", 0.5))

        if self._get("subpixel", False):
            candidates = [self._refine(work, c) for c in candidates]

        candidates.sort(key=lambda c: c.score, reverse=True)
        return candidates[: int(self._get("max_circles", 1))]

    def process_predictions(self, net_output, raw_image_path, raw_image_mat,
                            image_transform, save_image_path):
        """ECOS framework hook: run predict(), write the annotated image and a
        ``{save_image_path}.json`` sidecar. Verdict (predict_index=0=OK):
        - default: at least ``opt.min_count`` circles found in total;
        - when ``opt.require_per_roi`` is true: at least one circle in *every*
          configured ROI (the "one hole per ROI" pattern).
        """
        raw = cv2.imread(raw_image_path)
        circles = self.predict(raw)

        rois = self._resolve_rois(self._get("roi"))
        if self._get("require_per_roi", False):
            hit = {c.meta.get("roi_index") for c in circles}
            ok = all(i in hit for i in range(len(rois)))
        else:
            ok = len(circles) >= int(self._get("min_count", 1))
        predict_index = 0 if ok else 1
        predict_class = self.opt.class_names[predict_index]

        calib = self._get("calibration", 0.0)
        circle_json = []
        for c in circles:
            entry = {"roi_index": c.meta.get("roi_index", 0),
                     "cx": round(c.cx, 2), "cy": round(c.cy, 2), "r": round(c.r, 2),
                     "score": round(c.score, 3), "method": c.method}
            if calib and calib > 0:
                entry["diameter_mm"] = round(2 * c.r * calib, 4)
            circle_json.append(entry)

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

        save_txt_path = "{}.json".format(os.path.splitext(save_image_path)[0])
        with open(save_txt_path, "w") as f:
            json.dump({
                "save_image_path": save_image_path,
                "score": {"image_score": circles[0].score if circles else 0.0},
                "predict_index": predict_index,
                "class_name": predict_class,
                "num_rois": len(rois),
                "num_circles": len(circles),
                "circles": circle_json,
            }, f, indent=2)
        return save_image_path

    def update_opt(self, option, output_path="./opt.json"):
        """Update config from a Namespace/Dict and persist to JSON."""
        items = option.items() if isinstance(option, dict) else vars(option).items()
        for key, value in items:
            setattr(self.opt, key, value)
        with open(output_path, "w") as f:
            json.dump(self._opt_as_dict(), f, indent=2)
        self._detectors = None  # force rebuild with new params
        print("Configuration updated successfully!")

    # =====================================================================
    # PRIVATE
    # =====================================================================
    def _get(self, key, default=None):
        return getattr(self.opt, key, default) if not isinstance(self.opt, dict) else self.opt.get(key, default)

    def _opt_as_dict(self):
        return dict(self.opt) if isinstance(self.opt, dict) else vars(self.opt)

    def _method_list(self):
        m = self._get("method", "contour")
        if isinstance(m, str):
            return [m]
        return list(m)

    @staticmethod
    def _per_roi(value, i):
        """Pick the radius-band value for ROI ``i`` from a scalar (broadcast to
        every ROI) or a list (entry ``i`` sizes ROI ``i``). A 1-element list
        broadcasts; a list shorter than ``i`` clamps to its last entry (so a
        partial list degrades gracefully instead of raising). ``None`` and plain
        scalars pass straight through.
        """
        if isinstance(value, (list, tuple)):
            if len(value) == 0:
                return None
            return value[min(i, len(value) - 1)]
        return value

    @staticmethod
    def _resolve_rois(raw) -> list:
        """Normalise the roi config into a list of ROIs (``[None]`` = full frame).

        Accepts: ``None``; one ``{"x","y","w","h"}`` or ``[x1,y1,x2,y2]``; or a
        list of either. A bare 4-number list is treated as a single ``xyxy`` ROI,
        so ``[[..],[..]]`` is the only way to express several ROIs (each ROI is a
        dict or its own 4-list) - this keeps the single-ROI config unchanged.
        """
        if raw is None:
            return [None]
        if isinstance(raw, dict):
            return [_normalize_roi(raw)]
        if isinstance(raw, (list, tuple)):
            if len(raw) == 0:
                return [None]
            if len(raw) == 4 and all(isinstance(v, (int, float)) for v in raw):
                return [_normalize_roi(raw)]      # a single [x1,y1,x2,y2]
            return [_normalize_roi(r) for r in raw if r is not None]
        return [None]

    def _make_seed(self, radial_det, ctx: DetectContext):
        """Coarse-to-fine: seed the radial detector with a cheap first pass
        (contour by default) so it can refine from a real starting circle."""
        seed_method = radial_det.p.get("seed_method")
        if not seed_method:
            return None
        seed_params = radial_det.p.get("seed_params") or {}
        try:
            det = build_detector(seed_method, seed_params)
        except KeyError:
            return None
        found = det.detect(ctx)
        if not found:
            return None
        found = self._rescore(ctx.gray, found)
        found.sort(key=lambda c: c.score, reverse=True)
        return found[0]

    def _rescore(self, gray, circles):
        rays = int(self._get("score_rays", 72))
        band = float(self._get("score_band", 0.15))
        min_grad = float(self._get("score_min_grad", 8.0))
        tol = float(self._get("score_tol", 0.06))
        for c in circles:
            c.meta.setdefault("native_score", round(c.score, 3))
            c.score = edge_support_score(gray, c.cx, c.cy, c.r, rays, band, min_grad, tol)
        return circles

    @staticmethod
    def _nms(circles, dist_frac):
        """Greedy suppression: drop a lower-scored circle whose centre is within
        ``dist_frac * r`` of an already-kept one (same physical circle found by
        multiple methods)."""
        kept = []
        for c in sorted(circles, key=lambda c: c.score, reverse=True):
            dup = False
            for k in kept:
                if np.hypot(c.cx - k.cx, c.cy - k.cy) < dist_frac * max(k.r, 1.0):
                    dup = True
                    break
            if not dup:
                kept.append(c)
        return kept

    def _refine(self, gray, c: Circle) -> Circle:
        rays = int(self._get("subpixel_rays", 36))
        rng = float(self._get("subpixel_range", 6.0))
        cx, cy, r = refine_center_subpixel_radial(gray, c.cx, c.cy, c.r, rays, rng, trim=True)
        # Guard: keep the refinement only if it stayed sane (radial fits can
        # latch onto a neighbouring edge). Otherwise keep the original.
        if abs(cx - c.cx) <= 0.25 * c.r and abs(cy - c.cy) <= 0.25 * c.r and r > 1:
            return Circle(cx, cy, r, c.score, c.method, {**c.meta, "refined": True})
        return c

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

    def _draw_and_save(self, image, circles, save_path):
        out = image.copy()
        rois = self._resolve_rois(self._get("roi"))
        for ri, roi in enumerate(rois):
            if roi is not None:
                cv2.rectangle(out, (roi["x"], roi["y"]),
                              (roi["x"] + roi["w"], roi["y"] + roi["h"]), (255, 0, 0), 2)
                if len(rois) > 1:
                    cv2.putText(out, f"#{ri + 1}", (roi["x"] + 3, roi["y"] + 18),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1, cv2.LINE_AA)
        for i, c in enumerate(circles):
            cx, cy, r = c.as_int()
            tag = f"#{i+1}" if len(rois) == 1 else f"roi{c.meta.get('roi_index', 0)}"
            cv2.circle(out, (cx, cy), r, (0, 255, 0), 2, cv2.LINE_AA)
            cv2.drawMarker(out, (cx, cy), (0, 0, 255), cv2.MARKER_CROSS, 12, 2)
            cv2.putText(out, f"{tag} r={r} s={c.score:.2f} {c.method}",
                        (cx - r, cy - r - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        (0, 255, 0), 1, cv2.LINE_AA)
        save_dir = os.path.dirname(os.path.abspath(save_path))
        if save_dir:
            os.makedirs(save_dir, exist_ok=True)
        cv2.imwrite(save_path, out)
