#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

import collections
import copy
import logging
import os
from collections import defaultdict
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from termcolor import colored
from torch.nn.parallel import DataParallel, DistributedDataParallel

from cvpods.utils import PathManager


class Checkpointer(object):
    """
    A checkpointer that can save/load model as well as extra checkpointable
    objects.
    """

    def __init__(
            self,
            model: nn.Module,
            save_dir: str = "",
            resume: bool = False,
            *,
            save_to_disk: bool = True,
            **checkpointables: object,
    ):
        """
        Args:
            model (nn.Module): model.
            save_dir (str): a directory to save and find checkpoints.
            save_to_disk (bool): if True, save checkpoint to disk, otherwise
                disable saving for this checkpointer.
            checkpointables (object): any checkpointable objects, i.e., objects
                that have the `state_dict()` and `load_state_dict()` method. For
                example, it can be used like
                `Checkpointer(model, "dir", optimizer=optimizer)`.
        """
        if isinstance(model, (DistributedDataParallel, DataParallel)):
            model = model.module
        self.model = model
        self.checkpointables = copy.copy(checkpointables)
        self.resume = resume
        self.logger = logging.getLogger(__name__)
        self.save_dir = save_dir
        self.save_to_disk = save_to_disk

    def save(self, name: str, tag_checkpoint: bool = True, **kwargs: dict):
        """
        Dump model and checkpointables to a file.

        Args:
            name (str): name of the file.
            kwargs (dict): extra arbitrary data to save.
        """
        if not self.save_dir or not self.save_to_disk:
            return

        data = {}
        data["model"] = self.model.state_dict()
        for key, obj in self.checkpointables.items():
            data[key] = obj.state_dict()
        data.update(kwargs)

        basename = "{}.pth".format(name)
        save_file = os.path.join(self.save_dir, basename)
        assert os.path.basename(save_file) == basename, basename
        self.logger.info("Saving checkpoint to {}".format(save_file))
        with PathManager.open(save_file, "wb") as f:
            torch.save(data, f)

        if tag_checkpoint:
            self.tag_last_checkpoint(basename)

    def load(self, path: str, load_mapping=None):
        """
        Load from the given checkpoint. When path points to network file, this
        function has to be called on all ranks.

        Args:
            path (str): path or url to the checkpoint. If empty, will not load
                anything.
        Returns:
            dict:
                extra data loaded from the checkpoint that has not been
                processed. For example, those saved with
                :meth:`.save(**extra_data)`.
        """
        if not path:
            # no checkpoint provided
            self.logger.info(
                "No checkpoint found. Initializing model from scratch"
            )
            return {}
        self.logger.info("Loading checkpoint from {}".format(path))
        if not os.path.isfile(path):
            path = PathManager.get_local_path(path)
            assert PathManager.isfile(path), "Checkpoint {} not found!".format(path)

        checkpoint = self._load_file(path)
        self._load_model(checkpoint, load_mapping)
        if self.resume:
            for key, obj in self.checkpointables.items():
                if key in checkpoint:
                    self.logger.info("Loading {} from {}".format(key, path))
                    obj.load_state_dict(checkpoint.pop(key))
            # return any further checkpoint data
            return checkpoint
        else:
            return {}

    def has_checkpoint(self):
        """
        Returns:
            bool: whether a checkpoint exists in the target directory.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        return PathManager.exists(save_file)

    def get_checkpoint_file(self):
        """
        Returns:
            str: The latest checkpoint file in target directory.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        try:
            with PathManager.open(save_file, "r") as f:
                last_saved = f.read().strip()
        except IOError:
            # if file doesn't exist, maybe because it has just been
            # deleted by a separate process
            return ""
        return os.path.join(self.save_dir, last_saved)

    def get_all_checkpoint_files(self):
        """
        Returns:
            list: All available checkpoint files (.pth files) in target
                directory.
        """
        all_model_checkpoints = [
            os.path.join(self.save_dir, file)
            for file in PathManager.ls(self.save_dir)
            if PathManager.isfile(os.path.join(self.save_dir, file)) and file.endswith(".pth")
        ]
        return all_model_checkpoints

    def resume_or_load(self, path: str, *, resume: bool = True, load_mapping=None):
        """
        If `resume` is True, this method attempts to resume from the last
        checkpoint, if exists. Otherwise, load checkpoint from the given path.
        This is useful when restarting an interrupted training job.

        Args:
            path (str): path to the checkpoint.
            resume (bool): if True, resume from the last checkpoint if it exists.

        Returns:
            same as :meth:`load`.
        """
        if resume and self.has_checkpoint():
            path = self.get_checkpoint_file()
        return self.load(path, load_mapping)

    def tag_last_checkpoint(self, last_filename_basename: str):
        """
        Tag the last checkpoint.

        Args:
            last_filename_basename (str): the basename of the last filename.
        """
        save_file = os.path.join(self.save_dir, "last_checkpoint")
        with PathManager.open(save_file, "w") as f:
            f.write(last_filename_basename)

    def _load_file(self, f: str):
        """
        Load a checkpoint file. Can be overwritten by subclasses to support
        different formats.
        Args:
            f (str): a locally mounted file path.
        Returns:
            dict: with keys "model" and optionally others that are saved by
                the checkpointer dict["model"] must be a dict which maps strings
                to torch.Tensor or numpy arrays.
        """
        return torch.load(f, map_location=torch.device("cpu"))

    # def _load_model(self, checkpoint: Any, load_mapping=None):
    #     """
    #     Load weights from a checkpoint.
    #
    #     Args:
    #         checkpoint (Any): checkpoint contains the weights.
    #     """
    #     checkpoint_state_dict = checkpoint.pop("model")
    #
    #     self._convert_ndarray_to_tensor(checkpoint_state_dict)
    #
    #     # if the state_dict comes from a model that was wrapped in a
    #     # DataParallel or DistributedDataParallel during serialization,
    #     # remove the "module" prefix before performing the matching.
    #     _strip_prefix_if_present(checkpoint_state_dict, "module.")
    #
    #     # work around https://github.com/pytorch/pytorch/issues/24139
    #     model_state_dict = self.model.state_dict()
    #     for k in list(checkpoint_state_dict.keys()):
    #         if k in model_state_dict:
    #             shape_model = tuple(model_state_dict[k].shape)
    #             shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
    #             if shape_model != shape_checkpoint:
    #                 self.logger.warning(
    #                     "'{}' has shape {} in the checkpoint but {} in the "
    #                     "model! Skipped.".format(
    #                         k, shape_checkpoint, shape_model
    #                     )
    #                 )
    #                 checkpoint_state_dict.pop(k)
    #
    #     incompatible = self.model.load_state_dict(
    #         checkpoint_state_dict, strict=False
    #     )
    #     if incompatible.missing_keys:
    #         self.logger.info(
    #             get_missing_parameters_message(incompatible.missing_keys)
    #         )
    #     if incompatible.unexpected_keys:
    #         self.logger.info(
    #             get_unexpected_parameters_message(incompatible.unexpected_keys)
    #         )

    def _load_model(self, checkpoint, load_mapping=None):

        # some default Preprocessing
        checkpoint_state_dict = checkpoint.pop("model")

        self._convert_ndarray_to_tensor(checkpoint_state_dict)

        _strip_prefix_if_present(checkpoint_state_dict, "module.")

        """
        Strategy: suppose that the models that we will create will have prefixes appended
        to each of its keys, for example due to an extra level of nesting that the original
        pre-trained weights from ImageNet won't contain. For example, model.state_dict()
        might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains
        res2.conv1.weight. We thus want to match both parameters together.
        For that, we look for each model weight, look among all loaded keys if there is one
        that is a suffix of the current weight name, and use it if that's the case.
        If multiple matches exist, take the one with longest size
        of the corresponding name. For example, for the same model as before, the pretrained
        weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case,
        we want to match backbone[0].body.conv1.weight to conv1.weight, and
        backbone[0].body.res2.conv1.weight to res2.conv1.weight.
        """

        loaded_state_dict = checkpoint_state_dict
        model_state_dict = self.model.state_dict()
        current_keys = sorted(list(model_state_dict.keys()))
        loaded_keys = sorted(list(loaded_state_dict.keys()))
        # get a matrix of string matches, where each (i, j) entry correspond to the size of the
        # loaded_key string, if it matches
        # NOTE: Kaihua Tang, since some modules of current model will be initialized from assigned layer of
        # loaded model, we use load_mapping to do such operation
        mapped_current_keys = current_keys.copy()

        if load_mapping is not None:
            for i, key in enumerate(mapped_current_keys):
                for source_key, target_key in load_mapping.items():
                    try:
                        idx = key.index(source_key)
                    except ValueError:
                        continue
                    # print(key) 
                    # print(' ' * idx + source_key)
                    if idx is not None:
                        if key[idx - 1] != '.':
                            continue
                        elif len(key) != idx + len(source_key):
                            if key[idx + len(source_key)] != '.': 
                                continue
                        mapped_current_keys[i] = key.replace(source_key, target_key)
                        self.logger.info(
                            "MAPPING {} in current model to {} in loaded model.".format(key, mapped_current_keys[i]))

        match_matrix = [
            len(j) if i.endswith(j) and (i.index(j) == 0 or i[i.index(j) - 1] == '.') else 0 for i in mapped_current_keys for j in loaded_keys
        ]
        match_matrix = torch.as_tensor(match_matrix).view(
            len(current_keys), len(loaded_keys)
        )
        max_match_size, idxs = match_matrix.max(1)
        # remove indices that correspond to no-match
        idxs[max_match_size == 0] = -1

        # used for logging
        max_size = max([len(key) for key in current_keys]) if current_keys else 1
        max_size_loaded = max([len(key) for key in loaded_keys]) if loaded_keys else 1

        rematch_msg = "REMATCHING by mapping:\n"
        match_msg = "MATCHING by mapping:\n"
        log_str_template = "{} loaded from {} of shape {}\n"

        shape_unmatch_msg = "Some model parameters' shapes are mismatched with the checkpoint:\n"
        log_template = "{}{} shape miss-match {}{}\n"

        unmatch_msg = "Some model parameters are not in the checkpoint:\n"
        for idx_new, idx_old in enumerate(idxs.tolist()):
            if idx_old == -1:
                key = current_keys[idx_new]
                unmatch_msg += colored(f"  {key} of shape {tuple(model_state_dict[key].shape)}\n", "blue")
                continue
            key = current_keys[idx_new]
            key_old = loaded_keys[idx_old]
            # apply the auot suffix match only the shape is also matched too.
            if str(model_state_dict[key].shape) == str(loaded_state_dict[key_old].shape):
                model_state_dict[key] = loaded_state_dict[key_old]

                match_msg += colored(log_str_template.format(
                            key,
                            key_old,
                            tuple(loaded_state_dict[key_old].shape),
                        ), "green")
                # add a control gate for this logger (it's too large)
                if ((not key.startswith('module.')) and key != key_old) or (
                        key.startswith('module.') and key[7:] != key_old):
                    rematch_msg += colored(log_str_template.format(
                            key,
                            key_old,
                            tuple(loaded_state_dict[key_old].shape),
                        ), "green")
            else:
                shape_unmatch_msg += colored(log_template.format(
                            key,
                            tuple(model_state_dict[key].shape),
                            key_old,
                            tuple(loaded_state_dict[key_old].shape),
                        ), "red")


        self.logger.info(match_msg)
        self.logger.info(rematch_msg)

        incompatible = self.model.load_state_dict(
            model_state_dict, strict=True
        )
        self.logger.info(unmatch_msg)
        self.logger.info(shape_unmatch_msg)

        if incompatible.missing_keys:
            self.logger.info(
                get_missing_parameters_message(incompatible.missing_keys)
            )
        if incompatible.unexpected_keys:
            self.logger.info(
                get_unexpected_parameters_message(incompatible.unexpected_keys)
            )

    def _convert_ndarray_to_tensor(self, state_dict: dict):
        """
        In-place convert all numpy arrays in the state_dict to torch tensor.
        Args:
            state_dict (dict): a state-dict to be loaded to the model.
        """
        # model could be an OrderedDict with _metadata attribute
        # (as returned by Pytorch's state_dict()). We should preserve these
        # properties.
        for k in list(state_dict.keys()):
            if "weight_order" in k:
                continue
            v = state_dict[k]
            if not isinstance(v, np.ndarray) and not isinstance(
                    v, torch.Tensor
            ):
                raise ValueError(
                    "Unsupported type found in checkpoint! {}: {}".format(
                        k, type(v)
                    )
                )
            if not isinstance(v, torch.Tensor):
                state_dict[k] = torch.from_numpy(v)


class PeriodicCheckpointer:
    """
    Save checkpoints periodically. When `.step(iteration)` is called, it will
    execute `checkpointer.save` on the given checkpointer, if iteration is a
    multiple of period or if `max_iter` is reached.
    """

    def __init__(self, checkpointer: Any, period: int, max_iter: int = None):
        """
        Args:
            checkpointer (Any): the checkpointer object used to save
            checkpoints.
            period (int): the period to save checkpoint.
            max_iter (int): maximum number of iterations. When it is reached,
                a checkpoint named "model_final" will be saved.
        """
        self.checkpointer = checkpointer
        self.period = int(period)
        self.max_iter = max_iter

    def step(self, iteration: int, **kwargs: Any):
        """
        Perform the appropriate action at the given iteration.

        Args:
            iteration (int): the current iteration, ranged in [0, max_iter-1].
            kwargs (Any): extra data to save, same as in
                :meth:`Checkpointer.save`.
        """
        iteration = int(iteration)
        additional_state = {"iteration": iteration}
        additional_state.update(kwargs)
        if (iteration + 1) % self.period == 0:
            self.checkpointer.save(
                "model_{:07d}".format(iteration), **additional_state
            )
        if iteration >= self.max_iter - 1:
            self.checkpointer.save("model_final", **additional_state)

    def save(self, name: str, **kwargs: Any):
        """
        Same argument as :meth:`Checkpointer.save`.
        Use this method to manually save checkpoints outside the schedule.

        Args:
            name (str): file name.
            kwargs (Any): extra data to save, same as in
                :meth:`Checkpointer.save`.
        """
        self.checkpointer.save(name, **kwargs)


def get_missing_parameters_message(keys: list):
    """
    Get a logging-friendly message to report parameter names (keys) that are in
    the model but not found in a checkpoint.
    Args:
        keys (list[str]): List of keys that were not found in the checkpoint.
    Returns:
        str: message.
    """
    groups = _group_checkpoint_keys(keys)
    msg = "Some model parameters are not in the checkpoint:\n"
    msg += "\n".join(
        "  " + colored(k + _group_to_str(v), "blue") for k, v in groups.items()
    )
    return msg


def get_unexpected_parameters_message(keys: list):
    """
    Get a logging-friendly message to report parameter names (keys) that are in
    the checkpoint but not found in the model.
    Args:
        keys (list[str]): List of keys that were not found in the model.
    Returns:
        str: message.
    """
    groups = _group_checkpoint_keys(keys)
    msg = "The checkpoint contains parameters not used by the model:\n"
    msg += "\n".join(
        "  " + colored(k + _group_to_str(v), "magenta")
        for k, v in groups.items()
    )
    return msg


def _strip_prefix_if_present(state_dict: collections.OrderedDict, prefix: str):
    """
    Strip the prefix in metadata, if any.
    Args:
        state_dict (OrderedDict): a state-dict to be loaded to the model.
        prefix (str): prefix.
    """
    keys = sorted(state_dict.keys())
    if not all(len(key) == 0 or key.startswith(prefix) for key in keys):
        return

    for key in keys:
        newkey = key[len(prefix):]
        state_dict[newkey] = state_dict.pop(key)

    # also strip the prefix in metadata, if any..
    try:
        metadata = state_dict._metadata
    except AttributeError:
        pass
    else:
        for key in list(metadata.keys()):
            # for the metadata dict, the key can be:
            # '': for the DDP module, which we want to remove.
            # 'module': for the actual model.
            # 'module.xx.xx': for the rest.

            if len(key) == 0:
                continue
            newkey = key[len(prefix):]
            metadata[newkey] = metadata.pop(key)


def _group_checkpoint_keys(keys: list):
    """
    Group keys based on common prefixes. A prefix is the string up to the final
    "." in each key.
    Args:
        keys (list[str]): list of parameter names, i.e. keys in the model
            checkpoint dict.
    Returns:
        dict[list]: keys with common prefixes are grouped into lists.
    """
    groups = defaultdict(list)
    for key in keys:
        pos = key.rfind(".")
        if pos >= 0:
            head, tail = key[:pos], [key[pos + 1:]]
        else:
            head, tail = key, []
        groups[head].extend(tail)
    return groups


def _group_to_str(group: list):
    """
    Format a group of parameter name suffixes into a loggable string.
    Args:
        group (list[str]): list of parameter name suffixes.
    Returns:
        str: formated string.
    """
    if len(group) == 0:
        return ""

    if len(group) == 1:
        return "." + group[0]

    return ".{" + ", ".join(group) + "}"
