import torch
from omegaconf import OmegaConf

from easy_local_features.submodules.git_alike.alike import ALike

from ..matching.nearest_neighbor import NearestNeighborMatcher
from ..utils import ops
from ..utils.download import downloadModel
from .basemodel import BaseExtractor, MethodType
from typing import TypedDict, Optional


class ALIKEConfig(TypedDict):
    model_name: str
    top_k: int
    scores_th: float
    n_limit: int
    sub_pixel: bool
    model_path: Optional[str]


models = {
    "alike-t": "https://github.com/Shiaoming/ALIKE/raw/main/models/alike-t.pth",
    "alike-s": "https://github.com/Shiaoming/ALIKE/raw/main/models/alike-s.pth",
    "alike-n": "https://github.com/Shiaoming/ALIKE/raw/main/models/alike-n.pth",
    "alike-l": "https://github.com/Shiaoming/ALIKE/raw/main/models/alike-l.pth",
}

configs = {
    "alike-t": {"c1": 8, "c2": 16, "c3": 32, "c4": 64, "dim": 64, "single_head": True, "radius": 2, "model_path": ""},
    "alike-s": {"c1": 8, "c2": 16, "c3": 48, "c4": 96, "dim": 96, "single_head": True, "radius": 2, "model_path": ""},
    "alike-n": {
        "c1": 16,
        "c2": 32,
        "c3": 64,
        "c4": 128,
        "dim": 128,
        "single_head": True,
        "radius": 2,
        "model_path": "",
    },
    "alike-l": {
        "c1": 32,
        "c2": 64,
        "c3": 128,
        "c4": 128,
        "dim": 128,
        "single_head": False,
        "radius": 2,
        "model_path": "",
    },
}


class ALIKE_baseline(BaseExtractor):
    METHOD_TYPE = MethodType.DETECT_DESCRIBE
    """ALIKE baseline implementation.
    model_name: str = 'alike-t' | 'alike-s' | 'alike-n' | 'alike-l'
    top_k: int = -1. Detect top K keypoints. -1 for threshold based mode, >0 for top K mode. (default: -1)
    scores_th: float = 0.2. Detector score threshold (default: 0.2).
    n_limit: int = 5000. Maximum number of keypoints to be detected (default: 5000).
    no_sub_pixel: bool = False. Do not detect sub-pixel keypoints (default: False).
    device: int = -1. Device to run the model on. -1 for CPU, >=0 for GPU. (default: -1)
    """

    default_conf: ALIKEConfig = {
        "model_name": "alike-t",
        "top_k": -1,
        "scores_th": 0.2,
        "n_limit": 2048,
        "sub_pixel": True,
        "model_path": None,
    }

    def __init__(self, conf: ALIKEConfig = {}):
        self.conf = conf = OmegaConf.merge(OmegaConf.create(self.default_conf), conf)
        model_name = conf.model_name
        top_k = conf.top_k
        scores_th = conf.scores_th
        n_limit = conf.n_limit
        sub_pixel = conf.sub_pixel
        model_path = conf.model_path
        self.DEV = torch.device("cpu")

        if model_name not in models:
            raise ValueError(f"Model name {model_name} not found in {models.keys()}")

        if model_path is None:
            url = models[model_name]
            model_path = downloadModel("alike", model_name, url)

        config = configs[model_name]
        config["model_path"] = model_path

        self.model = ALike(**config, top_k=top_k, scores_th=scores_th, n_limit=n_limit).to(self.DEV)

        self.sub_pixel = sub_pixel
        self.matcher = NearestNeighborMatcher()

    def detectAndCompute(self, image, return_dict=False):
        image = ops.prepareImage(image)
        image = image.to(self.DEV)
        pred = self.model(image, sub_pixel=self.sub_pixel)
        keypoints = pred["keypoints"]  # (N, 2)
        descriptors = pred["descriptors"]  # (N, 64)
        scores = pred["scores"]  # (N,)

        keypoints = keypoints.unsqueeze(0)
        descriptors = descriptors.unsqueeze(0)
        scores = scores.unsqueeze(0)

        if return_dict:
            return {"keypoints": keypoints, "descriptors": descriptors, "scores": scores}

        return keypoints, descriptors

    def detect(self, img):
        return self.detectAndCompute(img)[0]

    def compute(self, image, keypoints):
        raise NotImplementedError("ALIKE_baseline.compute is not implemented; use detectAndCompute instead.")

    def to(self, device):
        self.model.to(device)
        self.DEV = device

    @property
    def has_detector(self):
        return True
