import torch, abc, sys

# from torch_dimcheck import dimchecked
# from ...submodules.torch_dimcheck import dimchecked

# the class/object below is there for making type annotations like
# def my_function(args) -> NpArray[OutputType]
if sys.version_info >= (3, 7):
    class NpArray:
        def __class_getitem__(self, arg):
            pass
else:
    # 3.6 and below don't support __class_getitem__
    class _NpArray:
        def __getitem__(self, _idx):
            pass
    
    NpArray = _NpArray()

class Features:
    # @dimchecked
    def __init__(self, kp, desc, kp_logp):
        assert kp.device == desc.device
        assert kp.device == kp_logp.device

        self.kp      = kp
        self.desc    = desc
        self.kp_logp = kp_logp

    @property
    def n(self):
        return self.kp.shape[0]

    @property
    def device(self):
        return self.kp.device

    def detached_and_grad_(self):
        return Features(
            self.kp,
            self.desc.detach().requires_grad_(),
            self.kp_logp.detach().requires_grad_(),
        )

    def requires_grad_(self, is_on):
        self.desc.requires_grad_(is_on)
        self.kp_logp.requires_grad_(is_on)

    def grad_tensors(self):
        return [self.desc, self.kp_logp]

    def to(self, *args, **kwargs):
        return Features(
            self.kp.to(*args, **kwargs),
            self.desc.to(*args, **kwargs),
            self.kp_logp.to(*args, **kwargs) if self.kp_logp is not None else None,
        )

class MatchDistribution(abc.ABC):
    @abc.abstractmethod
    def sample(self):
        pass

    @abc.abstractmethod
    def mle(self):
        pass

    @abc.abstractmethod
    def dense_logp(self):
        pass

    @abc.abstractmethod
    def dense_p(self):
        pass

    @abc.abstractmethod
    def features_1(self) -> Features:
        pass

    @abc.abstractmethod
    def features_2(self) -> Features:
        pass

    @property
    def shape(self):
        return self.features_1().kp.shape[0], self.features_2().kp.shape[1]

    def matched_pairs(self, mle=False):
        matches = self.mle() if mle else self.sample()

        return MatchedPairs(
            self.features_1().kp,
            self.features_2().kp,
            matches,    
        )


class MatchedPairs:
    # @dimchecked
    def __init__(self, kps1, kps2, matches):
        self.kps1    = kps1
        self.kps2    = kps2
        self.matches = matches

    def to(self, *args, **kwargs):
        return MatchedPairs(
            self.kps1.to(*args, **kwargs),
            self.kps2.to(*args, **kwargs),
            self.matches.to(*args, **kwargs),
        )
