"""
Implementation of the line matching methods.
"""
import numpy as np
import cv2
import torch
import torch.nn.functional as F
from ..misc.geometry_utils import keypoints_to_grid

class WunschLineMatcher(object):
    """ Class matching two sets of line segments
        with the Needleman-Wunsch algorithm. """
    def __init__(self, cross_check=True, num_samples=5, min_dist_pts=8,
                 top_k_candidates=10, grid_size=4, sampling="regular",
                 line_score=False):
        self.cross_check = cross_check
        self.num_samples = num_samples
        self.min_dist_pts = min_dist_pts
        self.top_k_candidates = top_k_candidates
        self.grid_size = grid_size
        self.line_score = line_score  # True to compute saliency on a line
        self.sampling_mode = sampling
        if sampling not in ["regular", "d2_net", "asl_feat"]:
            raise ValueError("Wrong sampling mode: " + sampling)

    def compute_descriptors(self, line_seg1, desc1):
        img_size1 = (desc1.shape[2] * self.grid_size,
                     desc1.shape[3] * self.grid_size)
        device = desc1.device

        # Default case when an image has no lines
        if len(line_seg1) == 0:
            return np.empty((0), dtype=int)

        # Sample points regularly along each line
        if self.sampling_mode == "regular":
            line_points1, valid_points1 = self.sample_line_points(line_seg1)
        else:
            line_points1, valid_points1 = self.sample_salient_points(
                line_seg1, desc1, img_size1, self.sampling_mode)
        line_points1 = torch.tensor(line_points1.reshape(-1, 2),
                                    dtype=torch.float, device=device)

        # Extract the descriptors for each point
        grid1 = keypoints_to_grid(line_points1, img_size1)
        desc1_sampled = F.normalize(F.grid_sample(desc1, grid1)[0, :, :, 0], dim=0)
        return [desc1_sampled, valid_points1]

    def compute_matches_topk_gpu(self, descinfo1, descinfo2, topk=10):
        desc1, valid_points1 = descinfo1
        desc2, valid_points2 = descinfo2
        len1, len2 = valid_points1.shape[0], valid_points2.shape[0]
        scores = desc1.t() @ desc2
        scores[~valid_points1.flatten()] = -1
        scores[:, ~valid_points2.flatten()] = -1
        scores = scores.reshape(len1, self.num_samples,
                                len2, self.num_samples)
        scores = scores.permute(0, 2, 1, 3)
        matches = self.filter_and_match_lines_topk_gpu(scores, topk=topk)
        return matches

    def compute_matches_topk(self, descinfo1, descinfo2, topk=10):
        desc1, valid_points1 = descinfo1
        desc2, valid_points2 = descinfo2
        len1, len2 = valid_points1.shape[0], valid_points2.shape[0]

        # Precompute the distance between line points for every pair of lines
        # Assign a score of -1 for unvalid points
        scores = (desc1.t() @ desc2).cpu().numpy()
        scores[~valid_points1.flatten()] = -1
        scores[:, ~valid_points2.flatten()] = -1
        scores = scores.reshape(len1, self.num_samples,
                                len2, self.num_samples)
        scores = scores.transpose(0, 2, 1, 3)

        # Pre-filter the line candidates and find the best match for each line
        matches = self.filter_and_match_lines_topk(scores, topk=topk)
        return matches

    def compute_matches(self, descinfo1, descinfo2):
        desc1, valid_points1 = descinfo1
        desc2, valid_points2 = descinfo2
        len1, len2 = valid_points1.shape[0], valid_points2.shape[0]

        # Precompute the distance between line points for every pair of lines
        # Assign a score of -1 for unvalid points
        scores = (desc1.t() @ desc2).cpu().numpy()
        scores[~valid_points1.flatten()] = -1
        scores[:, ~valid_points2.flatten()] = -1
        scores = scores.reshape(len1, self.num_samples,
                                len2, self.num_samples)
        scores = scores.transpose(0, 2, 1, 3)

        # Pre-filter the line candidates and find the best match for each line
        matches = self.filter_and_match_lines(scores)

        # [Optionally] filter matches with mutual nearest neighbor filtering
        if self.cross_check:
            matches2 = self.filter_and_match_lines(
                scores.transpose(1, 0, 3, 2))
            mutual = matches2[matches] == np.arange(len1)
            matches[~mutual] = -1
        return matches

    def forward(self, line_seg1, line_seg2, desc1, desc2):
        """
            Find the best matches between two sets of line segments
            and their corresponding descriptors.
        """
        descinfo1 = self.compute_descriptors(line_seg1, desc1)
        descinfo2 = self.compute_descriptors(line_seg2, desc2)
        matches = self.compute_matches(descinfo1, descinfo2)
        return matches

    def d2_net_saliency_score(self, desc):
        """ Compute the D2-Net saliency score
            on a 3D or 4D descriptor. """
        is_3d = len(desc.shape) == 3
        b_size = len(desc)
        feat = F.relu(desc)

        # Compute the soft local max
        exp = torch.exp(feat)
        if is_3d:
            sum_exp = 3 * F.avg_pool1d(exp, kernel_size=3, stride=1,
                                       padding=1)
        else:
            sum_exp = 9 * F.avg_pool2d(exp, kernel_size=3, stride=1,
                                       padding=1)
        soft_local_max = exp / sum_exp

        # Compute the depth-wise maximum
        depth_wise_max = torch.max(feat, dim=1)[0]
        depth_wise_max = feat / depth_wise_max.unsqueeze(1)

        # Total saliency score
        score = torch.max(soft_local_max * depth_wise_max, dim=1)[0]
        normalization = torch.sum(score.reshape(b_size, -1), dim=1)
        if is_3d:
            normalization = normalization.reshape(b_size, 1)
        else:
            normalization = normalization.reshape(b_size, 1, 1)
        score = score / normalization
        return score

    def asl_feat_saliency_score(self, desc):
        """ Compute the ASLFeat saliency score on a 3D or 4D descriptor. """
        is_3d = len(desc.shape) == 3
        b_size = len(desc)

        # Compute the soft local peakiness
        if is_3d:
            local_avg = F.avg_pool1d(desc, kernel_size=3, stride=1, padding=1)
        else:
            local_avg = F.avg_pool2d(desc, kernel_size=3, stride=1, padding=1)
        soft_local_score = F.softplus(desc - local_avg)

        # Compute the depth-wise peakiness
        depth_wise_mean = torch.mean(desc, dim=1).unsqueeze(1)
        depth_wise_score = F.softplus(desc - depth_wise_mean)

        # Total saliency score
        score = torch.max(soft_local_score * depth_wise_score, dim=1)[0]
        normalization = torch.sum(score.reshape(b_size, -1), dim=1)
        if is_3d:
            normalization = normalization.reshape(b_size, 1)
        else:
            normalization = normalization.reshape(b_size, 1, 1)
        score = score / normalization
        return score

    def sample_salient_points(self, line_seg, desc, img_size,
                              saliency_type='d2_net'):
        """
        Sample the most salient points along each line segments, with a
        minimal distance between each point. Pad the remaining points.
        Inputs:
            line_seg: an Nx2x2 torch.Tensor.
            desc: a NxDxHxW torch.Tensor.
            image_size: the original image size.
            saliency_type: 'd2_net' or 'asl_feat'.
        Outputs:
            line_points: an Nxnum_samplesx2 np.array.
            valid_points: a boolean Nxnum_samples np.array.
        """
        device = desc.device
        if not self.line_score:
            # Compute the score map
            if saliency_type == "d2_net":
                score = self.d2_net_saliency_score(desc)
            else:
                score = self.asl_feat_saliency_score(desc)

        num_lines = len(line_seg)
        line_lengths = np.linalg.norm(line_seg[:, 0] - line_seg[:, 1], axis=1)

        # The number of samples depends on the length of the line
        num_samples_lst = np.clip(line_lengths // self.min_dist_pts,
                                  2, self.num_samples)
        line_points = np.empty((num_lines, self.num_samples, 2), dtype=float)
        valid_points = np.empty((num_lines, self.num_samples), dtype=bool)

        # Sample the score on a fixed number of points of each line
        n_samples_per_region = 4
        for n in np.arange(2, self.num_samples + 1):
            sample_rate = n * n_samples_per_region
            # Consider all lines where we can fit up to n points
            cur_mask = num_samples_lst == n
            cur_line_seg = line_seg[cur_mask]
            cur_num_lines = len(cur_line_seg)
            if cur_num_lines == 0:
                continue
            line_points_x = np.linspace(cur_line_seg[:, 0, 0],
                                        cur_line_seg[:, 1, 0],
                                        sample_rate, axis=-1)
            line_points_y = np.linspace(cur_line_seg[:, 0, 1],
                                        cur_line_seg[:, 1, 1],
                                        sample_rate, axis=-1)
            cur_line_points = np.stack([line_points_x, line_points_y],
                                       axis=-1).reshape(-1, 2)
            # cur_line_points is of shape (n_cur_lines * sample_rate, 2)
            cur_line_points = torch.tensor(cur_line_points, dtype=torch.float,
                                           device=device)
            grid_points = keypoints_to_grid(cur_line_points, img_size)

            if self.line_score:
                # The saliency score is high when the activation are locally
                # maximal along the line (and not in a square neigborhood)
                line_desc = F.grid_sample(desc, grid_points).squeeze()
                line_desc = line_desc.reshape(-1, cur_num_lines, sample_rate)
                line_desc = line_desc.permute(1, 0, 2)
                if saliency_type == "d2_net":
                    scores = self.d2_net_saliency_score(line_desc)
                else:
                    scores = self.asl_feat_saliency_score(line_desc)
            else:
                scores = F.grid_sample(score.unsqueeze(1),
                                       grid_points).squeeze()

            # Take the most salient point in n distinct regions
            scores = scores.reshape(-1, n, n_samples_per_region)
            best = torch.max(scores, dim=2, keepdim=True)[1].cpu().numpy()
            cur_line_points = cur_line_points.reshape(-1, n,
                                                      n_samples_per_region, 2)
            cur_line_points = np.take_along_axis(
                cur_line_points, best[..., None], axis=2)[:, :, 0]

            # Pad
            cur_valid_points = np.ones((cur_num_lines, self.num_samples),
                                       dtype=bool)
            cur_valid_points[:, n:] = False
            cur_line_points = np.concatenate([
                cur_line_points,
                np.zeros((cur_num_lines, self.num_samples - n, 2), dtype=float)],
                axis=1)

            line_points[cur_mask] = cur_line_points
            valid_points[cur_mask] = cur_valid_points

        return line_points, valid_points

    def sample_line_points(self, line_seg):
        """
        Regularly sample points along each line segments, with a minimal
        distance between each point. Pad the remaining points.
        Inputs:
            line_seg: an Nx2x2 torch.Tensor.
        Outputs:
            line_points: an Nxnum_samplesx2 np.array.
            valid_points: a boolean Nxnum_samples np.array.
        """
        num_lines = len(line_seg)
        line_lengths = np.linalg.norm(line_seg[:, 0] - line_seg[:, 1], axis=1)

        # Sample the points separated by at least min_dist_pts along each line
        # The number of samples depends on the length of the line
        num_samples_lst = np.clip(line_lengths // self.min_dist_pts,
                                  2, self.num_samples)
        line_points = np.empty((num_lines, self.num_samples, 2), dtype=float)
        valid_points = np.empty((num_lines, self.num_samples), dtype=bool)
        for n in np.arange(2, self.num_samples + 1):
            # Consider all lines where we can fit up to n points
            cur_mask = num_samples_lst == n
            cur_line_seg = line_seg[cur_mask]
            line_points_x = np.linspace(cur_line_seg[:, 0, 0],
                                        cur_line_seg[:, 1, 0],
                                        n, axis=-1)
            line_points_y = np.linspace(cur_line_seg[:, 0, 1],
                                        cur_line_seg[:, 1, 1],
                                        n, axis=-1)
            cur_line_points = np.stack([line_points_x, line_points_y], axis=-1)

            # Pad
            cur_num_lines = len(cur_line_seg)
            cur_valid_points = np.ones((cur_num_lines, self.num_samples),
                                       dtype=bool)
            cur_valid_points[:, n:] = False
            cur_line_points = np.concatenate([
                cur_line_points,
                np.zeros((cur_num_lines, self.num_samples - n, 2), dtype=float)],
                axis=1)

            line_points[cur_mask] = cur_line_points
            valid_points[cur_mask] = cur_valid_points

        return line_points, valid_points

    def filter_and_match_lines_topk_gpu(self, scores, topk=10):
        # Pre-filter the pairs and keep the top k best candidate lines
        line_scores1 = scores.max(3)[0]
        valid_scores1 = line_scores1 != -1
        line_scores1 = ((line_scores1 * valid_scores1).sum(2)
                        / valid_scores1.sum(2))
        line_scores2 = scores.max(2)[0]
        valid_scores2 = line_scores2 != -1
        line_scores2 = ((line_scores2 * valid_scores2).sum(2)
                        / valid_scores2.sum(2))
        line_scores = (line_scores1 + line_scores2) / 2
        topk_lines = torch.argsort(line_scores,
                                   axis=1)[:, -topk:]
        return topk_lines.cpu().numpy()

    def filter_and_match_lines_topk(self, scores, topk=10):
        # Pre-filter the pairs and keep the top k best candidate lines
        line_scores1 = scores.max(3)
        valid_scores1 = line_scores1 != -1
        line_scores1 = ((line_scores1 * valid_scores1).sum(2)
                        / valid_scores1.sum(2))
        line_scores2 = scores.max(2)
        valid_scores2 = line_scores2 != -1
        line_scores2 = ((line_scores2 * valid_scores2).sum(2)
                        / valid_scores2.sum(2))
        line_scores = (line_scores1 + line_scores2) / 2
        topk_lines = np.argsort(line_scores,
                                axis=1)[:, -topk:]
        return topk_lines

    def filter_and_match_lines(self, scores):
        """
        Use the scores to keep the top k best lines, compute the Needleman-
        Wunsch algorithm on each candidate pairs, and keep the highest score.
        Inputs:
            scores: a (N, M, n, n) np.array containing the pairwise scores
                    of the elements to match.
        Outputs:
            matches: a (N) np.array containing the indices of the best match
        """
        # Pre-filter the pairs and keep the top k best candidate lines
        line_scores1 = scores.max(3)
        valid_scores1 = line_scores1 != -1
        line_scores1 = ((line_scores1 * valid_scores1).sum(2)
                        / valid_scores1.sum(2))
        line_scores2 = scores.max(2)
        valid_scores2 = line_scores2 != -1
        line_scores2 = ((line_scores2 * valid_scores2).sum(2)
                        / valid_scores2.sum(2))
        line_scores = (line_scores1 + line_scores2) / 2
        topk_lines = np.argsort(line_scores,
                                axis=1)[:, -self.top_k_candidates:]
        # topk_lines.shape = (n_lines1, top_k_candidates)
        top_scores = np.take_along_axis(scores, topk_lines[:, :, None, None],
                                        axis=1)

        # Consider the reversed line segments as well
        top_scores = np.concatenate([top_scores, top_scores[..., ::-1]],
                                    axis=1)

        # Compute the line distance matrix with Needleman-Wunsch algo and
        # retrieve the closest line neighbor
        n_lines1, top2k, n, m = top_scores.shape
        top_scores = top_scores.reshape(n_lines1 * top2k, n, m)
        nw_scores = self.needleman_wunsch(top_scores)
        nw_scores = nw_scores.reshape(n_lines1, top2k)
        matches = np.mod(np.argmax(nw_scores, axis=1), top2k // 2)
        matches = topk_lines[np.arange(n_lines1), matches]
        return matches

    def needleman_wunsch(self, scores):
        """
        Batched implementation of the Needleman-Wunsch algorithm.
        The cost of the InDel operation is set to 0 by subtracting the gap
        penalty to the scores.
        Inputs:
            scores: a (B, N, M) np.array containing the pairwise scores
                    of the elements to match.
        """
        b, n, m = scores.shape

        # Recalibrate the scores to get a gap score of 0
        gap = 0.1
        nw_scores = scores - gap

        # Run the dynamic programming algorithm
        nw_grid = np.zeros((b, n + 1, m + 1), dtype=float)
        for i in range(n):
            for j in range(m):
                nw_grid[:, i + 1, j + 1] = np.maximum(
                    np.maximum(nw_grid[:, i + 1, j], nw_grid[:, i, j + 1]),
                    nw_grid[:, i, j] + nw_scores[:, i, j])

        return nw_grid[:, -1, -1]

    def get_pairwise_distance(self, line_seg1, line_seg2, desc1, desc2):
        """
            Compute the OPPOSITE of the NW score for pairs of line segments
            and their corresponding descriptors.
        """
        num_lines = len(line_seg1)
        assert num_lines == len(line_seg2), "The same number of lines is required in pairwise score."
        img_size1 = (desc1.shape[2] * self.grid_size,
                     desc1.shape[3] * self.grid_size)
        img_size2 = (desc2.shape[2] * self.grid_size,
                     desc2.shape[3] * self.grid_size)
        device = desc1.device

        # Sample points regularly along each line
        line_points1, valid_points1 = self.sample_line_points(line_seg1)
        line_points2, valid_points2 = self.sample_line_points(line_seg2)
        line_points1 = torch.tensor(line_points1.reshape(-1, 2),
                                    dtype=torch.float, device=device)
        line_points2 = torch.tensor(line_points2.reshape(-1, 2),
                                    dtype=torch.float, device=device)

        # Extract the descriptors for each point
        grid1 = keypoints_to_grid(line_points1, img_size1)
        grid2 = keypoints_to_grid(line_points2, img_size2)
        desc1 = F.normalize(F.grid_sample(desc1, grid1)[0, :, :, 0], dim=0)
        desc1 = desc1.reshape(-1, num_lines, self.num_samples)
        desc2 = F.normalize(F.grid_sample(desc2, grid2)[0, :, :, 0], dim=0)
        desc2 = desc2.reshape(-1, num_lines, self.num_samples)

        # Compute the distance between line points for every pair of lines
        # Assign a score of -1 for unvalid points
        scores = torch.einsum('dns,dnt->nst', desc1, desc2).cpu().numpy()
        scores = scores.reshape(num_lines * self.num_samples,
                                self.num_samples)
        scores[~valid_points1.flatten()] = -1
        scores = scores.reshape(num_lines, self.num_samples, self.num_samples)
        scores = scores.transpose(1, 0, 2).reshape(self.num_samples, -1)
        scores[:, ~valid_points2.flatten()] = -1
        scores = scores.reshape(self.num_samples, num_lines, self.num_samples)
        scores = scores.transpose(1, 0, 2)

        # Compute the NW score for each pair of lines
        pairwise_scores = np.array([self.needleman_wunsch(s) for s in scores])
        return -pairwise_scores
