"""
A two-view sparse feature matching pipeline.

This model contains sub-models for each step:
    feature extraction, feature matching, outlier filtering, pose estimation.
Each step is optional, and the features or matches can be provided as input.
Default: SuperPoint with nearest neighbor matching.

Convention for the matches: m0[i] is the index of the keypoint in image 1
that corresponds to the keypoint i in image 0. m0[i] = -1 if i is unmatched.
"""

import numpy as np
import torch

from .. import get_model
from .base_model import BaseModel


def keep_quadrant_kp_subset(keypoints, scores, descs, h, w):
    """Keep only keypoints in one of the four quadrant of the image."""
    h2, w2 = h // 2, w // 2
    w_x = np.random.choice([0, w2])
    w_y = np.random.choice([0, h2])
    valid_mask = ((keypoints[..., 0] >= w_x)
                  & (keypoints[..., 0] < w_x + w2)
                  & (keypoints[..., 1] >= w_y)
                  & (keypoints[..., 1] < w_y + h2))
    keypoints = keypoints[valid_mask][None]
    scores = scores[valid_mask][None]
    descs = descs.permute(0, 2, 1)[valid_mask].t()[None]
    return keypoints, scores, descs


def keep_random_kp_subset(keypoints, scores, descs, num_selected):
    """Keep a random subset of keypoints."""
    num_kp = keypoints.shape[1]
    selected_kp = torch.randperm(num_kp)[:num_selected]
    keypoints = keypoints[:, selected_kp]
    scores = scores[:, selected_kp]
    descs = descs[:, :, selected_kp]
    return keypoints, scores, descs


def keep_best_kp_subset(keypoints, scores, descs, num_selected):
    """Keep the top num_selected best keypoints."""
    sorted_indices = torch.sort(scores, dim=1)[1]
    selected_kp = sorted_indices[:, -num_selected:]
    keypoints = torch.gather(keypoints, 1,
                             selected_kp[:, :, None].repeat(1, 1, 2))
    scores = torch.gather(scores, 1, selected_kp)
    descs = torch.gather(descs, 2,
                         selected_kp[:, None].repeat(1, descs.shape[1], 1))
    return keypoints, scores, descs


class TwoViewPipeline(BaseModel):
    default_conf = {
        'extractor': {
            'name': 'superpoint',
            'trainable': False,
        },
        'use_lines': False,
        'use_points': True,
        'randomize_num_kp': False,
        'detector': {'name': None},
        'descriptor': {'name': None},
        'matcher': {'name': 'nearest_neighbor_matcher'},
        'filter': {'name': None},
        'solver': {'name': None},
        'ground_truth': {
            'from_pose_depth': False,
            'from_homography': False,
            'th_positive': 3,
            'th_negative': 5,
            'reward_positive': 1,
            'reward_negative': -0.25,
            'is_likelihood_soft': True,
            'p_random_occluders': 0,
            'n_line_sampled_pts': 50,
            'line_perp_dist_th': 5,
            'overlap_th': 0.2,
            'min_visibility_th': 0.5
        },
    }
    required_data_keys = ['image0', 'image1']
    strict_conf = False  # need to pass new confs to children models
    components = [
        'extractor', 'detector', 'descriptor', 'matcher', 'filter', 'solver']

    def _init(self, conf):
        if conf.extractor.name:
            self.extractor = get_model(conf.extractor.name)(conf.extractor)
        else:
            if self.conf.detector.name:
                self.detector = get_model(conf.detector.name)(conf.detector)
            else:
                self.required_data_keys += ['keypoints0', 'keypoints1']
            if self.conf.descriptor.name:
                self.descriptor = get_model(conf.descriptor.name)(
                    conf.descriptor)
            else:
                self.required_data_keys += ['descriptors0', 'descriptors1']

        if conf.matcher.name:
            self.matcher = get_model(conf.matcher.name)(conf.matcher)
        else:
            self.required_data_keys += ['matches0']

        if conf.filter.name:
            self.filter = get_model(conf.filter.name)(conf.filter)

        if conf.solver.name:
            self.solver = get_model(conf.solver.name)(conf.solver)

    def _forward(self, data):

        def process_siamese(data, i):
            data_i = {k[:-1]: v for k, v in data.items() if k[-1] == i}
            if self.conf.extractor.name:
                pred_i = self.extractor(data_i)
            else:
                pred_i = {}
                if self.conf.detector.name:
                    pred_i = self.detector(data_i)
                else:
                    for k in ['keypoints', 'keypoint_scores', 'descriptors',
                              'lines', 'line_scores', 'line_descriptors',
                              'valid_lines']:
                        if k in data_i:
                            pred_i[k] = data_i[k]
                if self.conf.descriptor.name:
                    pred_i = {
                        **pred_i, **self.descriptor({**data_i, **pred_i})}
            return pred_i

        pred0 = process_siamese(data, '0')
        pred1 = process_siamese(data, '1')

        pred = {**{k + '0': v for k, v in pred0.items()},
                **{k + '1': v for k, v in pred1.items()}}

        if self.conf.matcher.name:
            pred = {**pred, **self.matcher({**data, **pred})}

        if self.conf.filter.name:
            pred = {**pred, **self.filter({**data, **pred})}

        if self.conf.solver.name:
            pred = {**pred, **self.solver({**data, **pred})}

        return pred

    def loss(self, pred, data):
        losses = {}
        total = 0
        for k in self.components:
            if self.conf[k].name:
                try:
                    losses_ = getattr(self, k).loss(pred, {**pred, **data})
                except NotImplementedError:
                    continue
                losses = {**losses, **losses_}
                total = losses_['total'] + total
        return {**losses, 'total': total}

    def metrics(self, pred, data):
        metrics = {}
        for k in self.components:
            if self.conf[k].name:
                try:
                    metrics_ = getattr(self, k).metrics(pred, {**pred, **data})
                except NotImplementedError:
                    continue
                metrics = {**metrics, **metrics_}
        return metrics