import numpy as np
import paddlehub as phub
import StyleTransfer.srcTransformer.StyTR as StyTR
import StyleTransfer.srcTransformer.transformer as transformer
import tensorflow as tf
import tensorflow_hub as tfhub
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms

# TRANSFORMER

vgg_path = "StyleTransfer/srcTransformer/Transformer_models/vgg_normalised.pth"
decoder_path = "StyleTransfer/srcTransformer/Transformer_models/decoder_iter_160000.pth"
Trans_path = (
    "StyleTransfer/srcTransformer/Transformer_models/transformer_iter_160000.pth"
)
embedding_path = (
    "StyleTransfer/srcTransformer/Transformer_models/embedding_iter_160000.pth"
)


def style_transform(h, w):
    """
    This function creates a transformation for the style image,
    that crops it and formats it into a tensor.

    Parameters:
        h = height
        w = width
    Return:
        transform = transformation pipeline
    """
    transform_list = []
    transform_list.append(transforms.CenterCrop((h, w)))
    transform_list.append(transforms.ToTensor())
    transform = transforms.Compose(transform_list)
    return transform


def content_transform():
    """
    This function simply creates a transformation pipeline,
    that formats the content image into a tensor.

    Returns:
        transform = the transformation pipeline
    """
    transform_list = []
    transform_list.append(transforms.ToTensor())
    transform = transforms.Compose(transform_list)
    return transform


# This loads the network architecture already at building time
vgg = StyTR.vgg
vgg.load_state_dict(torch.load(vgg_path))
vgg = nn.Sequential(*list(vgg.children())[:44])
decoder = StyTR.decoder
Trans = transformer.Transformer()
embedding = StyTR.PatchEmbed()
# The (square) shape of the content and style image is fixed
content_size = 640
style_size = 640


def StyleTransformer(content_img: Image.Image, style_img: Image.Image) -> Image.Image:
    """
    This function creates the Transformer network and applies it on
    a content and style image to create a styled image.

    Parameters:
        content_img = the image with the content
        style_img   = the image with the style/pattern
    Returns:
        output      = an image that is a combination of both
    """

    decoder.eval()
    Trans.eval()
    vgg.eval()

    state_dict = torch.load(decoder_path)
    decoder.load_state_dict(state_dict)

    state_dict = torch.load(Trans_path)
    Trans.load_state_dict(state_dict)

    state_dict = torch.load(embedding_path)
    embedding.load_state_dict(state_dict)

    network = StyTR.StyTrans(vgg, decoder, embedding, Trans)
    network.eval()

    content_tf = content_transform()
    style_tf = style_transform(style_size, style_size)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    network.to(device)
    content = content_tf(content_img.convert("RGB"))
    style = style_tf(style_img.convert("RGB"))
    style = style.to(device).unsqueeze(0)
    content = content.to(device).unsqueeze(0)
    with torch.no_grad():
        output = network(content, style)
    output = output[0].cpu().squeeze()
    output = (
        output.mul(255)
        .add_(0.5)
        .clamp_(0, 255)
        .permute(1, 2, 0)
        .to("cpu", torch.uint8)
        .numpy()
    )
    return Image.fromarray(output)


# STYLE-FAST

style_transfer_model = tfhub.load(
    "https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2"
)


def StyleFAST(content_image: Image.Image, style_image: Image.Image) -> Image.Image:
    """
    This function applies a Fast image style transfer technique,
    which uses a pretrained model from tensorhub.

    Parameters:
        content_image   = the image with the content
        style_image     = the image with the style/pattern
    Returns:
        stylized_image  = an image that is a combination of both
    """
    content_image = (
        tf.convert_to_tensor(np.array(content_image), np.float32)[tf.newaxis, ...]
        / 255.0
    )
    style_image = (
        tf.convert_to_tensor(np.array(style_image), np.float32)[tf.newaxis, ...] / 255.0
    )
    output = style_transfer_model(content_image, style_image)
    stylized_image = output[0]
    return Image.fromarray(np.uint8(stylized_image[0] * 255))


# STYLE PROJECTION

stylepro_artistic = phub.Module(name="stylepro_artistic")


def styleProjection(
    content_image: Image.Image, style_image: Image.Image, alpha: float = 1.0
):
    """
    This function uses parameter free style transfer,
    based on a model from paddlehub.
    There is an optional weight parameter alpha, which
    allows to control the balance between image and style.

    Parameters:
        content_image   = the image with the content
        style_image     = the image with the style/pattern
        alpha           = weight for the image vs style.
                          This should be a float between 0 and 1.
    Returns:
        result          = an image that is a combination of both
    """
    result = stylepro_artistic.style_transfer(
        images=[
            {
                "content": np.array(content_image.convert("RGB"))[:, :, ::-1],
                "styles": [np.array(style_image.convert("RGB"))[:, :, ::-1]],
            }
        ],
        alpha=alpha,
    )

    return Image.fromarray(np.uint8(result[0]["data"])[:, :, ::-1]).convert("RGB")