#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import collections
import contextlib
import json
import logging
import os
import time
from functools import partial
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn
from classy_vision.generic.distributed_util import broadcast_object, is_primary
from fvcore.common.file_io import PathManager


try:
    import apex

    apex_available = True
except ImportError:
    apex_available = False


# constants:
CHECKPOINT_FILE = "checkpoint.torch"
CPU_DEVICE = torch.device("cpu")
GPU_DEVICE = torch.device("cuda")


def is_pos_int(number: int) -> bool:
    """
    Returns True if a number is a positive integer.
    """
    return type(number) == int and number >= 0


def is_pos_float(number: float) -> bool:
    """
    Returns True if a number is a positive float.
    """
    return type(number) == float and number >= 0.0


def is_pos_int_list(l: List) -> bool:
    """
    Returns True if a list contains positive integers
    """
    return type(l) == list and all(is_pos_int(n) for n in l)


def is_pos_int_tuple(t: Tuple) -> bool:
    """
    Returns True if a tuple contains positive integers
    """
    return type(t) == tuple and all(is_pos_int(n) for n in t)


def is_long_tensor(tensor: torch.Tensor) -> bool:
    """
    Returns True if a tensor is a long tensor.
    """
    if torch.is_tensor(tensor):
        return tensor.type().endswith("LongTensor")
    else:
        return False


def is_float_tensor(tensor: torch.Tensor) -> bool:
    """
    Returns True if a tensor is a float tensor.
    """
    if torch.is_tensor(tensor):
        return tensor.type().endswith("FloatTensor")
    else:
        return False


def is_double_tensor(tensor: torch.Tensor) -> bool:
    """
    Returns True if a tensor is a double tensor.
    """
    if torch.is_tensor(tensor):
        return tensor.type().endswith("DoubleTensor")
    else:
        return False


def is_leaf(module: nn.Module) -> bool:
    """
    Returns True if module is leaf in the graph.
    """
    assert isinstance(module, nn.Module), "module should be nn.Module"
    return len(list(module.children())) == 0 or hasattr(module, "_mask")


def is_on_gpu(model: torch.nn.Module) -> bool:
    """
    Returns True if all parameters of a model live on the GPU.
    """
    assert isinstance(model, torch.nn.Module)
    on_gpu = True
    has_params = False
    for param in model.parameters():
        has_params = True
        if not param.data.is_cuda:
            on_gpu = False
    return has_params and on_gpu


def is_not_none(sample: Any) -> bool:
    """
    Returns True if sample is not None and constituents are not none.
    """
    if sample is None:
        return False

    if isinstance(sample, (list, tuple)):
        if any(s is None for s in sample):
            return False

    if isinstance(sample, dict):
        if any(s is None for s in sample.values()):
            return False
    return True


def copy_model_to_gpu(model, loss=None):
    """
    Copies a model and (optional) loss to GPU and enables cudnn benchmarking.
    For multiple gpus training, the model in DistributedDataParallel for
    distributed training.
    """
    if not torch.backends.cudnn.deterministic:
        torch.backends.cudnn.benchmark = True
    model = model.cuda()

    if loss is not None:
        loss = loss.cuda()
        return model, loss
    else:
        return model


def recursive_copy_to_device(
    value: Any, non_blocking: bool, device: torch.device
) -> Any:
    """
    Recursively searches lists, tuples, dicts and copies any object which
    supports an object.to API (e.g. tensors) to device if possible.
    Other values are passed as-is in the result.

    Note:  These are all copies, so if there are two objects that reference
    the same object, then after this call, there will be two different objects
    referenced on the device.
    """
    if isinstance(value, list) or isinstance(value, tuple):
        device_val = []
        for val in value:
            device_val.append(
                recursive_copy_to_device(val, non_blocking=non_blocking, device=device)
            )

        return device_val if isinstance(value, list) else tuple(device_val)
    elif isinstance(value, collections.abc.Mapping):
        device_val = {}
        for key, val in value.items():
            device_val[key] = recursive_copy_to_device(
                val, non_blocking=non_blocking, device=device
            )

        return device_val
    elif callable(getattr(value, "to", None)):
        return value.to(device=device, non_blocking=non_blocking)

    return value


def recursive_copy_to_gpu(value: Any, non_blocking: bool = True) -> Any:
    """
    Recursively searches lists, tuples, dicts and copies tensors to GPU if
    possible. Non-tensor values are passed as-is in the result.
    Note:  These are all copies, so if there are two objects that reference
    the same object, then after this call, there will be two different objects
    referenced on the GPU.
    """
    return recursive_copy_to_device(
        value=value, non_blocking=non_blocking, device=GPU_DEVICE
    )


@contextlib.contextmanager
def numpy_seed(seed: Optional[int], *addl_seeds: int) -> None:
    """Context manager which seeds the NumPy PRNG with the specified seed and
    restores the state afterward"""
    if seed is None:
        yield
        return
    if len(addl_seeds) > 0:
        seed = int(hash((seed, *addl_seeds)) % 1e6)
    state = np.random.get_state()
    np.random.seed(seed)
    try:
        yield
    finally:
        np.random.set_state(state)


def get_checkpoint_dict(
    task, input_args: Optional[Dict], deep_copy: bool = False
) -> Dict[str, Any]:
    assert input_args is None or isinstance(
        input_args, dict
    ), f"Unexpected input_args of type: {type(input_args)}"
    return {
        "input_args": input_args,
        "classy_state_dict": task.get_classy_state(deep_copy=deep_copy),
    }


def load_and_broadcast_checkpoint(
    checkpoint_path: str, device: torch.device = CPU_DEVICE
) -> Optional[Dict]:
    """Loads a checkpoint on master and broadcasts it to all replicas.

    This is a collective operation which needs to be run in sync on all replicas.

    See :func:`load_checkpoint` for the arguments.
    """
    if is_primary():
        checkpoint = load_checkpoint(checkpoint_path, device)
    else:
        checkpoint = None
    logging.info(f"Broadcasting checkpoint loaded from {checkpoint_path}")
    return broadcast_object(checkpoint)


def load_checkpoint(
    checkpoint_path: str, device: torch.device = CPU_DEVICE
) -> Optional[Dict]:
    """Loads a checkpoint from the specified checkpoint path.

    Args:
        checkpoint_path: The path to load the checkpoint from. Can be a file or a
            directory. If it is a directory, the checkpoint is loaded from
            :py:data:`CHECKPOINT_FILE` inside the directory.
        device: device to load the checkpoint to

    Returns:
        The checkpoint, if it exists, or None.
    """
    if not checkpoint_path:
        return None

    assert device is not None, "Please specify what device to load checkpoint on"
    assert device.type in ["cpu", "cuda"], f"Unknown device: {device}"
    if device.type == "cuda":
        assert torch.cuda.is_available()

    if not PathManager.exists(checkpoint_path):
        logging.warning(f"Checkpoint path {checkpoint_path} not found")
        return None
    if PathManager.isdir(checkpoint_path):
        checkpoint_path = f"{checkpoint_path.rstrip('/')}/{CHECKPOINT_FILE}"

    if not PathManager.exists(checkpoint_path):
        logging.warning(f"Checkpoint file {checkpoint_path} not found.")
        return None

    logging.info(f"Attempting to load checkpoint from {checkpoint_path}")
    # load model on specified device and not on saved device for model and return
    # the checkpoint
    with PathManager.open(checkpoint_path, "rb") as f:
        checkpoint = torch.load(f, map_location=device)
    logging.info(f"Loaded checkpoint from {checkpoint_path}")
    return checkpoint


def update_classy_model(
    model, model_state_dict: Dict, reset_heads: bool, strict: bool = True
) -> bool:
    """
    Updates the model with the provided model state dictionary.

    Args:
        model: ClassyVisionModel instance to update
        model_state_dict: State dict, should be the output of a call to
            ClassyVisionModel.get_classy_state().
        reset_heads: if False, uses the heads' state from model_state_dict.
        strict: if True, strictly match the module/buffer keys in current model and
            pass-in model_state_dict
    """
    try:
        if reset_heads:
            current_model_state_dict = model.get_classy_state()
            # replace the checkpointed head states with source head states
            model_state_dict["model"]["heads"] = current_model_state_dict["model"][
                "heads"
            ]
        model.set_classy_state(model_state_dict, strict=strict)
        logging.info("Model state load successful")
        return True
    except Exception:
        logging.exception("Could not load the model state")
    return False


def update_classy_state(task, state_dict: Dict) -> bool:
    """
    Updates the task with the provided task dictionary.

    Args:
        task: ClassyTask instance to update
        state_dict: State dict, should be the output of a call to
            ClassyTask.get_classy_state().
    """
    logging.info("Loading classy state from checkpoint")

    try:
        task.set_classy_state(state_dict)
        logging.info("Checkpoint load successful")
        return True
    except Exception:
        logging.exception("Could not load the checkpoint")

    return False


def save_checkpoint(checkpoint_folder, state, checkpoint_file=CHECKPOINT_FILE):
    """
    Saves a state variable to the specified checkpoint folder. Returns the filename
    of the checkpoint if successful. Raises an exception otherwise.
    """

    # make sure that we have a checkpoint folder:
    if not PathManager.isdir(checkpoint_folder):
        try:
            PathManager.mkdirs(checkpoint_folder)
        except BaseException:
            logging.warning("Could not create folder %s." % checkpoint_folder)
            raise

    # write checkpoint atomically:
    try:
        full_filename = f"{checkpoint_folder}/{checkpoint_file}"
        with PathManager.open(full_filename, "wb") as f:
            torch.save(state, f)
        return full_filename
    except BaseException:
        logging.warning(
            "Unable to write checkpoint to %s." % checkpoint_folder, exc_info=True
        )
        raise


def flatten_dict(value_dict: Dict, prefix="", sep="_") -> Dict:
    """
    Flattens nested dict into (key, val) dict. Used for flattening meters
    structure, so that they can be visualized.
    """
    items = []
    for k, v in value_dict.items():
        key = prefix + sep + k if prefix else k
        if isinstance(v, collections.MutableMapping):
            items.extend(flatten_dict(value_dict=v, prefix=key, sep=sep).items())
        else:
            items.append((key, v))
    return dict(items)


def load_json(json_path):
    """
    Loads a json config from a file.
    """
    assert os.path.exists(json_path), "Json file %s not found" % json_path
    json_file = open(json_path)
    json_config = json_file.read()
    json_file.close()
    try:
        config = json.loads(json_config)
    except BaseException as err:
        raise Exception("Failed to validate config with error: %s" % str(err))

    return config


@contextlib.contextmanager
def torch_seed(seed: Optional[int]):
    """Context manager which seeds the PyTorch PRNG with the specified seed and
    restores the state afterward. Setting seed to None is equivalent to running
    the code without the context manager."""
    if seed is None:
        yield
        return
    state = torch.get_rng_state()
    torch.manual_seed(seed)
    try:
        yield
    finally:
        torch.set_rng_state(state)


def convert_to_one_hot(targets: torch.Tensor, classes) -> torch.Tensor:
    """
    This function converts target class indices to one-hot vectors,
    given the number of classes.

    """
    assert (
        torch.max(targets).item() < classes
    ), "Class Index must be less than number of classes"
    one_hot_targets = torch.zeros(
        (targets.shape[0], classes), dtype=torch.long, device=targets.device
    )
    one_hot_targets.scatter_(1, targets.long(), 1)
    return one_hot_targets


def maybe_convert_to_one_hot(
    target: torch.Tensor, model_output: torch.Tensor
) -> torch.Tensor:
    """
    This function infers whether target is integer or 0/1 encoded
    and converts it to 0/1 encoding if necessary.
    """
    target_shape_list = list(target.size())

    if len(target_shape_list) == 1 or (
        len(target_shape_list) == 2 and target_shape_list[1] == 1
    ):
        target = convert_to_one_hot(target.view(-1, 1), model_output.shape[1])

    # target are not necessarily hard 0/1 encoding. It can be soft
    # (i.e. fractional) in some cases, such as mixup label
    assert (
        target.shape == model_output.shape
    ), "Target must of the same shape as model_output."

    return target


def get_model_dummy_input(
    model,
    input_shape: Any,
    input_key: Union[str, List[str]],
    batchsize: int = 1,
    non_blocking: bool = False,
) -> Any:

    # input_shape with type dict of dict
    # e.g. {"key_1": {"key_1_1": [2, 3], "key_1_2": [4, 5, 6], "key_1_3": []}
    if isinstance(input_shape, dict):
        input = {}
        for key, value in input_shape.items():
            input[key] = get_model_dummy_input(
                model, value, input_key, batchsize, non_blocking
            )
    elif isinstance(input_key, list):
        # av mode, with multiple input keys
        input = {}
        for i, key in enumerate(input_key):
            shape = (batchsize,) + tuple(input_shape[i])
            cur_input = torch.zeros(shape)
            if next(model.parameters()).is_cuda:
                cur_input = cur_input.cuda(non_blocking=non_blocking)
            input[key] = cur_input
    else:
        # add a dimension to represent minibatch axis
        shape = (batchsize,) + tuple(input_shape)
        input = torch.zeros(shape)
        if next(model.parameters()).is_cuda:
            input = input.cuda(non_blocking=non_blocking)
        if input_key:
            input = {input_key: input}
    return input


def get_batchsize_per_replica(x: Union[Tuple, List, Dict]) -> int:
    """
    Some layer may take tuple/list/dict/list[dict] as input in forward function. We
    recursively dive into the tuple/list until we meet a tensor and infer the batch size
    """
    while isinstance(x, (list, tuple)):
        assert len(x) > 0, "input x of tuple/list type must have at least one element"
        x = x[0]

    if isinstance(x, (dict,)):
        # index zero is always equal to batch size. select an arbitrary key.
        key_list = list(x.keys())
        x = x[key_list[0]]

    return x.size()[0]


def split_batchnorm_params(model: nn.Module):
    """Finds the set of BatchNorm parameters in the model.

    Recursively traverses all parameters in the given model and returns a tuple
    of lists: the first element is the set of batchnorm parameters, the second
    list contains all other parameters of the model."""
    batchnorm_params = []
    other_params = []
    for module in model.modules():
        # If module has children (i.e. internal node of constructed DAG) then
        # only add direct parameters() to the list of params, else go over
        # children node to find if they are BatchNorm or have "bias".
        if list(module.children()) != []:
            for params in module.parameters(recurse=False):
                if params.requires_grad:
                    other_params.append(params)
        elif isinstance(module, nn.modules.batchnorm._BatchNorm):
            for params in module.parameters():
                if params.requires_grad:
                    batchnorm_params.append(params)
        else:
            for params in module.parameters():
                if params.requires_grad:
                    other_params.append(params)
    return batchnorm_params, other_params


class Timer:
    """Timer context manager to get the elapsed time for a code block.

    Example:
        .. code-block:: python

            with Timer() as timer:
                do_something()
            elapsed_time = timer.elapsed_time
    """

    def __init__(self):
        self.start = 0
        self.elapsed_time = 0

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, *args, **kwargs):
        self.elapsed_time = time.perf_counter() - self.start


@contextlib.contextmanager
def _train_mode(model: nn.Module, train_mode: bool):
    """Context manager which sets the train mode of a model. After returning, it
    restores the state of every sub-module individually."""
    train_modes = {}
    for name, module in model.named_modules():
        train_modes[name] = module.training
    try:
        model.train(train_mode)
        yield
    finally:
        for name, module in model.named_modules():
            module.training = train_modes[name]


def log_class_usage(component_type, klass):
    """This function is used to log the usage of different Classy components."""
    identifier = "ClassyVision"
    if klass and hasattr(klass, "__name__"):
        identifier += f".{component_type}.{klass.__name__}"
    torch._C._log_api_usage_once(identifier)


def get_torch_version():
    """Get the torch version as [major, minor].

    All comparisons must be done with the two version values. Revisions are not
    supported.
    """
    version_list = torch.__version__.split(".")[:2]
    return [int(version_str) for version_str in version_list]


train_model = partial(_train_mode, train_mode=True)
train_model.__doc__ = """Context manager which puts the model in train mode.

    After returning, it restores the state of every sub-module individually.
    """


eval_model = partial(_train_mode, train_mode=False)
eval_model.__doc__ = """Context manager which puts the model in eval mode.

    After returning, it restores the state of every sub-module individually.
    """


def master_params(optimizer):
    """Generator to iterate over all parameters in the optimizer param_groups.

    When apex is available, uses that to guarantee we get the FP32 copy of the
    parameters when O2 is enabled. Otherwise, iterate ourselves."""
    if apex_available:
        yield from apex.amp.master_params(optimizer)
    else:
        for group in optimizer.param_groups:
            for p in group["params"]:
                yield p
