"""
Deep-learning matchers: SuperPoint+LightGlue ("SuperGlue" in the UI/output
naming, per the existing convention -- LightGlue is the permissively
licensed, actively maintained successor actually running under the hood)
and LoFTR (dense, pairwise).

Template-side features/tensors are precomputed once at startup and reused
for every request; only the query image is processed per request. Since
there are only a handful of templates, even LoFTR (which has no reusable
per-image descriptor and must re-run a full forward pass per pair) is cheap
here -- it does NOT scale the way it would against hundreds of images.
"""

import gc
import logging

import cv2
import numpy as np
import torch

import config

logger = logging.getLogger(__name__)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_superpoint = None
_lightglue = None
_loftr = None

_template_superpoint_feats = {}   # template_name -> {"feats", "mask"}
_template_loftr_tensors = {}       # template_name -> {"tensor", "mask"}


def get_superpoint():
    global _superpoint
    if _superpoint is None:
        logger.info("Loading SuperPoint on %s...", DEVICE)
        from lightglue import SuperPoint
        _superpoint = SuperPoint(
            max_num_keypoints=config.SUPERPOINT_MAX_KEYPOINTS
        ).eval().to(DEVICE)
    return _superpoint


def get_lightglue():
    global _lightglue
    if _lightglue is None:
        logger.info("Loading LightGlue on %s...", DEVICE)
        from lightglue import LightGlue
        _lightglue = LightGlue(features="superpoint").eval().to(DEVICE)
    return _lightglue


def get_loftr():
    global _loftr
    if _loftr is None:
        logger.info("Loading LoFTR on %s...", DEVICE)
        import kornia.feature as KF
        _loftr = KF.LoFTR(pretrained="outdoor").eval().to(DEVICE)
    return _loftr


def unload_models():
    """Drops the SuperPoint/LightGlue/LoFTR model objects (not the small
    per-template feature/tensor caches, which stay put) and frees their CUDA
    memory. They're plain lazy singletons (see get_superpoint/get_lightglue/
    get_loftr above), so the next call to any of those simply reloads --
    exactly like a fresh process start. Used to make room for SAM-based
    flower counting on an 8GB card that can't hold everything at once."""
    global _superpoint, _lightglue, _loftr
    freed = _superpoint is not None or _lightglue is not None or _loftr is not None
    _superpoint = None
    _lightglue = None
    _loftr = None
    if freed:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return freed


def masked_gray_tensor(bgr, mask, max_dim=None, fill_value=0.5):
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0

    if max_dim is not None:
        h, w = gray.shape
        scale = max_dim / max(h, w)
        if scale < 1.0:
            new_w, new_h = int(w * scale), int(h * scale)
            gray = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_AREA)
            mask = cv2.resize(mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST)

    gray[mask == 0] = fill_value
    tensor = torch.from_numpy(gray)[None, None].to(DEVICE)
    return tensor, mask


def keep_points_inside_mask(pts, mask):
    if len(pts) == 0:
        return np.zeros(0, dtype=bool)
    h, w = mask.shape
    xs = np.clip(pts[:, 0].round().astype(int), 0, w - 1)
    ys = np.clip(pts[:, 1].round().astype(int), 0, h - 1)
    return mask[ys, xs] > 0


# ---------------------------------------------------------------
# SuperPoint + LightGlue
# ---------------------------------------------------------------

def superpoint_extract(bgr, mask):
    # ASSUMPTION: cap at MAX_IMAGE_DIM even though uploads are already
    # resized upstream -- this is what used to run at full (sometimes
    # multi-thousand-pixel) upload resolution and was the main driver of
    # the multi-GB memory spikes that got the process OOM-killed.
    tensor, resized_mask = masked_gray_tensor(bgr, mask, max_dim=config.MAX_IMAGE_DIM)
    with torch.no_grad():
        feats = get_superpoint().extract(tensor)
    return feats, resized_mask


def set_template_superpoint(template_name, bgr, mask):
    feats, resized_mask = superpoint_extract(bgr, mask)
    _template_superpoint_feats[template_name] = {"feats": feats, "mask": resized_mask}


def _lightglue_pair(feats_q, feats_t, mask_q, mask_t):
    from lightglue.utils import rbd

    with torch.no_grad():
        matches01 = get_lightglue()({"image0": feats_q, "image1": feats_t})

    feats_q_, feats_t_, matches01_ = [rbd(x) for x in [feats_q, feats_t, matches01]]
    matches = matches01_["matches"]
    if matches.shape[0] == 0:
        return 0, 0.0

    kpts_q = feats_q_["keypoints"][matches[..., 0]].cpu().numpy()
    kpts_t = feats_t_["keypoints"][matches[..., 1]].cpu().numpy()

    keep = keep_points_inside_mask(kpts_q, mask_q) & keep_points_inside_mask(kpts_t, mask_t)
    kpts_q, kpts_t = kpts_q[keep], kpts_t[keep]

    if len(kpts_q) < 4:
        return 0, 0.0

    src = kpts_q.reshape(-1, 1, 2).astype(np.float32)
    dst = kpts_t.reshape(-1, 1, 2).astype(np.float32)

    _, ransac_mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
    if ransac_mask is None:
        return 0, 0.0

    inlier_count = int(ransac_mask.sum())
    confidence_pct = (inlier_count / len(src)) * 100 if len(src) > 0 else 0.0
    return inlier_count, confidence_pct


def superglue_match_against_templates(bgr, mask):
    feats_q, mask_q = superpoint_extract(bgr, mask)
    results = []
    for template_name, data in _template_superpoint_feats.items():
        inlier_count, confidence_pct = _lightglue_pair(
            feats_q, data["feats"], mask_q, data["mask"]
        )
        results.append({
            "template": template_name,
            "score": inlier_count,
            "confidence": round(confidence_pct, 1),
        })
    results.sort(key=lambda r: r["score"], reverse=True)
    return results


# ---------------------------------------------------------------
# LoFTR
# ---------------------------------------------------------------

def set_template_loftr(template_name, bgr, mask):
    tensor, resized_mask = masked_gray_tensor(bgr, mask, max_dim=config.LOFTR_MAX_DIM)
    _template_loftr_tensors[template_name] = {"tensor": tensor, "mask": resized_mask}


def _loftr_pair(tensor_q, tensor_t, mask_q, mask_t):
    with torch.no_grad():
        out = get_loftr()({"image0": tensor_q, "image1": tensor_t})

    conf = out["confidence"].cpu().numpy()
    keep_conf = conf >= config.LOFTR_CONFIDENCE_THRESHOLD

    kpts_q = out["keypoints0"].cpu().numpy()[keep_conf]
    kpts_t = out["keypoints1"].cpu().numpy()[keep_conf]

    keep = keep_points_inside_mask(kpts_q, mask_q) & keep_points_inside_mask(kpts_t, mask_t)
    kpts_q, kpts_t = kpts_q[keep], kpts_t[keep]

    if len(kpts_q) < 4:
        return 0, 0.0

    src = kpts_q.reshape(-1, 1, 2).astype(np.float32)
    dst = kpts_t.reshape(-1, 1, 2).astype(np.float32)

    _, ransac_mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
    if ransac_mask is None:
        return 0, 0.0

    inlier_count = int(ransac_mask.sum())
    confidence_pct = (inlier_count / len(src)) * 100 if len(src) > 0 else 0.0
    return inlier_count, confidence_pct


def loftr_match_against_templates(bgr, mask):
    tensor_q, mask_q = masked_gray_tensor(bgr, mask, max_dim=config.LOFTR_MAX_DIM)
    results = []
    for template_name, data in _template_loftr_tensors.items():
        inlier_count, confidence_pct = _loftr_pair(
            tensor_q, data["tensor"], mask_q, data["mask"]
        )
        results.append({
            "template": template_name,
            "score": inlier_count,
            "confidence": round(confidence_pct, 1),
        })
        # LoFTR has no reusable per-image descriptor -- this is a full dense
        # CNN+transformer forward pass PER TEMPLATE (6 of them per query,
        # back to back). Left uncleared, each pass's intermediate
        # activations fragment the allocator further than SuperPoint/
        # LightGlue's much lighter per-template cost ever does, and on an
        # 8GB card that's enough on its own to tip into CUDA OOM a few
        # templates in.
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    results.sort(key=lambda r: r["score"], reverse=True)
    return results