"""
Implements the full pipeline from raw images to line matches.
"""
import time
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.functional import softmax

from .model_util import get_model
from .loss import get_loss_and_weights
from .metrics import super_nms
from .line_detection import LineSegmentDetectionModule
from .line_matching import WunschLineMatcher
from ..train import convert_junc_predictions
from ..misc.train_utils import adapt_checkpoint
from .line_detector import line_map_to_segments


class LineMatcher(object):
    """Full line matcher including line detection and matching
    with the Needleman-Wunsch algorithm."""

    def __init__(
        self,
        model_cfg,
        ckpt_path,
        device,
        line_detector_cfg,
        line_matcher_cfg,
        multiscale=False,
        scales=[1.0, 2.0],
    ):
        # Get loss weights if dynamic weighting
        _, loss_weights = get_loss_and_weights(model_cfg, device)
        self.device = device

        # Initialize the cnn backbone
        self.model = get_model(model_cfg, loss_weights)
        checkpoint = torch.load(ckpt_path, map_location=self.device)
        checkpoint = adapt_checkpoint(checkpoint["model_state_dict"])
        self.model.load_state_dict(checkpoint)
        self.model = self.model.to(self.device)
        self.model = self.model.eval()

        self.grid_size = model_cfg["grid_size"]
        self.junc_detect_thresh = model_cfg["detection_thresh"]
        self.max_num_junctions = model_cfg.get("max_num_junctions", 300)

        # Initialize the line detector
        self.line_detector = LineSegmentDetectionModule(**line_detector_cfg)
        self.multiscale = multiscale
        self.scales = scales

        # Initialize the line matcher
        self.line_matcher = WunschLineMatcher(**line_matcher_cfg)

        # Print some debug messages
        for key, val in line_detector_cfg.items():
            print(f"[Debug] {key}: {val}")
        # print("[Debug] detect_thresh: %f" % (line_detector_cfg["detect_thresh"]))
        # print("[Debug] num_samples: %d" % (line_detector_cfg["num_samples"]))

    # Perform line detection and descriptor inference on a single image
    def line_detection(
        self, input_image, valid_mask=None, desc_only=False, profile=False
    ):
        # Restrict input_image to 4D torch tensor
        if (not len(input_image.shape) == 4) or (
            not isinstance(input_image, torch.Tensor)
        ):
            raise ValueError("[Error] the input image should be a 4D torch tensor")

        # Move the input to corresponding device
        input_image = input_image.to(self.device)

        # Forward of the CNN backbone
        start_time = time.time()
        with torch.no_grad():
            net_outputs = self.model(input_image)

        outputs = {"descriptor": net_outputs["descriptors"]}

        if not desc_only:
            junc_np = convert_junc_predictions(
                net_outputs["junctions"],
                self.grid_size,
                self.junc_detect_thresh,
                self.max_num_junctions,
            )
            if valid_mask is None:
                junctions = np.where(junc_np["junc_pred_nms"].squeeze())
            else:
                junctions = np.where(junc_np["junc_pred_nms"].squeeze() * valid_mask)
            junctions = np.concatenate(
                [junctions[0][..., None], junctions[1][..., None]], axis=-1
            )

            if net_outputs["heatmap"].shape[1] == 2:
                # Convert to single channel directly from here
                heatmap = (
                    softmax(net_outputs["heatmap"], dim=1)[:, 1:, :, :]
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                )
            else:
                heatmap = (
                    torch.sigmoid(net_outputs["heatmap"])
                    .cpu()
                    .numpy()
                    .transpose(0, 2, 3, 1)
                )
            heatmap = heatmap[0, :, :, 0]

            # Run the line detector.
            line_map, junctions, heatmap = self.line_detector.detect(
                junctions, heatmap, device=self.device
            )
            if isinstance(line_map, torch.Tensor):
                line_map = line_map.cpu().numpy()
            if isinstance(junctions, torch.Tensor):
                junctions = junctions.cpu().numpy()
            outputs["heatmap"] = heatmap.cpu().numpy()
            outputs["junctions"] = junctions

            # If it's a line map with multiple detect_thresh and inlier_thresh
            if len(line_map.shape) > 2:
                num_detect_thresh = line_map.shape[0]
                num_inlier_thresh = line_map.shape[1]
                line_segments = []
                for detect_idx in range(num_detect_thresh):
                    line_segments_inlier = []
                    for inlier_idx in range(num_inlier_thresh):
                        line_map_tmp = line_map[detect_idx, inlier_idx, :, :]
                        line_segments_tmp = line_map_to_segments(
                            junctions, line_map_tmp
                        )
                        line_segments_inlier.append(line_segments_tmp)
                    line_segments.append(line_segments_inlier)
            else:
                line_segments = line_map_to_segments(junctions, line_map)

            outputs["line_segments"] = line_segments

        end_time = time.time()

        if profile:
            outputs["time"] = end_time - start_time

        return outputs

    # Perform line detection and descriptor inference at multiple scales
    def multiscale_line_detection(
        self,
        input_image,
        valid_mask=None,
        desc_only=False,
        profile=False,
        scales=[1.0, 2.0],
        aggregation="mean",
    ):
        # Restrict input_image to 4D torch tensor
        if (not len(input_image.shape) == 4) or (
            not isinstance(input_image, torch.Tensor)
        ):
            raise ValueError("[Error] the input image should be a 4D torch tensor")

        # Move the input to corresponding device
        input_image = input_image.to(self.device)
        img_size = input_image.shape[2:4]
        desc_size = tuple(np.array(img_size) // 4)

        # Run the inference at multiple image scales
        start_time = time.time()
        junctions, heatmaps, descriptors = [], [], []
        for s in scales:
            # Resize the image
            resized_img = F.interpolate(input_image, scale_factor=s, mode="bilinear")

            # Forward of the CNN backbone
            with torch.no_grad():
                net_outputs = self.model(resized_img)

            descriptors.append(
                F.interpolate(
                    net_outputs["descriptors"], size=desc_size, mode="bilinear"
                )
            )

            if not desc_only:
                junc_prob = convert_junc_predictions(
                    net_outputs["junctions"], self.grid_size
                )["junc_pred"]
                junctions.append(
                    cv2.resize(
                        junc_prob.squeeze(),
                        (img_size[1], img_size[0]),
                        interpolation=cv2.INTER_LINEAR,
                    )
                )

                if net_outputs["heatmap"].shape[1] == 2:
                    # Convert to single channel directly from here
                    heatmap = softmax(net_outputs["heatmap"], dim=1)[:, 1:, :, :]
                else:
                    heatmap = torch.sigmoid(net_outputs["heatmap"])
                heatmaps.append(F.interpolate(heatmap, size=img_size, mode="bilinear"))

        # Aggregate the results
        if aggregation == "mean":
            # Aggregation through the mean activation
            descriptors = torch.stack(descriptors, dim=0).mean(0)
        else:
            # Aggregation through the max activation
            descriptors = torch.stack(descriptors, dim=0).max(0)[0]
        outputs = {"descriptor": descriptors}

        if not desc_only:
            if aggregation == "mean":
                junctions = np.stack(junctions, axis=0).mean(0)[None]
                heatmap = torch.stack(heatmaps, dim=0).mean(0)[0, 0, :, :]
                heatmap = heatmap.cpu().numpy()
            else:
                junctions = np.stack(junctions, axis=0).max(0)[None]
                heatmap = torch.stack(heatmaps, dim=0).max(0)[0][0, 0, :, :]
                heatmap = heatmap.cpu().numpy()

            # Extract junctions
            junc_pred_nms = super_nms(
                junctions[..., None],
                self.grid_size,
                self.junc_detect_thresh,
                self.max_num_junctions,
            )
            if valid_mask is None:
                junctions = np.where(junc_pred_nms.squeeze())
            else:
                junctions = np.where(junc_pred_nms.squeeze() * valid_mask)
            junctions = np.concatenate(
                [junctions[0][..., None], junctions[1][..., None]], axis=-1
            )

            # Run the line detector.
            line_map, junctions, heatmap = self.line_detector.detect(
                junctions, heatmap, device=self.device
            )
            if isinstance(line_map, torch.Tensor):
                line_map = line_map.cpu().numpy()
            if isinstance(junctions, torch.Tensor):
                junctions = junctions.cpu().numpy()
            outputs["heatmap"] = heatmap.cpu().numpy()
            outputs["junctions"] = junctions

            # If it's a line map with multiple detect_thresh and inlier_thresh
            if len(line_map.shape) > 2:
                num_detect_thresh = line_map.shape[0]
                num_inlier_thresh = line_map.shape[1]
                line_segments = []
                for detect_idx in range(num_detect_thresh):
                    line_segments_inlier = []
                    for inlier_idx in range(num_inlier_thresh):
                        line_map_tmp = line_map[detect_idx, inlier_idx, :, :]
                        line_segments_tmp = line_map_to_segments(
                            junctions, line_map_tmp
                        )
                        line_segments_inlier.append(line_segments_tmp)
                    line_segments.append(line_segments_inlier)
            else:
                line_segments = line_map_to_segments(junctions, line_map)

            outputs["line_segments"] = line_segments

        end_time = time.time()

        if profile:
            outputs["time"] = end_time - start_time

        return outputs

    def __call__(self, images, valid_masks=[None, None], profile=False):
        # Line detection and descriptor inference on both images
        if self.multiscale:
            forward_outputs = [
                self.multiscale_line_detection(
                    images[0], valid_masks[0], profile=profile, scales=self.scales
                ),
                self.multiscale_line_detection(
                    images[1], valid_masks[1], profile=profile, scales=self.scales
                ),
            ]
        else:
            forward_outputs = [
                self.line_detection(images[0], valid_masks[0], profile=profile),
                self.line_detection(images[1], valid_masks[1], profile=profile),
            ]
        line_seg1 = forward_outputs[0]["line_segments"]
        line_seg2 = forward_outputs[1]["line_segments"]
        desc1 = forward_outputs[0]["descriptor"]
        desc2 = forward_outputs[1]["descriptor"]

        # Match the lines in both images
        start_time = time.time()
        matches = self.line_matcher.forward(line_seg1, line_seg2, desc1, desc2)
        end_time = time.time()

        outputs = {"line_segments": [line_seg1, line_seg2], "matches": matches}

        if profile:
            outputs["line_detection_time"] = (
                forward_outputs[0]["time"] + forward_outputs[1]["time"]
            )
            outputs["line_matching_time"] = end_time - start_time

        return outputs