import torch

from ..utils.base_model import BaseModel


def find_nn(sim, ratio_thresh, distance_thresh):
    sim_nn, ind_nn = sim.topk(2 if ratio_thresh else 1, dim=-1, largest=True)
    dist_nn = 2 * (1 - sim_nn)
    mask = torch.ones(ind_nn.shape[:-1], dtype=torch.bool, device=sim.device)
    if ratio_thresh:
        mask = mask & (dist_nn[..., 0] <= (ratio_thresh**2) * dist_nn[..., 1])
    if distance_thresh:
        mask = mask & (dist_nn[..., 0] <= distance_thresh**2)
    matches = torch.where(mask, ind_nn[..., 0], ind_nn.new_tensor(-1))
    scores = torch.where(mask, (sim_nn[..., 0] + 1) / 2, sim_nn.new_tensor(0))
    return matches, scores


def mutual_check(m0, m1):
    inds0 = torch.arange(m0.shape[-1], device=m0.device)
    loop = torch.gather(m1, -1, torch.where(m0 > -1, m0, m0.new_tensor(0)))
    ok = (m0 > -1) & (inds0 == loop)
    m0_new = torch.where(ok, m0, m0.new_tensor(-1))
    return m0_new


class NearestNeighbor(BaseModel):
    default_conf = {
        "ratio_threshold": None,
        "distance_threshold": None,
        "do_mutual_check": True,
    }
    required_inputs = ["descriptors0", "descriptors1"]

    def _init(self, conf):
        pass

    def _forward(self, data):
        if (
            data["descriptors0"].size(-1) == 0
            or data["descriptors1"].size(-1) == 0
        ):
            matches0 = torch.full(
                data["descriptors0"].shape[:2],
                -1,
                device=data["descriptors0"].device,
            )
            return {
                "matches0": matches0,
                "matching_scores0": torch.zeros_like(matches0),
            }
        ratio_threshold = self.conf["ratio_threshold"]
        if (
            data["descriptors0"].size(-1) == 1
            or data["descriptors1"].size(-1) == 1
        ):
            ratio_threshold = None
        sim = torch.einsum(
            "bdn,bdm->bnm", data["descriptors0"], data["descriptors1"]
        )
        matches0, scores0 = find_nn(
            sim, ratio_threshold, self.conf["distance_threshold"]
        )
        if self.conf["do_mutual_check"]:
            matches1, scores1 = find_nn(
                sim.transpose(1, 2),
                ratio_threshold,
                self.conf["distance_threshold"],
            )
            matches0 = mutual_check(matches0, matches1)
        return {
            "matches0": matches0,
            "matching_scores0": scores0,
        }