"""Identity lanes and the decision that fuses them.

Two lanes answer "which catalog row is this crop", and one policy turns their
evidence into accepted / ambiguous / unknown / rejected:

- `LocalLane` reads printed geometry (RootSIFT over a reference pyramid,
  similarity-constrained RANSAC, a hue check). It is the only primitive that
  ranks fine-grained siblings correctly: 40/42 enrolled facings first at full
  resolution, against 32/42 for Van's composite.
- `TextLane` reads the package text and compares it to each catalog row's
  signature (its printed name plus the OCR of its enrolled packshot). It is the
  only measured mechanism that rejects a near-unknown sibling, and it can name
  catalog rows that have no enrolled image at all.

Why a policy and not a threshold: the matcher's raw score is not comparable
across references of different resolution, and a single unverified match used
to clear the old floor. Measured on the corrected image6 labels, the old gate
accepts 22-28 of 48 unknowns; the local gate alone reaches 79% precision at 43%
coverage, and adding the text lane reaches 95% precision at 59% coverage.

`ponytail:` the policy is a hand-set pair of lines tuned on one photo, not a
calibrated model. E-4 replaces it once Tier-2 held-out labels exist.
"""

from __future__ import annotations

import os
import sys
import unicodedata
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import cv2  # noqa: E402
import numpy as np  # noqa: E402

from rebuild_catalog import sku_of  # noqa: E402

# The local lane's accept line, in the units of `confidence()` below. Measured
# on the corrected image6 labels: 2.0 keeps 95% of shown labels correct at 59%
# coverage, 1.0 keeps 87% at 66%. The stricter line is the default because a
# wrong label on a shelf is worse than a missing one.
LOCAL_ACCEPT = 2.0
# The text lane's own accept line: a bigram overlap this high, with this much
# margin over the runner-up, names the row on its own. Measured: the text lane
# picks the right SKU on 41-42 of the 55 facings where it reads anything.
TEXT_ACCEPT = 0.5
TEXT_MARGIN = 0.15
# Reference pyramid heights. A ~1000 px packshot against a 110-250 px crop is a
# 4-8x ratio, and ORB's 8x1.2 pyramid tops out at 3.6x, which is why ORB and
# BRISK find zero inliers at live resolution. Scaling each reference to the
# crop's own size removes the ratio instead of widening the pyramid.
REFERENCE_HEIGHTS = (200, 400)
# Lowe's ratio test.
RATIO_THRESH = 0.75
# RANSAC reprojection tolerance, as a fraction of the reference's long side.
REPROJ_FRACTION = 0.02
# Hue histogram: chromatic pixels only, so a white background cannot vote.
HUE_BINS = 18
HUE_SATURATION_MIN = 60
HUE_VALUE_MIN = 50


@dataclass
class Evidence:
    """What one reference view says about one crop."""

    inliers: int = 0
    cover: float = 0.0   # inlier hull area / reference area
    qcover: float = 0.0  # inlier hull area / crop area
    color: float = 0.0   # hue histogram intersection


def _hue_histogram(image: Any) -> np.ndarray:
    """Hue distribution over the chromatic pixels, L1-normalized."""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    chromatic = ((hsv[..., 1] > HUE_SATURATION_MIN) & (hsv[..., 2] > HUE_VALUE_MIN))
    hist = cv2.calcHist([hsv], [0], chromatic.astype(np.uint8), [HUE_BINS], [0, 180]).ravel()
    return hist / (hist.sum() + 1e-6)


def _root_sift(descriptors: Optional[np.ndarray]) -> Optional[np.ndarray]:
    """RootSIFT: L1-normalize then square-root. Cheap, and it beats raw SIFT."""
    if descriptors is None:
        return None
    normalized = descriptors / (np.abs(descriptors).sum(1, keepdims=True) + 1e-7)
    return np.sqrt(normalized).astype(np.float32)


class LocalLane:
    """RootSIFT over a reference pyramid, similarity-RANSAC, hue re-rank.

    Scores every enrolled reference view and reports the best view per SKU, so
    a multi-view catalog cannot make one product compete with itself.
    """

    name = "local"

    def __init__(
        self,
        stock_dir: str,
        *,
        heights: Sequence[int] = REFERENCE_HEIGHTS,
        nfeatures: int = 2000,
    ) -> None:
        self.stock_dir = stock_dir
        self.sift = cv2.SIFT_create(nfeatures=nfeatures)
        self.matcher = cv2.BFMatcher(cv2.NORM_L2)
        self.views: List[Dict[str, Any]] = []
        self.sku_ids: List[str] = []
        self._load(heights)

    def _load(self, heights: Sequence[int]) -> None:
        if not os.path.exists(self.stock_dir):
            return
        names = sorted(
            name for name in os.listdir(self.stock_dir)
            if name.lower().endswith((".jpg", ".jpeg", ".png"))
        )
        for name in names:
            image = cv2.imread(os.path.join(self.stock_dir, name))
            if image is None:
                continue
            sku = sku_of(name)
            if sku not in self.sku_ids:
                self.sku_ids.append(sku)
            for height in heights:
                scaled = cv2.resize(
                    image,
                    (max(1, int(image.shape[1] * height / image.shape[0])), int(height)),
                    interpolation=cv2.INTER_AREA,
                )
                gray = cv2.cvtColor(scaled, cv2.COLOR_BGR2GRAY)
                keypoints, descriptors = self.sift.detectAndCompute(gray, None)
                self.views.append(
                    {
                        "sku": sku,
                        "keypoints": keypoints,
                        "descriptors": _root_sift(descriptors),
                        "shape": gray.shape,
                        "hist": _hue_histogram(scaled),
                    }
                )

    def view_count(self) -> int:
        return len(self.views)

    def evidence(self, crop: Any) -> Dict[str, Evidence]:
        """Best evidence per canonical SKU. Empty when the crop has no features."""
        if crop is None or getattr(crop, "size", 0) == 0:
            return {}
        gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
        keypoints, descriptors = self.sift.detectAndCompute(gray, None)
        descriptors = _root_sift(descriptors)
        if descriptors is None or len(descriptors) < 2:
            return {}
        query_hist = _hue_histogram(crop)
        crop_area = float(gray.shape[0] * gray.shape[1])

        best: Dict[str, Evidence] = {}
        for view in self.views:
            reference = view["descriptors"]
            if reference is None or len(reference) < 2:
                continue
            pairs = self.matcher.knnMatch(descriptors, reference, k=2)
            good = [
                first for first, second in (pair for pair in pairs if len(pair) == 2)
                if first.distance < RATIO_THRESH * second.distance
            ]
            if len(good) < 4:
                continue
            source = np.float32([keypoints[m.queryIdx].pt for m in good])
            target = np.float32([view["keypoints"][m.trainIdx].pt for m in good])
            _, mask = cv2.estimateAffinePartial2D(
                source, target, method=cv2.RANSAC,
                ransacReprojThreshold=REPROJ_FRACTION * max(view["shape"]),
            )
            if mask is None:
                continue
            inliers = mask.ravel().astype(bool)
            count = int(inliers.sum())
            if count < 3:
                continue
            reference_area = float(view["shape"][0] * view["shape"][1])
            evidence = Evidence(
                inliers=count,
                cover=float(cv2.contourArea(cv2.convexHull(target[inliers])) / reference_area),
                qcover=float(cv2.contourArea(cv2.convexHull(source[inliers])) / crop_area),
                color=float(np.minimum(query_hist, view["hist"]).sum()),
            )
            current = best.get(view["sku"])
            if current is None or evidence.inliers > current.inliers:
                best[view["sku"]] = evidence
        return best

    @staticmethod
    def rank_key(evidence: Evidence) -> float:
        """Order candidates. The hue term is what separates colour-coded variants
        (pink ゼロ from red ハーフ), which grayscale SIFT cannot see at all."""
        return evidence.inliers * (0.25 + evidence.color)

    @staticmethod
    def confidence(evidence: Evidence, runner_up: Evidence) -> float:
        """How much this crop's best candidate stands out, in inlier units.

        The runner-up is subtracted in inlier space, so a sibling that matches
        the same artwork cannot be separated by a high score alone; the hue
        agreement then scales the result, because a match that disagrees on
        colour is not the same package.
        """
        margin = 1.0 - runner_up.inliers / max(evidence.inliers, 1)
        return evidence.inliers * margin * evidence.color

    def rank(self, crop: Any) -> List[Tuple[str, float]]:
        """Candidates best-first, each with the confidence the gate reads.

        Order and value are different quantities on purpose: the order comes
        from `rank_key` (inliers weighted by hue), the value from `confidence`
        (the margin over the runner-up). Sorting by the value would reorder the
        candidates, which is how a lane that ranks 40/42 correctly ends up
        gating on the wrong one.
        """
        evidence = self.evidence(crop)
        if not evidence:
            return []
        ranked = sorted(evidence.items(), key=lambda kv: self.rank_key(kv[1]), reverse=True)
        return [
            (
                sku,
                self.confidence(
                    item,
                    ranked[index + 1][1] if index + 1 < len(ranked) else Evidence(),
                ),
            )
            for index, (sku, item) in enumerate(ranked)
        ]

    def scores(self, crop: Any) -> Dict[str, float]:
        """`rank()` as a dict, in rank order. Convenience for callers that only
        need the values; the policy uses `rank()` so the order survives."""
        return {sku: value for sku, value in self.rank(crop)}


# Katakana folding: OCR reads small kana and the long-vowel mark inconsistently,
# so キユーピー and キユピー have to land on the same bigrams.
_KANA_FOLD = str.maketrans(
    "ッャュョァィゥェォヮヵヶユ", "ツヤユヨアイウエオワカケユ"
)


def normalize_text(text: str) -> str:
    """NFKC, uppercase, kana-folded, alphanumeric and `%` only."""
    text = unicodedata.normalize("NFKC", text).upper()
    text = text.translate(_KANA_FOLD)
    text = text.replace("キユーピー", "キユピ").replace("キューピー", "キユピ")
    return "".join(ch for ch in text if ch.isalnum() or ch == "%")


def bigrams(text: str) -> set:
    return {text[i:i + 2] for i in range(len(text) - 1)}


def text_similarity(read: str, signature: str) -> float:
    """Share of the read text's bigrams that the signature explains.

    Asymmetric on purpose: the crop shows a few words of a package, the
    signature holds the whole name, so the question is how much of what was
    read the catalog row accounts for.
    """
    read_grams = bigrams(normalize_text(read))
    if not read_grams:
        return 0.0
    return len(read_grams & bigrams(normalize_text(signature))) / max(1, len(read_grams))


class TextLane:
    """Package text against each catalog row's signature.

    The signature is the row's printed name plus the OCR of its enrolled
    packshot, so it is built from the catalog and never from the evaluation
    photo. Rows with no image still get a signature from their name, which is
    how the lane can name a SKU no matcher can see.
    """

    name = "text"

    def __init__(
        self, signatures: Dict[str, str], ocr_engine: Any = None, lock: Any = None
    ) -> None:
        self.signatures = {str(sku): text for sku, text in signatures.items()}
        self.ocr_engine = ocr_engine
        # The engine's OCR lock. Vision tolerates concurrency but the CPU does
        # not, and this lane runs one OCR per crop on the classify pool, so it
        # shares the same lock as tag reading instead of adding a second one.
        self.lock = lock

    @classmethod
    def from_catalog(
        cls,
        catalog_rows: Sequence[Dict[str, Any]],
        stock_dir: str,
        ocr_engine: Any = None,
        lock: Any = None,
    ) -> "TextLane":
        signatures: Dict[str, str] = {}
        for row in catalog_rows:
            sku = str(row.get("sku_id") or "")
            if not sku:
                continue
            signature = str(row.get("product_name") or "")
            if ocr_engine is not None and getattr(ocr_engine, "available", lambda: False)():
                path = os.path.join(stock_dir, f"{sku}.jpg")
                if os.path.exists(path):
                    image = cv2.imread(path)
                    if image is not None:
                        try:
                            result = ocr_engine.read(image)
                        except Exception:
                            result = None
                        if result is not None:
                            signature += " " + " ".join(line.text for line in result.lines)
            signatures[sku] = signature
        return cls(signatures, ocr_engine=ocr_engine, lock=lock)

    def read(self, crop: Any) -> str:
        """The text on one crop, upsampled when the crop is small.

        Below ~200 px tall the recognizer misses the small variant words that
        are the whole point of this lane, so the crop is doubled first.
        """
        if self.ocr_engine is None or crop is None or getattr(crop, "size", 0) == 0:
            return ""
        if not getattr(self.ocr_engine, "available", lambda: False)():
            return ""
        image = crop
        if image.shape[0] < 200:
            image = cv2.resize(image, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)
        try:
            if self.lock is not None:
                with self.lock:
                    result = self.ocr_engine.read(image)
            else:
                result = self.ocr_engine.read(image)
        except Exception:
            return ""
        return " ".join(line.text for line in result.lines)

    def scores(self, crop: Any) -> Dict[str, float]:
        """Bigram similarity per catalog row. Empty when nothing was read."""
        return self.scores_from_text(self.read(crop))

    def scores_from_text(self, read: str) -> Dict[str, float]:
        if len(normalize_text(read)) < 4:
            return {}
        return {
            sku: text_similarity(read, signature)
            for sku, signature in self.signatures.items()
        }


def _ranked(scores: Any) -> List[Tuple[str, float]]:
    """Candidates best-first.

    A lane hands over an already-ordered sequence and it is used as given: the
    order is the lane's ranking, and re-sorting by the confidence value would
    silently replace it. A plain dict is sorted by value, which is the only
    order a dict has.
    """
    if scores is None:
        return []
    if isinstance(scores, dict):
        items = scores.items()
        return sorted(
            ((str(sku), float(value)) for sku, value in items),
            key=lambda kv: kv[1],
            reverse=True,
        )
    return [(str(sku), float(value)) for sku, value in scores]


class DecisionPolicy:
    """Turn two lanes' evidence into one identity decision.

    The local lane answers first, because it is the only lane that ranks
    siblings correctly. The text lane then does one of two jobs: it vetoes a
    local answer whose package text names a different catalog row, and it
    answers on its own when the local lane has nothing. Both jobs are measured:
    the veto takes precision from 62% to 82% for one lost correct label, and
    the lane adds ten correct labels at live resolution, including SKUs with no
    enrolled image.
    """

    def __init__(
        self,
        *,
        local_accept: float = LOCAL_ACCEPT,
        text_accept: float = TEXT_ACCEPT,
        text_margin: float = TEXT_MARGIN,
        method: str = "rootsift",
    ) -> None:
        self.local_accept = float(local_accept)
        self.text_accept = float(text_accept)
        self.text_margin = float(text_margin)
        self.method = method

    def decide(
        self,
        local_rank: Any,
        text_rank: Any = None,
        *,
        details: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """One decision, in the shape `decide_identity` already returns."""
        from identity import (
            IDENTITY_ACCEPTED,
            IDENTITY_AMBIGUOUS,
            IDENTITY_REJECTED,
            IDENTITY_UNKNOWN,
        )

        local = _ranked(local_rank)
        text = _ranked(text_rank)
        text_top = text[0] if text else None
        text_runner = text[1][1] if len(text) > 1 else 0.0
        text_names = (
            text_top[0]
            if text_top
            and text_top[1] >= self.text_accept
            and (text_top[1] - text_runner) >= self.text_margin
            else None
        )

        answer: Optional[str] = None
        source: Optional[str] = None
        if local and local[0][1] >= self.local_accept:
            candidate = local[0][0]
            # A local match whose package text names a different row is the
            # near-unknown case no threshold can reject: the sibling shares the
            # artwork, so the geometry is real and only the words differ.
            if text_top and text_top[0] != candidate and text_top[1] > local_text_floor(text, candidate):
                answer = None
            else:
                answer, source = candidate, "local"
        if answer is None and text_names:
            answer, source = text_names, "text"

        ranked = [
            {"sku_id": sku, "score": round(value, 2)} for sku, value in local[:5]
        ]
        if not ranked and text:
            ranked = [{"sku_id": sku, "score": round(value, 3)} for sku, value in text[:5]]

        if answer is not None:
            level = IDENTITY_ACCEPTED
            message = "Accepted unique catalog match"
        elif local and local[0][1] > 0.0:
            level = IDENTITY_AMBIGUOUS if len(local) > 1 else IDENTITY_UNKNOWN
            message = "No candidate cleared the accept line"
        elif text:
            level = IDENTITY_UNKNOWN
            message = "Text read, but it does not name one catalog row"
        else:
            level = IDENTITY_REJECTED
            message = "No matching evidence"

        top_score = local[0][1] if local else (text_top[1] if text_top else 0.0)
        return {
            "status": "success" if level == IDENTITY_ACCEPTED else level,
            "method": self.method,
            "best_sku": answer,
            "score": round(float(top_score), 2),
            "confidence": round(float(top_score), 4),
            "identity_level": level,
            "ranked_candidates": ranked,
            "details": (details or {}).get(answer, {}) if answer else {},
            "message": message,
            "source": source,
        }


def local_text_floor(text: List[Tuple[str, float]], candidate: str) -> float:
    """The text score the local candidate itself explains.

    A veto needs the text to *prefer* another row, not merely to mention it: a
    crop that reads "マヨネーズ" scores every mayonnaise row, and the local
    candidate is usually one of them. Only a row the text explains better than
    the local answer's own row is a contradiction.
    """
    for sku, value in text:
        if sku == candidate:
            return value
    return 0.0