import torch
import cv2

from ..matching.nearest_neighbor import NearestNeighborMatcher
from omegaconf import OmegaConf
from .basemodel import BaseExtractor, MethodType
from ..utils.download import getCache
from ..utils import ops
from typing import TypedDict


class DEALConfig(TypedDict):
    model_name: str
    top_k: int
    detection_threshold: float
    nms_radius: int
    force_num_keypoints: bool


class DEAL_baseline(BaseExtractor):
    METHOD_TYPE = MethodType.DETECT_DESCRIBE
    default_conf: DEALConfig = {
        "model_name": "deal",
        "top_k": 2048,
        "detection_threshold": 0.2,
        "nms_radius": 4,
        "force_num_keypoints": False,
    }

    def __init__(self, conf: DEALConfig = {}):
        self.conf = conf = OmegaConf.merge(OmegaConf.create(self.default_conf), conf)

        self.max_kps = conf.top_k
        self.DEV = torch.device("cpu")
        cache = getCache("DEAL")
        self.sift = cv2.SIFT_create(nfeatures=self.max_kps, contrastThreshold=0.04, edgeThreshold=10)
        self.deal = torch.hub.load("verlab/DEAL_NeurIPS_2021", "DEAL", True, cache)
        self.deal.device = self.DEV
        self.deal.net.to(self.DEV)
        self._fix_stn_device(self.DEV)
        self.deal.net.eval()

        self.matcher = NearestNeighborMatcher()

    def detectAndCompute(self, img, return_dict=False):
        gray = ops.to_cv(ops.prepareImage(img), to_gray=True)

        with torch.no_grad():
            kps = self.detect(gray)
            kps, desc = self.compute(gray, kps)

        # compute() already returns tensors with batch dim
        if return_dict:
            return {"keypoints": kps, "descriptors": desc}

        return kps, desc

    def detect(self, img):
        gray = ops.to_cv(ops.prepareImage(img), to_gray=True)

        with torch.no_grad():
            kps = self.sift.detect(gray, None)

        return torch.tensor([kp.pt for kp in kps], dtype=torch.float32).to(self.DEV).unsqueeze(0)

    def compute(self, img, kps):
        gray = ops.to_cv(ops.prepareImage(img), to_gray=True)

        # Convert tensor keypoints to cv2.KeyPoint list for DEAL's SIFT-based compute
        if isinstance(kps, torch.Tensor):
            pts = kps.squeeze(0) if kps.ndim == 3 else kps  # [N, 2]
            cv_kps = [cv2.KeyPoint(float(pt[0]), float(pt[1]), 0) for pt in pts]
        elif isinstance(kps[0], cv2.KeyPoint):
            cv_kps = kps
        else:
            cv_kps = [cv2.KeyPoint(float(kp[0]), float(kp[1]), 0) for kp in kps]

        with torch.no_grad():
            desc = self.deal.compute(gray, cv_kps)

        desc = torch.from_numpy(desc).to(self.DEV).unsqueeze(0)

        kps_tensor = torch.tensor([kp.pt for kp in cv_kps], dtype=torch.float32).to(self.DEV).unsqueeze(0)
        return kps_tensor, desc

    def _fix_stn_device(self, device):
        """Update SpatialTransformer's hardcoded device and internal tensors.
        The upstream hub code creates SpatialTransformer with self.device = 'cuda'
        and stores self.pi as a plain tensor (not a buffer), so nn.Module.to() won't move them."""
        dev = torch.device(device) if isinstance(device, str) else device
        # Fix every module in the tree that has a .device attribute
        for mod in self.deal.net.modules():
            if hasattr(mod, "device") and isinstance(mod.device, torch.device):
                mod.device = dev
            if hasattr(mod, "pi") and isinstance(mod.pi, torch.Tensor):
                mod.pi = mod.pi.to(dev)

    def to(self, device):
        self.deal.device = device
        self.deal.net.to(device)
        self._fix_stn_device(device)
        self.DEV = device
        return self

    @property
    def has_detector(self):
        return True
