import os
from collections import OrderedDict
from typing import Optional, Tuple, Union

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import nn
from torchvision.utils import make_grid
from torchvision.utils import save_image

from boilr.nn.utils import named_leaf_modules

img_folder = None


def _unique_filename(fname: str, extension: str) -> str:

    def exists(fname):
        return os.path.exists(fname + '.' + extension)

    if not exists(fname):
        return fname
    for i in range(10000):
        t = fname + '_' + str(i)
        if not exists(t):
            return t
    raise RuntimeError("too many files ({})".format(i))


def plot_imgs(imgs: torch.Tensor,
              name: Optional[str] = None,
              extension: Optional[str] = 'png',
              colorbar: Optional[bool] = True,
              overwrite: Optional[bool] = False):
    """Saves collection of images.

    Plots collection of 1-channel images (as 3D tensor, or 4D tensor with size
    1 on the 2nd dimension) and saves it as png. If any image extends beyond
    [0, 1], all are normalized such that the minimum and maximum are 0 and 1.

    If overwrite is False, it automatically appends an integer to the filename
    to make it unique.

    Args:
        imgs (torch.Tensor):
        name (str, optional):
        extension (str, optional):
        colorbar (bool, optional):
        overwrite (bool, optional):
    """
    if imgs.dim() == 4 and imgs.size(1) == 1:
        imgs = imgs.squeeze(1)
    if imgs.dim() != 3:
        msg = ("input tensor must be 3D, or 4D with size 1 on the 2nd "
               "dimension, but has shape {}".format(imgs.shape))
        raise RuntimeError(msg)
    if img_folder is None:
        raise RuntimeError("Image folder not set")
    fname = name.replace(' ', '_')
    fname = os.path.join(img_folder, fname)
    if not overwrite:
        fname = _unique_filename(fname, extension)
    fname = fname + '.' + extension
    imgs = imgs.detach().cpu().unsqueeze(1)  # (N, 1, H, W)
    n_imgs = imgs.size(0)
    _, c = balanced_approx_factorization(n_imgs)  # grid arrangement

    # Get minimum and maximum
    low = imgs.min().item()
    high = imgs.max().item()

    # Normalize if images extend beyond the range [0, 1]
    normalize = low < 0. or high > 1.
    if normalize:
        imgs = (imgs - low) / (high - low)

    # Compute pad value, either 0 or 1
    pad_value = img_grid_pad_value(imgs)

    # Images are now in [0, 1], arrange them into a grid as they are
    # Grid has shape (3, grid_h, grid_w) and has values in [0, 1]
    grid = make_grid(imgs, nrow=c, pad_value=pad_value, normalize=False)

    if colorbar:
        # make_grid made it RGB, now take one channel only
        grid = grid[0]

        # Rescale images to original interval (now including the padding)
        if normalize:
            grid = grid * (high - low) + low

        # Save grid of images with colorbar
        vmin = low if normalize else 0.
        vmax = high if normalize else 1.
        plt.imshow(grid, vmin=vmin, vmax=vmax, cmap='gray')
        plt.colorbar()
        plt.title(name)
        plt.savefig(fname)
        plt.close()

    else:
        from PIL import Image

        # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
        grid = grid.mul(255).add_(0.5).clamp_(0, 255).to('cpu', torch.uint8)
        grid = grid.permute(1, 2, 0).numpy()

        # Save image with PIL
        im = Image.fromarray(grid)
        im.save(fname, format=None)


def img_grid_pad_value(imgs: torch.Tensor,
                       thresh: Optional[float] = .2) -> float:
    """Returns padding value (black or white) for a grid of images.

    Hack to visualize boundaries between images with torchvision's
    save_image(). If the median border value of all images is below the
    threshold, use white, otherwise black (which is the default).

    Args:
        imgs (torch.Tensor): A 4d tensor
        thresh (float, optional): Threshold in (0, 1).

    Returns:
        pad_value (float): The padding value
    """

    assert imgs.dim() == 4
    imgs = imgs.clamp(min=0., max=1.)
    assert 0. < thresh < 1.

    imgs = imgs.mean(1)  # reduce to 1 channel
    h = imgs.size(1)
    w = imgs.size(2)
    borders = list()
    borders.append(imgs[:, 0].flatten())
    borders.append(imgs[:, h - 1].flatten())
    borders.append(imgs[:, 1:h - 1, 0].flatten())
    borders.append(imgs[:, 1:h - 1, w - 1].flatten())
    borders = torch.cat(borders)
    if torch.median(borders) < thresh:
        return 1.0
    return 0.0


def save_image_grid(images: torch.Tensor,
                    filename: str,
                    nrows: Optional[int] = 8):
    """Saves a grid of images.

    Wrapper around torchvision save_image. The grid lines are either black or
    white depending on the image colors.

    Args:
        images (Tensor): batch of images (N, channels, H, W)
        filename (str): name of output file
        nrows (int, optional): number of rows and columns. The number of images
            is the square of this.
    """
    pad = img_grid_pad_value(images)
    save_image(images, filename, nrow=nrows, pad_value=pad)


def save_image_grid_reconstructions(inputs: torch.Tensor, recons: torch.Tensor,
                                    filename: str):
    """Saves a grid of images with inputs and reconstructions.

    This is meant for autoencoder-like models. The number of images is inferred
    from the input sizes. Inputs and reconstructions must have the same shape.

    Args:
        inputs (Tensor): batch of input images
        recons (Tensor): batch of reconstructions
        filename (str): name of output file
    """
    assert inputs.shape == recons.shape
    n_img = inputs.shape[0]
    n = int(np.sqrt(2 * n_img))
    imgs = torch.stack([inputs.cpu(), recons.cpu()])
    imgs = imgs.permute(1, 0, 2, 3, 4)
    imgs = imgs.reshape(n**2, inputs.size(1), inputs.size(2), inputs.size(3))
    save_image_grid(imgs, filename, nrows=n)


def balanced_approx_factorization(x: int,
                                  ratio: Optional[Union[float, int]] = 1
                                 ) -> Tuple[int, int]:
    """Approximately factorizes an integer into integers.

    Useful for example to plot images in a grid.

    Args:
        x (int): Number to be approximately factorized
        ratio (float or int, optional): Ratio columns/rows

    Returns:
        dims (tuple): A tuple (rows, columns)
    """

    # We want c/r to be approx equal to ratio, and r*c to be approx equal to x
    # ==> r = x/c = x/(ratio*r)
    # ==> r = sqrt(x/ratio) and c = sqrt(x*ratio)
    assert type(x) == int or x.dtype == int
    c = int(np.ceil(np.sqrt(x * ratio)))
    r = int(np.ceil(x / c))
    return r, c


def balanced_factorization(x: int) -> Tuple[int, int]:
    """Approximately factorizes an integer into integers.

    Args:
        x (int): Number to be approximately factorized

    Returns:
        dims (tuple): A tuple (rows, columns)
    """

    assert type(x) == int or x.dtype == int
    a = int(np.floor(np.sqrt(x)))
    while True:
        b = x / a
        if b.is_integer():
            return int(b), a
        a += 1


def clean_axes():
    """Cleans current axes."""
    plt.gca().tick_params(
        axis='both',  # changes apply to the x-axis
        which='both',  # both major and minor ticks
        left=False,  # ticks along the left edge are off
        right=False,
        bottom=False,
        top=False,
        labelleft=False,  # labels along the left edge are off
        labelbottom=False)


def _save_activation_hook(ord_dict, name):

    def hook(model, inp, ret):
        if isinstance(ret, tuple):
            try:
                ret = torch.cat(ret, dim=1)
            except TypeError:
                ret = ret[0]
            except RuntimeError as e:
                print("WARNING:", e)
                return
        ord_dict[name] = ret.detach()

    return hook


def set_up_saving_all_activations(model: nn.Module) -> OrderedDict:
    """Registers forward hooks to save all activations in a model.

    It returns a dictionary that will be populated by the forward hooks at
    the next forward pass.

    Args:
        model (torch.nn.Module)

    Returns:
        activations (dict): A dict that will be populated with all activations
            at the next forward pass.
    """
    all_activations = OrderedDict()
    for module_name, module in named_leaf_modules(model):
        module.register_forward_hook(
            _save_activation_hook(all_activations, module_name))
    return all_activations


#########################
# Test

if __name__ == '__main__':

    ### Test plot imgs

    img_folder = ''

    low = -20.
    high = 30.
    imgs = torch.rand(4, 1, 8, 8) * (high - low) + low
    plot_imgs(imgs, name='test_colorbar_false', colorbar=False)
    plot_imgs(imgs, name='test_colorbar_true', colorbar=True)

    low = .7
    high = 2.
    imgs = torch.rand(4, 1, 8, 8) * (high - low) + low
    plot_imgs(imgs, name='test_colorbar_false', colorbar=False)
    plot_imgs(imgs, name='test_colorbar_true', colorbar=True)

    low = .5
    high = .9
    imgs = torch.rand(4, 1, 8, 8) * (high - low) + low
    plot_imgs(imgs, name='test_colorbar_false', colorbar=False)
    plot_imgs(imgs, name='test_colorbar_true', colorbar=True)

    ### Test img grid pad value
    img_grid_pad_value(torch.rand(6, 3, 32, 32), thresh=.3)
