import torch

from .nn import NN2
from darkfeat import DarkFeat


class NNMatching(torch.nn.Module):
    def __init__(self, model_path=""):
        super().__init__()
        self.nn = NN2().eval()
        self.darkfeat = DarkFeat(model_path).eval()

    def forward(self, data):
        """Run DarkFeat and nearest neighborhood matching
        Args:
          data: dictionary with minimal keys: ['image0', 'image1']
        """
        pred = {}

        # Extract DarkFeat (keypoints, scores, descriptors)
        if "keypoints0" not in data:
            pred0 = self.darkfeat({"image": data["image0"]})
            # print({k+'0': v[0].shape for k, v in pred0.items()})
            pred = {**pred, **{k + "0": [v] for k, v in pred0.items()}}
        if "keypoints1" not in data:
            pred1 = self.darkfeat({"image": data["image1"]})
            pred = {**pred, **{k + "1": [v] for k, v in pred1.items()}}

        # Batch all features
        # We should either have i) one image per batch, or
        # ii) the same number of local features for all images in the batch.
        data = {**data, **pred}

        for k in data:
            if isinstance(data[k], (list, tuple)):
                data[k] = torch.stack(data[k])

        # Perform the matching
        pred = {**pred, **self.nn(data)}

        return pred