"""
authors: Richard Osuala, Zuzanna Szafranowska
BCN-AIM 2021
"""

import logging
import os
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel


class BaseGenerator(nn.Module):
    def __init__(
        self,
        nz: int,
        ngf: int,
        nc: int,
        ngpu: int,
        leakiness: float = 0.2,
        bias: bool = False,
    ):
        super(BaseGenerator, self).__init__()
        self.nz = nz
        self.ngf = ngf
        self.nc = nc
        self.ngpu = ngpu
        self.leakiness = leakiness
        self.bias = bias
        self.main = None

    def forward(self, input):
        raise NotImplementedError


class Generator(BaseGenerator):
    def __init__(
        self,
        nz: int,
        ngf: int,
        nc: int,
        ngpu: int,
        image_size: int,
        conditional: bool,
        leakiness: float,
        bias: bool = False,
        n_cond: int = 10,
        is_condition_categorical: bool = False,
        num_embedding_dimensions: int = 50,
    ):
        super(Generator, self).__init__(
            nz=nz,
            ngf=ngf,
            nc=nc,
            ngpu=ngpu,
            leakiness=leakiness,
            bias=bias,
        )
        # if is_condition_categorical is False, we model the condition as continous input to the network
        self.is_condition_categorical = is_condition_categorical

        # n_cond is only used if is_condition_categorical is True.
        self.num_embedding_input = n_cond

        # num_embedding_dimensions is only used if is_condition_categorical is True.
        # num_embedding_dimensions standard would be dim(z), but atm we have a nn.Linear after
        # nn.Embedding that upscales the dimension to self.nz. Using same value of num_embedding_dims in D and G.
        self.num_embedding_dimensions = num_embedding_dimensions

        # whether the is a conditional input into the GAN i.e. cGAN
        self.conditional: bool = conditional

        # The image size (supported params should be 128 or 64)
        self.image_size = image_size

        if self.image_size == 128:
            self.first_layers = nn.Sequential(
                # input is Z, going into a convolution
                nn.ConvTranspose2d(
                    self.nz * self.nc, self.ngf * 16, 4, 1, 0, bias=self.bias
                ),
                nn.BatchNorm2d(self.ngf * 16),
                nn.ReLU(True),
                # state size. (ngf*16) x 4 x 4
                nn.ConvTranspose2d(
                    self.ngf * 16, self.ngf * 8, 4, 2, 1, bias=self.bias
                ),
                nn.BatchNorm2d(self.ngf * 8),
                nn.ReLU(True),
            )
        elif self.image_size == 64:
            self.first_layers = nn.Sequential(
                # input is Z, going into a convolution
                nn.ConvTranspose2d(
                    self.nz * self.nc, self.ngf * 8, 4, 1, 0, bias=self.bias
                ),
                nn.BatchNorm2d(self.ngf * 8),
                nn.ReLU(True),
            )
        else:
            raise ValueError(
                f"Allowed image sizes are 128 and 64. You provided {self.image_size}. Please adjust."
            )

        self.main = nn.Sequential(
            *self.first_layers.children(),
            # state size. (ngf*8) x 8 x 8
            nn.ConvTranspose2d(self.ngf * 8, self.ngf * 4, 4, 2, 1, bias=self.bias),
            nn.BatchNorm2d(self.ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 16 x 16
            nn.ConvTranspose2d(self.ngf * 4, self.ngf * 2, 4, 2, 1, bias=self.bias),
            nn.BatchNorm2d(self.ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 32 x 32
            nn.ConvTranspose2d(self.ngf * 2, self.ngf, 4, 2, 1, bias=self.bias),
            nn.BatchNorm2d(self.ngf),
            nn.ReLU(True),
            # state size. (ngf) x 64 x 64
            # Note that out_channels=1 instead of out_channels=self.nc.
            # This is due to conditional input channel of our grayscale images
            nn.ConvTranspose2d(
                in_channels=self.ngf,
                out_channels=1,
                kernel_size=4,
                stride=2,
                padding=1,
                bias=self.bias,
            ),
            nn.Tanh(),
            # state size. (nc) x 128 x 128
        )

        if self.is_condition_categorical:
            self.embed_nn = nn.Sequential(
                # e.g. condition -> int -> embedding -> fcl -> feature map -> concat with image -> conv layers..
                # embedding layer
                nn.Embedding(
                    num_embeddings=self.num_embedding_input,
                    embedding_dim=self.num_embedding_dimensions,
                ),
                # target output dim of dense layer is batch_size x self.nz x 1 x 1
                # input is dimension of the embedding layer output
                nn.Linear(
                    in_features=self.num_embedding_dimensions, out_features=self.nz
                ),
                # nn.BatchNorm1d(self.nz),
                nn.LeakyReLU(self.leakiness, inplace=True),
            )
        else:
            self.embed_nn = nn.Sequential(
                # target output dim of dense layer is: nz x 1 x 1
                # input is dimension of the numbers of embedding
                nn.Linear(in_features=1, out_features=self.nz),
                # TODO Ablation: How does BatchNorm1d affect the conditional model performance?
                nn.BatchNorm1d(self.nz),
                nn.LeakyReLU(self.leakiness, inplace=True),
            )

    def forward(self, x, conditions=None):
        if self.conditional:
            # combining condition labels and input images via a new image channel
            if not self.is_condition_categorical:
                # If labels are continuous (not modelled as categorical), use floats instead of integers for labels.
                # Also adjust dimensions to (batch_size x 1) as needed for input into linear layer
                # labels should already be of type float, no change expected in .float() conversion (it is only a safety check)

                # Just for testing:
                conditions *= 0
                conditions += 1

                conditions = conditions.view(conditions.size(0), -1).float()
            embedded_conditions = self.embed_nn(conditions)
            embedded_conditions_with_random_noise_dim = embedded_conditions.view(
                -1, self.nz, 1, 1
            )
            x = torch.cat([x, embedded_conditions_with_random_noise_dim], 1)
        return self.main(x)


def interval_mapping(image, from_min, from_max, to_min, to_max):
    # map values from [from_min, from_max] to [to_min, to_max]
    # image: input array
    from_range = from_max - from_min
    to_range = to_max - to_min
    # scale to interval [0,1]
    scaled = np.array((image - from_min) / float(from_range), dtype=float)
    # multiply by range and add minimum to get interval [min,range+min]
    return to_min + (scaled * to_range)


def image_generator(model_path, device, nz, ngf, nc, ngpu, num_samples):
    # instantiate the model
    logging.debug("Instantiating model...")
    netG = Generator(
        nz=nz,
        ngf=ngf,
        nc=nc,
        ngpu=ngpu,
        image_size=128,
        leakiness=0.1,
        conditional=False,
    )
    if device.type == "cuda":
        netG.cuda()

    # load the model's weights from state_dict *'.pt file
    logging.debug(f"Loading model weights from {model_path} ...")

    checkpoint = torch.load(model_path, map_location=device)
    try:
        netG.load_state_dict(state_dict=checkpoint["generator"])
    except KeyError:
        raise KeyError(
            f"checkpoint['generator_state_dict'] was not found."
        )  # checkpoint={checkpoint}")
    logging.debug(f"Using retrieved model from generator_state_dict checkpoint")
    netG.eval()

    # generate the images
    logging.debug(f"Generating {num_samples} images using {device}...")
    z = torch.randn(num_samples, nz, 1, 1, device=device)
    images = netG(z).detach().cpu().numpy()
    image_list = []
    for j, img_ in enumerate(images):
        image_list.append(img_)
    return image_list


def save_generated_images(image_list, path):
    logging.debug(f"Saving generated images now in {path}")
    for i, img_ in enumerate(image_list):
        Path(path).mkdir(parents=True, exist_ok=True)
        img_path = f"{path}/{i}.png"
        img_ = interval_mapping(img_.transpose(1, 2, 0), -1.0, 0.0, 0, 255)
        img_ = img_.astype("uint8")
        cv2.imwrite(img_path, img_)
    logging.debug(f"Saved generated images to {path}")


def return_images(image_list):
    # logging.debug(f"Returning generated images as {type(image_list)}.")
    processed_image_list = []
    for i, img_ in enumerate(image_list):
        img_ = interval_mapping(img_.transpose(1, 2, 0), -1.0, 0.0, 0, 255)
        img_ = img_.astype("uint8")
        processed_image_list.append(img_)
    return processed_image_list


def generate(model_file, num_samples, output_path, save_images: bool):
    """This function generates synthetic images of mammography regions of interest"""
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        ngpu = 0
        if device == "cuda":
            ngpu = 1
        image_list = image_generator(model_file, device, 100, 64, 1, ngpu, num_samples)
        if save_images:
            save_generated_images(image_list, output_path)
        else:
            return return_images(image_list)
    except Exception as e:
        logging.error(
            f"Error while trying to generate {num_samples} images with model {model_file}: {e}"
        )
        raise e