import logging
import math
import time

import torch

from .alnet import ALNet
from .soft_detect import DKD

configs = {
    "alike-t": {"c1": 8, "c2": 16, "c3": 32, "c4": 64, "dim": 64, "single_head": True, "radius": 2, "model_path": ""},
    "alike-s": {"c1": 8, "c2": 16, "c3": 48, "c4": 96, "dim": 96, "single_head": True, "radius": 2, "model_path": ""},
    "alike-n": {
        "c1": 16,
        "c2": 32,
        "c3": 64,
        "c4": 128,
        "dim": 128,
        "single_head": True,
        "radius": 2,
        "model_path": "",
    },
    "alike-l": {
        "c1": 32,
        "c2": 64,
        "c3": 128,
        "c4": 128,
        "dim": 128,
        "single_head": False,
        "radius": 2,
        "model_path": "",
    },
}


class ALike(ALNet):
    def __init__(
        self,
        # ================================== feature encoder
        c1: int = 32,
        c2: int = 64,
        c3: int = 128,
        c4: int = 128,
        dim: int = 128,
        single_head: bool = False,
        # ================================== detect parameters
        radius: int = 2,
        top_k: int = 500,
        scores_th: float = 0.5,
        n_limit: int = 5000,
        device: str = "cpu",
        model_path: str = "",
    ):
        super().__init__(c1, c2, c3, c4, dim, single_head)
        self.radius = radius
        self.top_k = top_k
        self.n_limit = n_limit
        self.scores_th = scores_th
        self.dkd = DKD(radius=self.radius, top_k=self.top_k, scores_th=self.scores_th, n_limit=self.n_limit)
        self.device = device

        if model_path != "":
            state_dict = torch.load(model_path, self.device)
            self.load_state_dict(state_dict)
            self.to(self.device)
            self.eval()
            logging.info(f"Loaded model parameters from {model_path}")
            logging.info(
                f"Number of model parameters: {sum(p.numel() for p in self.parameters() if p.requires_grad) / 1e3}KB"
            )
        else:
            # no model
            logging.info(f"No model loaded, please load model parameters manually")

    def extract_dense_map(self, image, ret_dict=False):
        # ====================================================
        # check image size, should be integer multiples of 2^5
        # if it is not a integer multiples of 2^5, padding zeros
        device = image.device
        b, c, h, w = image.shape
        h_ = math.ceil(h / 32) * 32 if h % 32 != 0 else h
        w_ = math.ceil(w / 32) * 32 if w % 32 != 0 else w
        if h_ != h:
            h_padding = torch.zeros(b, c, h_ - h, w, device=device)
            image = torch.cat([image, h_padding], dim=2)
        if w_ != w:
            w_padding = torch.zeros(b, c, h_, w_ - w, device=device)
            image = torch.cat([image, w_padding], dim=3)
        # ====================================================

        scores_map, descriptor_map = super().forward(image)

        # ====================================================
        if h_ != h or w_ != w:
            descriptor_map = descriptor_map[:, :, :h, :w]
            scores_map = scores_map[:, :, :h, :w]  # Bx1xHxW
        # ====================================================

        # BxCxHxW
        descriptor_map = torch.nn.functional.normalize(descriptor_map, p=2, dim=1)

        if ret_dict:
            return {
                "descriptor_map": descriptor_map,
                "scores_map": scores_map,
            }
        else:
            return descriptor_map, scores_map

    def forward(self, image: torch.Tensor, sort=False, sub_pixel=False):
        """
        :param img: [0,1] torch image tensor, BxCxHxW
        :param image_size_max: maximum image size, otherwise, the image will be resized
        :param sort: sort keypoints by scores
        :param sub_pixel: whether to use sub-pixel accuracy
        :return: a dictionary with 'keypoints', 'descriptors', 'scores', and 'time'
        """
        # H, W, three = image.shape
        H, W = image.shape[-2:]

        # ==================== extract keypoints
        start = time.time()

        with torch.no_grad():
            descriptor_map, scores_map = self.extract_dense_map(image)
            keypoints, descriptors, scores, _ = self.dkd(scores_map, descriptor_map, sub_pixel=sub_pixel)
            keypoints, descriptors, scores = keypoints[0], descriptors[0], scores[0]
            keypoints = (keypoints + 1) / 2 * keypoints.new_tensor([[W - 1, H - 1]])

        if sort:
            indices = torch.argsort(scores, descending=True)
            keypoints = keypoints[indices]
            descriptors = descriptors[indices]
            scores = scores[indices]

        end = time.time()

        return {
            "keypoints": keypoints,
            "descriptors": descriptors,
            "scores": scores,
            "scores_map": scores_map,
            "time": end - start,
        }
