"""Tensor-utils
"""
import io
import math
from contextlib import redirect_stdout
from pathlib import Path

# from copy import copy
from threading import Thread

import numpy as np
import torch
import torch.nn as nn
from skimage import io as skio
from torch import autograd
from torch.autograd import Variable
from torch.nn import init

from climategan.utils import all_texts_to_array


def transforms_string(ts):
    return " -> ".join([t.__class__.__name__ for t in ts.transforms])


def init_weights(net, init_type="normal", init_gain=0.02, verbose=0, caller=""):
    """Initialize network weights.
    Parameters:
        net (network)     -- network to be initialized
        init_type (str)   -- the name of an initialization method:
                             normal | xavier | kaiming | orthogonal
        init_gain (float) -- scaling factor for normal, xavier and orthogonal.

    We use 'normal' in the original pix2pix and CycleGAN paper.
    But xavier and kaiming might work better for some applications.
    Feel free to try yourself.
    """

    if not init_type:
        print(
            "init_weights({}): init_type is {}, defaulting to normal".format(
                caller + " " + net.__class__.__name__, init_type
            )
        )
        init_type = "normal"
    if not init_gain:
        print(
            "init_weights({}): init_gain is {}, defaulting to normal".format(
                caller + " " + net.__class__.__name__, init_type
            )
        )
        init_gain = 0.02

    def init_func(m):
        classname = m.__class__.__name__
        if classname.find("BatchNorm2d") != -1:
            if hasattr(m, "weight") and m.weight is not None:
                init.normal_(m.weight.data, 1.0, init_gain)
            if hasattr(m, "bias") and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif hasattr(m, "weight") and (
            classname.find("Conv") != -1 or classname.find("Linear") != -1
        ):
            if init_type == "normal":
                init.normal_(m.weight.data, 0.0, init_gain)
            elif init_type == "xavier":
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == "xavier_uniform":
                init.xavier_uniform_(m.weight.data, gain=1.0)
            elif init_type == "kaiming":
                init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
            elif init_type == "orthogonal":
                init.orthogonal_(m.weight.data, gain=init_gain)
            elif init_type == "none":  # uses pytorch's default init method
                m.reset_parameters()
            else:
                raise NotImplementedError(
                    "initialization method [%s] is not implemented" % init_type
                )
            if hasattr(m, "bias") and m.bias is not None:
                init.constant_(m.bias.data, 0.0)

    if verbose > 0:
        print("initialize %s with %s" % (net.__class__.__name__, init_type))
    net.apply(init_func)


def domains_to_class_tensor(domains, one_hot=False):
    """Converts a list of strings to a 1D Tensor representing the domains

    domains_to_class_tensor(["sf", "rn"])
    >>> torch.Tensor([2, 1])

    Args:
        domain (list(str)): each element of the list should be in {rf, rn, sf, sn}
        one_hot (bool, optional): whether or not to 1-h encode class labels.
            Defaults to False.
    Raises:
        ValueError: One of the domains listed is not in {rf, rn, sf, sn}

    Returns:
        torch.Tensor: 1D tensor mapping a domain to an int (not 1-hot) or 1-hot
            domain labels in a 2D tensor
    """

    mapping = {"r": 0, "s": 1}

    if not all(domain in mapping for domain in domains):
        raise ValueError(
            "Unknown domains {} should be in {}".format(domains, list(mapping.keys()))
        )

    target = torch.tensor([mapping[domain] for domain in domains])

    if one_hot:
        one_hot_target = torch.FloatTensor(len(target), 2)  # 2 domains
        one_hot_target.zero_()
        one_hot_target.scatter_(1, target.unsqueeze(1), 1)
        # https://discuss.pytorch.org/t/convert-int-into-one-hot-format/507
        target = one_hot_target
    return target


def fake_domains_to_class_tensor(domains, one_hot=False):
    """Converts a list of strings to a 1D Tensor representing the fake domains
    (real or sim only)

    fake_domains_to_class_tensor(["s", "r"], False)
    >>> torch.Tensor([0, 2])


    Args:
        domain (list(str)): each element of the list should be in {r, s}
        one_hot (bool, optional): whether or not to 1-h encode class labels.
            Defaults to False.
    Raises:
        ValueError: One of the domains listed is not in {rf, rn, sf, sn}

    Returns:
        torch.Tensor: 1D tensor mapping a domain to an int (not 1-hot) or
            a 2D tensor filled with 0.25 to fool the classifier (equiprobability
            for each domain).
    """
    if one_hot:
        target = torch.FloatTensor(len(domains), 2)
        target.fill_(0.5)

    else:
        mapping = {"r": 1, "s": 0}

        if not all(domain in mapping for domain in domains):
            raise ValueError(
                "Unknown domains {} should be in {}".format(
                    domains, list(mapping.keys())
                )
            )

        target = torch.tensor([mapping[domain] for domain in domains])
    return target


def show_tanh_tensor(tensor):
    import skimage

    if isinstance(tensor, torch.Tensor):
        image = tensor.permute(1, 2, 0).detach().numpy()
    else:
        image = tensor
        if image.shape[-1] != 3:
            image = image.transpose(1, 2, 0)

    if image.min() < 0 and image.min() > -1:
        image = image / 2 + 0.5
    elif image.min() < -1:
        raise ValueError("can't handle this data")

    skimage.io.imshow(image)


def normalize_tensor(t):
    """
    Brings any tensor to the [0; 1] range.

    Args:
        t (torch.Tensor): input to normalize

    Returns:
        torch.Tensor: t projected to [0; 1]
    """
    t = t - torch.min(t)
    t = t / torch.max(t)
    return t


def get_normalized_depth_t(tensor, domain, normalize=False, log=True):
    assert not (normalize and log)
    if domain == "r":
        # megadepth depth
        tensor = tensor.unsqueeze(0)
        tensor = tensor - torch.min(tensor)
        tensor = torch.true_divide(tensor, torch.max(tensor))

    elif domain == "s":
        # from 3-channel depth encoding from Unity simulator to 1-channel [0-1] values
        tensor = decode_unity_depth_t(tensor, log=log, normalize=normalize)

    elif domain == "kitti":
        tensor = tensor / 100
        if not log:
            tensor = 1 / tensor
            if normalize:
                tensor = tensor - tensor.min()
                tensor = tensor / tensor.max()
        else:
            tensor = torch.log(tensor)

        tensor = tensor.unsqueeze(0)

    return tensor


def decode_bucketed_depth(tensor, opts):
    # tensor is size 1 x C x H x W
    assert tensor.shape[0] == 1
    idx = torch.argmax(tensor.squeeze(0), dim=0)  # channels become dim 0 with squeeze
    linspace_args = (
        opts.gen.d.classify.linspace.min,
        opts.gen.d.classify.linspace.max,
        opts.gen.d.classify.linspace.buckets,
    )
    indexer = torch.linspace(*linspace_args)
    log_depth = indexer[idx.long()].to(torch.float32)  # H x W
    depth = torch.exp(log_depth)
    return depth.unsqueeze(0).unsqueeze(0).to(tensor.device)


def decode_unity_depth_t(unity_depth, log=True, normalize=False, numpy=False, far=1000):
    """Transforms the 3-channel encoded depth map from our Unity simulator
    to 1-channel depth map containing metric depth values.
    The depth is encoded in the following way:
    - The information from the simulator is (1 - LinearDepth (in [0,1])).
        far corresponds to the furthest distance to the camera included in the
        depth map.
        LinearDepth * far gives the real metric distance to the camera.
    - depth is first divided in 31 slices encoded in R channel with values ranging
        from 0 to 247
    - each slice is divided again in 31 slices, whose value is encoded in G channel
    - each of the G slices is divided into 256 slices, encoded in B channel

    In total, we have a discretization of depth into N = 31*31*256 - 1 possible values,
    covering a range of far/N meters.

    Note that, what we encode here is 1 - LinearDepth so that the furthest point is
    [0,0,0] (that is sky) and the closest point[255,255,255]

    The metric distance associated to a pixel whose depth is (R,G,B) is :
        d = (far/N) * [((255 - R)//8)*256*31 + ((255 - G)//8)*256 + (255 - B)]

    * torch.Tensor in [0, 1] as torch.float32 if numpy == False

    * else numpy.array in [0, 255] as np.uint8

    Args:
        unity_depth (torch.Tensor): one depth map obtained from our simulator
        numpy (bool, optional): Whether to return a float tensor or an int array.
         Defaults to False.
        far: far parameter of the camera in Unity simulator.

    Returns:
        [torch.Tensor or numpy.array]: decoded depth
    """
    R = unity_depth[:, :, 0]
    G = unity_depth[:, :, 1]
    B = unity_depth[:, :, 2]

    R = ((247 - R) / 8).type(torch.IntTensor)
    G = ((247 - G) / 8).type(torch.IntTensor)
    B = (255 - B).type(torch.IntTensor)
    depth = ((R * 256 * 31 + G * 256 + B).type(torch.FloatTensor)) / (256 * 31 * 31 - 1)
    depth = depth * far
    if not log:
        depth = 1 / depth
    depth = depth.unsqueeze(0)  # (depth * far).unsqueeze(0)

    if log:
        depth = torch.log(depth)
    if normalize:
        depth = depth - torch.min(depth)
        depth /= torch.max(depth)
    if numpy:
        depth = depth.data.cpu().numpy()
        return depth.astype(np.uint8).squeeze()
    return depth


def to_inv_depth(log_depth, numpy=False):
    """Convert log depth tensor to inverse depth image for display

    Args:
        depth (Tensor): log depth float tensor
    """
    depth = torch.exp(log_depth)
    # visualize prediction using inverse depth, so that we don't need sky
    # segmentation (if you want to use RGB map for visualization,
    # you have to run semantic segmentation to mask the sky first
    # since the depth of sky is random from CNN)
    inv_depth = 1 / depth
    inv_depth /= torch.max(inv_depth)
    if numpy:
        inv_depth = inv_depth.data.cpu().numpy()
    # you might also use percentile for better visualization

    return inv_depth


def shuffle_batch_tuple(mbt):
    """shuffle the order of domains in the batch

    Args:
        mbt (tuple): multi-batch tuple

    Returns:
        list: randomized list of domain-specific batches
    """
    assert isinstance(mbt, (tuple, list))
    assert len(mbt) > 0
    perm = np.random.permutation(len(mbt))
    return [mbt[i] for i in perm]


def slice_batch(batch, slice_size):
    assert slice_size > 0
    for k, v in batch.items():
        if isinstance(v, dict):
            for task, d in v.items():
                batch[k][task] = d[:slice_size]
        else:
            batch[k] = v[:slice_size]
    return batch


def save_tanh_tensor(image, path):
    """Save an image which can be numpy or tensor, 2 or 3 dims (no batch)
    to path.

    Args:
        image (np.array or torch.Tensor): image to save
        path (pathlib.Path or str): where to save the image
    """
    path = Path(path)
    if isinstance(image, torch.Tensor):
        image = image.detach().cpu().numpy()
        if image.shape[-1] != 3 and image.shape[0] == 3:
            image = np.transpose(image, (1, 2, 0))
    if image.min() < 0 and image.min() > -1:
        image = image / 2 + 0.5
    elif image.min() < -1:
        image -= image.min()
        image /= image.max()
        # print("Warning: scaling image data in save_tanh_tensor")

    skio.imsave(path, (image * 255).astype(np.uint8))


def save_batch(multi_domain_batch, root="./", step=0, num_threads=5):
    root = Path(root)
    root.mkdir(parents=True, exist_ok=True)
    images_to_save = {"paths": [], "images": []}
    for domain, batch in multi_domain_batch.items():
        y = batch["data"].get("y")
        x = batch["data"]["x"]
        if y is not None:
            paths = batch["paths"]["x"]
            imtensor = torch.cat([x, y], dim=-1)
            for i, im in enumerate(imtensor):
                imid = Path(paths[i]).stem[:10]
                images_to_save["paths"] += [
                    root / "im_{}_{}_{}.png".format(step, domain, imid)
                ]
                images_to_save["images"].append(im)
    if num_threads > 0:
        threaded_write(images_to_save["images"], images_to_save["paths"], num_threads)
    else:
        for im, path in zip(images_to_save["images"], images_to_save["paths"]):
            save_tanh_tensor(im, path)


def threaded_write(images, paths, num_threads=5):
    t_im = []
    t_p = []
    for im, p in zip(images, paths):
        t_im.append(im)
        t_p.append(p)
        if len(t_im) == num_threads:
            ts = [
                Thread(target=save_tanh_tensor, args=(_i, _p))
                for _i, _p in zip(t_im, t_p)
            ]
            list(map(lambda t: t.start(), ts))
            list(map(lambda t: t.join(), ts))
            t_im = []
            t_p = []
    if t_im:
        ts = [
            Thread(target=save_tanh_tensor, args=(_i, _p)) for _i, _p in zip(t_im, t_p)
        ]
        list(map(lambda t: t.start(), ts))
        list(map(lambda t: t.join(), ts))


def get_num_params(model):
    total_params = sum(p.numel() for p in model.parameters())
    return total_params


def vgg_preprocess(batch):
    """Preprocess batch to use VGG model"""
    tensortype = type(batch.data)
    (r, g, b) = torch.chunk(batch, 3, dim=1)
    batch = torch.cat((b, g, r), dim=1)  # convert RGB to BGR
    batch = (batch + 1) * 255 * 0.5  # [-1, 1] -> [0, 255]
    mean = tensortype(batch.data.size()).cuda()
    mean[:, 0, :, :] = 103.939
    mean[:, 1, :, :] = 116.779
    mean[:, 2, :, :] = 123.680
    batch = batch.sub(Variable(mean))  # subtract mean
    return batch


def zero_grad(model: nn.Module):
    """
    Sets gradients to None. Mode efficient than model.zero_grad()
    or opt.zero_grad() according to https://www.youtube.com/watch?v=9mS1fIYj1So

    Args:
        model (nn.Module): model to zero out
    """
    for p in model.parameters():
        p.grad = None


# Take the prediction of fake and real images from the combined batch
def divide_pred(disc_output):
    """
    Divide a multiscale discriminator's output into 2 sets of tensors,
    expecting the input to the discriminator to be a concatenation
    on the batch axis of real and fake (or fake and real) images,
    effectively doubling the batch size for better batchnorm statistics

    Args:
        disc_output (list | torch.Tensor): Discriminator output to split

    Returns:
        list | torch.Tensor[type]: pair of split outputs
    """
    # https://github.com/NVlabs/SPADE/blob/master/models/pix2pix_model.py
    # the prediction contains the intermediate outputs of multiscale GAN,
    # so it's usually a list
    if type(disc_output) == list:
        half1 = []
        half2 = []
        for p in disc_output:
            half1.append([tensor[: tensor.size(0) // 2] for tensor in p])
            half2.append([tensor[tensor.size(0) // 2 :] for tensor in p])
    else:
        half1 = disc_output[: disc_output.size(0) // 2]
        half2 = disc_output[disc_output.size(0) // 2 :]

    return half1, half2


def is_tpu_available():
    _torch_tpu_available = False
    try:
        import torch_xla.core.xla_model as xm  # type: ignore

        if "xla" in str(xm.xla_device()):
            _torch_tpu_available = True
        else:
            _torch_tpu_available = False
    except ImportError:
        _torch_tpu_available = False

    return _torch_tpu_available


def get_WGAN_gradient(input, output):
    # github code reference:
    # https://github.com/caogang/wgan-gp/blob/master/gan_cifar10.py
    # Calculate the gradient that WGAN-gp needs
    grads = autograd.grad(
        outputs=output,
        inputs=input,
        grad_outputs=torch.ones(output.size()).cuda(),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    grads = grads.view(grads.size(0), -1)
    gp = ((grads.norm(2, dim=1) - 1) ** 2).mean()
    return gp


def print_num_parameters(trainer, force=False):
    if trainer.verbose == 0 and not force:
        return
    print("-" * 35)
    if trainer.G.encoder is not None:
        print(
            "{:21}:".format("num params encoder"),
            f"{get_num_params(trainer.G.encoder):12,}",
        )
    for d in trainer.G.decoders.keys():
        print(
            "{:21}:".format(f"num params decoder {d}"),
            f"{get_num_params(trainer.G.decoders[d]):12,}",
        )

    print(
        "{:21}:".format("num params painter"),
        f"{get_num_params(trainer.G.painter):12,}",
    )

    if trainer.D is not None:
        for d in trainer.D.keys():
            print(
                "{:21}:".format(f"num params discrim {d}"),
                f"{get_num_params(trainer.D[d]):12,}",
            )

    print("-" * 35)


def srgb2lrgb(x):
    x = normalize(x)
    im = ((x + 0.055) / 1.055) ** (2.4)
    im[x <= 0.04045] = x[x <= 0.04045] / 12.92
    return im


def lrgb2srgb(ims):
    if len(ims.shape) == 3:
        ims = [ims]
        stack = False
    else:
        ims = list(ims)
        stack = True

    outs = []
    for im in ims:

        out = torch.zeros_like(im)
        for k in range(3):
            temp = im[k, :, :]

            out[k, :, :] = 12.92 * temp * (temp <= 0.0031308) + (
                1.055 * torch.pow(temp, (1 / 2.4)) - 0.055
            ) * (temp > 0.0031308)
        outs.append(out)

    if stack:
        return torch.stack(outs)

    return outs[0]


def normalize(t, mini=0.0, maxi=1.0):
    """
    Normalizes a tensor to [0, 1].
    If the tensor has more than 3 dimensions, the first one
    is assumed to be the batch dimension and the tensor is
    normalized per batch element, not across the batches.

    Args:
        t (torch.Tensor): Tensor to normalize
        mini (float, optional): Min allowed value. Defaults to 0.
        maxi (float, optional): Max allowed value. Defaults to 1.

    Returns:
        torch.Tensor: The normalized tensor
    """
    if len(t.shape) == 3:
        return mini + (maxi - mini) * (t - t.min()) / (t.max() - t.min())

    batch_size = t.shape[0]
    extra_dims = [1] * (t.ndim - 1)
    min_t = t.reshape(batch_size, -1).min(1)[0].reshape(batch_size, *extra_dims)
    t = t - min_t
    max_t = t.reshape(batch_size, -1).max(1)[0].reshape(batch_size, *extra_dims)
    t = t / max_t
    return mini + (maxi - mini) * t


def retrieve_sky_mask(seg):
    """
    get the binary mask for the sky given a segmentation tensor
    of logits (N x C x H x W) or labels (N x H x W)

    Args:
        seg (torch.Tensor): Segmentation map

    Returns:
        torch.Tensor: Sky mask
    """
    if len(seg.shape) == 4:  # Predictions
        seg_ind = torch.argmax(seg, dim=1)
    else:
        seg_ind = seg

    sky_mask = seg_ind == 9
    return sky_mask


def all_texts_to_tensors(texts, width=640, height=40):
    """
    Creates a list of tensors with texts from PIL images

    Args:
        texts (list(str)): texts to write
        width (int, optional): width of individual texts. Defaults to 640.
        height (int, optional): height of individual texts. Defaults to 40.

    Returns:
        list(torch.Tensor): len(texts) tensors 3 x height x width
    """
    arrays = all_texts_to_array(texts, width, height)
    arrays = [array.transpose(2, 0, 1) for array in arrays]
    return [torch.tensor(array) for array in arrays]


def write_architecture(trainer):
    stem = "archi"
    out = Path(trainer.opts.output_path)

    # encoder
    with open(out / f"{stem}_encoder.txt", "w") as f:
        f.write(str(trainer.G.encoder))

    # decoders
    for k, v in trainer.G.decoders.items():
        with open(out / f"{stem}_decoder_{k}.txt", "w") as f:
            f.write(str(v))

    # painter
    if get_num_params(trainer.G.painter) > 0:
        with open(out / f"{stem}_painter.txt", "w") as f:
            f.write(str(trainer.G.painter))

    # discriminators
    if get_num_params(trainer.D) > 0:
        for k, v in trainer.D.items():
            with open(out / f"{stem}_discriminator_{k}.txt", "w") as f:
                f.write(str(v))

    with io.StringIO() as buf, redirect_stdout(buf):
        print_num_parameters(trainer)
        output = buf.getvalue()
        with open(out / "archi_num_params.txt", "w") as f:
            f.write(output)


def rand_perlin_2d(shape, res, fade=lambda t: 6 * t**5 - 15 * t**4 + 10 * t**3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])

    grid = (
        torch.stack(
            torch.meshgrid(
                torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])
            ),
            dim=-1,
        )
        % 1
    )
    angles = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1)
    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1)

    tile_grads = (
        lambda slice1, slice2: gradients[slice1[0] : slice1[1], slice2[0] : slice2[1]]
        .repeat_interleave(d[0], 0)
        .repeat_interleave(d[1], 1)
    )
    dot = lambda grad, shift: (  # noqa: E731
        torch.stack(
            (
                grid[: shape[0], : shape[1], 0] + shift[0],
                grid[: shape[0], : shape[1], 1] + shift[1],
            ),
            dim=-1,
        )
        * grad[: shape[0], : shape[1]]
    ).sum(dim=-1)

    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])
    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])
    t = fade(grid[: shape[0], : shape[1]])
    return math.sqrt(2) * torch.lerp(
        torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]
    )


def mix_noise(x, mask, res=(8, 3), weight=0.1):
    noise = rand_perlin_2d(x.shape[-2:], res).unsqueeze(0).unsqueeze(0).to(x.device)
    noise = noise - noise.min()
    mask = mask.repeat(1, 3, 1, 1).to(x.device).to(torch.float16)
    y = mask * (weight * noise + (1 - weight) * x) + (1 - mask) * x
    return y


def tensor_ims_to_np_uint8s(ims):
    """
    transform a CHW of NCHW tensor into a list of np.uint8 [0, 255]
    image arrays

    Args:
        ims (torch.Tensor | list): [description]
    """
    if not isinstance(ims, list):
        assert isinstance(ims, torch.Tensor)
        if ims.ndim == 3:
            ims = [ims]

    nps = []
    for t in ims:
        if t.shape[0] == 3:
            t = t.permute(1, 2, 0)
        else:
            assert t.shape[-1] == 3

        n = t.cpu().numpy()
        n = (n + 1) / 2 * 255
        nps.append(n.astype(np.uint8))

    return nps[0] if len(nps) == 1 else nps


def tensor_to_uint8_numpy_image(tensor):
    """
    Turns a BxCxHxW tensor into a numpy image:
    * normalize
    * to [0, 255]
    * detach
    * channels last
    * to uin8
    * to cpu
    * to numpy

    Args:
        tensor (torch.Tensor): Tensor to transform

    Returns:
        np.array: BxHxWxC np.uint8 array in [0, 255]
    """
    return (
        normalize(tensor, 0, 255)  # [0, 255]
        .detach()  # detach from graph if needed
        .permute(0, 2, 3, 1)  # BxHxWxC
        .to(torch.uint8)  # uint8
        .cpu()  # cpu
        .numpy()  # numpy array
    )