"""

Misc functions, including distributed helpers.

Mostly copy-paste from torchvision references.

"""

import os

import subprocess

import time

from collections import defaultdict, deque

import datetime

import pickle

from typing import Optional, List

import torch

import torch.distributed as dist

from torch import Tensor

# needed due to empty tensor bug in pytorch and torchvision 0.5

import torchvision

if float(torchvision.__version__[:3]) < 0.7:

    from torchvision.ops import _new_empty_tensor

    from torchvision.ops.misc import _output_size

class SmoothedValue(object):

    """Track a series of values and provide access to smoothed values over a

    window or the global series average.

    """

    def __init__(self, window_size=20, fmt=None):

        if fmt is None:

            fmt = "{median:.4f} ({global_avg:.4f})"

        self.deque = deque(maxlen=window_size)

        self.total = 0.0

        self.count = 0

        self.fmt = fmt

    def update(self, value, n=1):

        self.deque.append(value)

        self.count += n

        self.total += value * n

    def synchronize_between_processes(self):

        """

        Warning: does not synchronize the deque!

        """

        if not is_dist_avail_and_initialized():

            return

        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')

        dist.barrier()

        dist.all_reduce(t)

        t = t.tolist()

        self.count = int(t[0])

        self.total = t[1]

    @property

    def median(self):

        d = torch.tensor(list(self.deque))

        return d.median().item()

    @property

    def avg(self):

        d = torch.tensor(list(self.deque), dtype=torch.float32)

        return d.mean().item()

    @property

    def global_avg(self):

        return self.total / self.count

    @property

    def max(self):

        return max(self.deque)

    @property

    def value(self):

        return self.deque[-1]

    def __str__(self):

        return self.fmt.format(

            median=self.median,

            avg=self.avg,

            global_avg=self.global_avg,

            max=self.max,

            value=self.value)

def all_gather(data):

    """

    Run all_gather on arbitrary picklable data (not necessarily tensors)

    Args:

        data: any picklable object

    Returns:

        list[data]: list of data gathered from each rank

    """

    world_size = get_world_size()

    if world_size == 1:

        return [data]

    # serialized to a Tensor

    buffer = pickle.dumps(data)

    storage = torch.ByteStorage.from_buffer(buffer)

    tensor = torch.ByteTensor(storage).to("cuda")

    # obtain Tensor size of each rank

    local_size = torch.tensor([tensor.numel()], device="cuda")

    size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)]

    dist.all_gather(size_list, local_size)

    size_list = [int(size.item()) for size in size_list]

    max_size = max(size_list)

    # receiving Tensor from all ranks

    # we pad the tensor because torch all_gather does not support

    # gathering tensors of different shapes

    tensor_list = []

    for _ in size_list:

        tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda"))

    if local_size != max_size:

        padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda")

        tensor = torch.cat((tensor, padding), dim=0)

    dist.all_gather(tensor_list, tensor)

    data_list = []

    for size, tensor in zip(size_list, tensor_list):

        buffer = tensor.cpu().numpy().tobytes()[:size]

        data_list.append(pickle.loads(buffer))

    return data_list

def reduce_dict(input_dict, average=True):

    """

    Args:

        input_dict (dict): all the values will be reduced

        average (bool): whether to do average or sum

    Reduce the values in the dictionary from all processes so that all processes

    have the averaged results. Returns a dict with the same fields as

    input_dict, after reduction.

    """

    world_size = get_world_size()

    if world_size < 2:

        return input_dict

    with torch.no_grad():

        names = []

        values = []

        # sort the keys so that they are consistent across processes

        for k in sorted(input_dict.keys()):

            names.append(k)

            values.append(input_dict[k])

        values = torch.stack(values, dim=0)

        dist.all_reduce(values)

        if average:

            values /= world_size

        reduced_dict = {k: v for k, v in zip(names, values)}

    return reduced_dict

class MetricLogger(object):

    def __init__(self, delimiter="\t"):

        self.meters = defaultdict(SmoothedValue)

        self.delimiter = delimiter

    def update(self, **kwargs):

        for k, v in kwargs.items():

            if isinstance(v, torch.Tensor):

                v = v.item()

            assert isinstance(v, (float, int))

            self.meters[k].update(v)

    def __getattr__(self, attr):

        if attr in self.meters:

            return self.meters[attr]

        if attr in self.__dict__:

            return self.__dict__[attr]

        raise AttributeError("'{}' object has no attribute '{}'".format(

            type(self).__name__, attr))

    def __str__(self):

        loss_str = []

        for name, meter in self.meters.items():

            loss_str.append(

                "{}: {}".format(name, str(meter))

            )

        return self.delimiter.join(loss_str)

    def synchronize_between_processes(self):

        for meter in self.meters.values():

            meter.synchronize_between_processes()

    def add_meter(self, name, meter):

        self.meters[name] = meter

    def log_every(self, iterable, print_freq, header=None):

        i = 0

        if not header:

            header = ''

        start_time = time.time()

        end = time.time()

        iter_time = SmoothedValue(fmt='{avg:.4f}')

        data_time = SmoothedValue(fmt='{avg:.4f}')

        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'

        if torch.cuda.is_available():

            log_msg = self.delimiter.join([

                header,

                '[{0' + space_fmt + '}/{1}]',

                'eta: {eta}',

                '{meters}',

                'time: {time}',

                'data: {data}',

                'max mem: {memory:.0f}'

            ])

        else:

            log_msg = self.delimiter.join([

                header,

                '[{0' + space_fmt + '}/{1}]',

                'eta: {eta}',

                '{meters}',

                'time: {time}',

                'data: {data}'

            ])

        MB = 1024.0 * 1024.0

        for obj in enumerate(iterable):

            data_time.update(time.time() - end)

            yield obj

            iter_time.update(time.time() - end)

            if i % print_freq == 0 or i == len(iterable) - 1:

                eta_seconds = iter_time.global_avg * (len(iterable) - i)

                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))

                if torch.cuda.is_available():

                    print(log_msg.format(

                        i, len(iterable), eta=eta_string,

                        meters=str(self),

                        time=str(iter_time), data=str(data_time),

                        memory=torch.cuda.max_memory_allocated() / MB))

                else:

                    print(log_msg.format(

                        i, len(iterable), eta=eta_string,

                        meters=str(self),

                        time=str(iter_time), data=str(data_time)))

            i += 1

            end = time.time()

        total_time = time.time() - start_time

        total_time_str = str(datetime.timedelta(seconds=int(total_time)))

        print('{} Total time: {} ({:.4f} s / it)'.format(

            header, total_time_str, total_time / len(iterable)))

def get_sha():

    cwd = os.path.dirname(os.path.abspath(__file__))

    def _run(command):

        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()

    sha = 'N/A'

    diff = "clean"

    branch = 'N/A'

    try:

        sha = _run(['git', 'rev-parse', 'HEAD'])

        subprocess.check_output(['git', 'diff'], cwd=cwd)

        diff = _run(['git', 'diff-index', 'HEAD'])

        diff = "has uncommited changes" if diff else "clean"

        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])

    except Exception:

        pass

    message = f"sha: {sha}, status: {diff}, branch: {branch}"

    return message

def collate_fn(batch):

    batch = list(zip(*batch))

    batch[0] = nested_tensor_from_tensor_list(batch[0])

    return tuple(batch)

def _max_by_axis(the_list):

    # type: (List[List[int]]) -> List[int]

    maxes = the_list[0]

    for sublist in the_list[1:]:

        for index, item in enumerate(sublist):

            maxes[index] = max(maxes[index], item)

    return maxes

class NestedTensor(object):

    def __init__(self, tensors, mask: Optional[Tensor]):

        self.tensors = tensors

        self.mask = mask

    def to(self, device):

        # type: (Device) -> NestedTensor # noqa

        cast_tensor = self.tensors.to(device)

        mask = self.mask

        if mask is not None:

            assert mask is not None

            cast_mask = mask.to(device)

        else:

            cast_mask = None

        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):

        return self.tensors, self.mask

    def __repr__(self):

        return str(self.tensors)

def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):

    # TODO make this more general

    if tensor_list[0].ndim == 3:

        if torchvision._is_tracing():

            # nested_tensor_from_tensor_list() does not export well to ONNX

            # call _onnx_nested_tensor_from_tensor_list() instead

            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # TODO make it support different-sized images

        max_size = _max_by_axis([list(img.shape) for img in tensor_list])

        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))

        batch_shape = [len(tensor_list)] + max_size

        b, c, h, w = batch_shape

        dtype = tensor_list[0].dtype

        device = tensor_list[0].device

        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)

        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)

        for img, pad_img, m in zip(tensor_list, tensor, mask):

            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)

            m[: img.shape[1], :img.shape[2]] = False

    else:

        raise ValueError('not supported')

    return NestedTensor(tensor, mask)

# _onnx_nested_tensor_from_tensor_list() is an implementation of

# nested_tensor_from_tensor_list() that is supported by ONNX tracing.

@torch.jit.unused

def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:

    max_size = []

    for i in range(tensor_list[0].dim()):

        max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64)

        max_size.append(max_size_i)

    max_size = tuple(max_size)

    # work around for

    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)

    # m[: img.shape[1], :img.shape[2]] = False

    # which is not yet supported in onnx

    padded_imgs = []

    padded_masks = []

    for img in tensor_list:

        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]

        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))

        padded_imgs.append(padded_img)

        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)

        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)

        padded_masks.append(padded_mask.to(torch.bool))

    tensor = torch.stack(padded_imgs)

    mask = torch.stack(padded_masks)

    return NestedTensor(tensor, mask=mask)

def setup_for_distributed(is_master):

    """

    This function disables printing when not in master process

    """

    import builtins as __builtin__

    builtin_print = __builtin__.print

    def print(*args, **kwargs):

        force = kwargs.pop('force', False)

        if is_master or force:

            builtin_print(*args, **kwargs)

    __builtin__.print = print

def is_dist_avail_and_initialized():

    if not dist.is_available():

        return False

    if not dist.is_initialized():

        return False

    return True

def get_world_size():

    if not is_dist_avail_and_initialized():

        return 1

    return dist.get_world_size()

def get_rank():

    if not is_dist_avail_and_initialized():

        return 0

    return dist.get_rank()

def is_main_process():

    return get_rank() == 0

def save_on_master(*args, **kwargs):

    if is_main_process():

        torch.save(*args, **kwargs)

def init_distributed_mode(args):

    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:

        args.rank = int(os.environ["RANK"])

        args.world_size = int(os.environ['WORLD_SIZE'])

        args.gpu = int(os.environ['LOCAL_RANK'])

    elif 'SLURM_PROCID' in os.environ:

        args.rank = int(os.environ['SLURM_PROCID'])

        args.gpu = args.rank % torch.cuda.device_count()

    else:

        print('Not using distributed mode')

        args.distributed = False

        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)

    args.dist_backend = 'nccl'

    print('| distributed init (rank {}): {}'.format(

        args.rank, args.dist_url), flush=True)

    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,

                                         world_size=args.world_size, rank=args.rank)

    torch.distributed.barrier()

    setup_for_distributed(args.rank == 0)

@torch.no_grad()

def accuracy(output, target, topk=(1,)):

    """Computes the precision@k for the specified values of k"""

    if target.numel() == 0:

        return [torch.zeros([], device=output.device)]

    maxk = max(topk)

    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)

    pred = pred.t()

    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []

    for k in topk:

        correct_k = correct[:k].view(-1).float().sum(0)

        res.append(correct_k.mul_(100.0 / batch_size))

    return res

def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):

    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor

    """

    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.

    This will eventually be supported natively by PyTorch, and this

    class can go away.

    """

    if float(torchvision.__version__[:3]) < 0.7:

        if input.numel() > 0:

            return torch.nn.functional.interpolate(

                input, size, scale_factor, mode, align_corners

            )

        output_shape = _output_size(2, input, size, scale_factor)

        output_shape = list(input.shape[:-2]) + list(output_shape)

        return _new_empty_tensor(input, output_shape)

    else:

        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)