import torch

from hloc import logger

from ..utils.base_model import BaseModel


class XFeatDense(BaseModel):
    default_conf = {
        "keypoint_threshold": 0.005,
        "max_keypoints": 8000,
    }
    required_inputs = [
        "image0",
        "image1",
    ]

    def _init(self, conf):
        self.net = torch.hub.load(
            "verlab/accelerated_features",
            "XFeat",
            pretrained=True,
            top_k=self.conf["max_keypoints"],
        )
        logger.info("Load XFeat(dense) model done.")

    def _forward(self, data):
        # Compute coarse feats
        out0 = self.net.detectAndComputeDense(
            data["image0"], top_k=self.conf["max_keypoints"]
        )
        out1 = self.net.detectAndComputeDense(
            data["image1"], top_k=self.conf["max_keypoints"]
        )

        # Match batches of pairs
        idxs_list = self.net.batch_match(
            out0["descriptors"], out1["descriptors"]
        )
        B = len(data["image0"])

        # Refine coarse matches
        # this part is harder to batch, currently iterate
        matches = []
        for b in range(B):
            matches.append(
                self.net.refine_matches(
                    out0, out1, matches=idxs_list, batch_idx=b
                )
            )
        # we use results from one batch
        matches = matches[0]
        pred = {
            "keypoints0": matches[:, :2],
            "keypoints1": matches[:, 2:],
            "mconf": torch.ones_like(matches[:, 0]),
        }
        return pred