r""" Convolutional Hough Matching Networks """

import torch.nn as nn
import torch

from . import chmlearner as chmlearner
from .base import backbone


class CHMNet(nn.Module):
    def __init__(self, ktype):
        super(CHMNet, self).__init__()

        self.backbone = backbone.resnet101(pretrained=True)
        self.learner = chmlearner.CHMLearner(ktype, feat_dim=1024)

    def forward(self, src_img, trg_img):
        src_feat, trg_feat = self.extract_features(src_img, trg_img)
        correlation  = self.learner(src_feat, trg_feat)
        return correlation

    def extract_features(self, src_img, trg_img):
        feat = self.backbone.conv1.forward(torch.cat([src_img, trg_img], dim=1))
        feat = self.backbone.bn1.forward(feat)
        feat = self.backbone.relu.forward(feat)
        feat = self.backbone.maxpool.forward(feat)

        for idx in range(1, 5):
            feat = self.backbone.__getattr__('layer%d' % idx)(feat)

            if idx == 3:
                src_feat = feat.narrow(1, 0, feat.size(1) // 2).clone()
                trg_feat = feat.narrow(1, feat.size(1) // 2, feat.size(1) // 2).clone()
                return src_feat, trg_feat

    def training_objective(cls, prd_kps, trg_kps, npts):
        l2dist = (prd_kps - trg_kps).pow(2).sum(dim=1)
        loss = []
        for dist, npt in zip(l2dist, npts):
            loss.append(dist[:npt].mean())
        return torch.stack(loss).mean()