import os
from posixpath import basename

import torch
import gradio as gr

import os
import sys
import numpy as np

from e4e.models.psp import pSp
from util import *
from huggingface_hub import hf_hub_download

import os
import sys
import tempfile
import shutil
from argparse import Namespace
from pathlib import Path
import shutil

import dlib
import numpy as np
import torchvision.transforms as transforms
from torchvision import utils
from PIL import Image

from model.sg2_model import Generator
from generate_videos import generate_frames, video_from_interpolations, vid_to_gif

model_dir = "models"
os.makedirs(model_dir, exist_ok=True)

model_repos = {"e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
               "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
               "base": ("akhaliq/jojogan-stylegan2-ffhq-config-f", "stylegan2-ffhq-config-f.pt"),
               "anime": ("rinong/stylegan-nada-models", "anime.pt"),
               "joker": ("rinong/stylegan-nada-models", "joker.pt"),
            #    "simpson": ("rinong/stylegan-nada-models", "simpson.pt"),
            #    "ssj": ("rinong/stylegan-nada-models", "ssj.pt"),
            #    "white_walker": ("rinong/stylegan-nada-models", "white_walker.pt"),
            #    "zuckerberg": ("rinong/stylegan-nada-models", "zuckerberg.pt"),
            #    "cubism": ("rinong/stylegan-nada-models", "cubism.pt"),
            #    "disney_princess": ("rinong/stylegan-nada-models", "disney_princess.pt"),
            #    "edvard_munch": ("rinong/stylegan-nada-models", "edvard_munch.pt"),
            #    "van_gogh": ("rinong/stylegan-nada-models", "van_gogh.pt"),
            #    "oil": ("rinong/stylegan-nada-models", "oil.pt"),
            #    "rick_morty": ("rinong/stylegan-nada-models", "rick_morty.pt"),
            #    "botero": ("rinong/stylegan-nada-models", "botero.pt"),
            #    "crochet": ("rinong/stylegan-nada-models", "crochet.pt"),
            #    "modigliani": ("rinong/stylegan-nada-models", "modigliani.pt"),
            #    "shrek": ("rinong/stylegan-nada-models", "shrek.pt"),
            #    "sketch": ("rinong/stylegan-nada-models", "sketch.pt"),
            #    "thanos": ("rinong/stylegan-nada-models", "thanos.pt"),
               }

def get_models():
    os.makedirs(model_dir, exist_ok=True)

    model_paths = {}

    for model_name, repo_details in model_repos.items():
        download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
        model_paths[model_name] = download_path

    return model_paths

model_paths = get_models()

class ImageEditor(object):
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        latent_size = 512
        n_mlp = 8
        channel_mult = 2
        model_size = 1024

        self.generators = {}

        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]

        for model in self.model_list:
            g_ema = Generator(
                model_size, latent_size, n_mlp, channel_multiplier=channel_mult
            ).to(self.device)

            checkpoint = torch.load(model_paths[model], map_location=self.device)

            g_ema.load_state_dict(checkpoint['g_ema'])

            self.generators[model] = g_ema

        self.experiment_args = {"model_path": model_paths["e4e"]}
        self.experiment_args["transform"] = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        self.resize_dims = (256, 256)

        model_path = self.experiment_args["model_path"]

        ckpt = torch.load(model_path, map_location="cpu")
        opts = ckpt["opts"]

        opts["checkpoint_path"] = model_path
        opts = Namespace(**opts)

        self.e4e_net = pSp(opts, self.device)
        self.e4e_net.eval()

        self.shape_predictor = dlib.shape_predictor(
            model_paths["dlib"]
        )

        print("setup complete")

    def get_style_list(self):
        # style_list = ['all', 'list - enter below']
        style_list = []

        for key in self.generators:
            style_list.append(key)

        return style_list

    def invert_image(self, input_image):
        input_image = self.run_alignment(str(input_image))
        
        input_image = input_image.resize(self.resize_dims)

        img_transforms = self.experiment_args["transform"]
        transformed_image = img_transforms(input_image)

        with torch.no_grad():
            images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
            result_image, latent = images[0], latents[0]

        inverted_latent = latent.unsqueeze(0).unsqueeze(1)

        return inverted_latent
    
    def get_generators_for_styles(self, output_styles, loop_styles=False):

        # if style_string:
        #     styles = style_string.split(",")
        #     for style in styles:
        #         if style not in self.model_list:
        #             raise ValueError(f"Encountered style '{style}' in the input style list which is not an available option.")
        # else:
        #     styles = style_checkbox_list

        if "base" in output_styles:                                    # always start with base if chosen
            output_styles.insert(0, output_styles.pop(output_styles.index("base")))
        if loop_styles:
            output_styles.append(output_styles[0])

        return [self.generators[style] for style in output_styles]

    def edit_image(self, input, output_styles):
        return self.predict(input, output_styles)

    def edit_video(self, input, output_styles, with_editing, video_format, loop_styles):
        return self.predict(input, output_styles, True, with_editing, video_format, loop_styles)

    def predict(
        self,
        input,                  # Input image path
        output_styles,          # Style checkbox options.
        generate_video = False, # Generate a video instead of an output image
        with_editing   = False, # Apply latent space editing to the generated video
        video_format   = "mp4", # Choose gif to display in browser, mp4 for higher-quality downloadable video
        loop_styles    = False, # Loop back to the initial style
    ):  

        # @title Align image
        out_dir = Path(tempfile.mkdtemp())
        out_path = out_dir / "out.jpg"
        
        inverted_latent = self.invert_image(input)
        generators = self.get_generators_for_styles(output_styles, loop_styles)

        if not generate_video:
            with torch.no_grad():
                img_list = []
                for g_ema in generators:
                    img, _ = g_ema(inverted_latent, input_is_latent=True, truncation=1, randomize_noise=False)
                    img_list.append(img)
                
                out_img = torch.cat(img_list, axis=0)
                utils.save_image(out_img, out_path, nrow=int(np.sqrt(out_img.size(0))), normalize=True, scale_each=True, range=(-1, 1))

            return str(out_path)
        
        return self.generate_vid(generators, inverted_latent, out_dir, video_format, with_editing)

    def generate_vid(self, generators, latent, out_dir, video_format, with_editing):      
        np_latent = latent.squeeze(0).cpu().detach().numpy()
        args = {
                'fps': 24,
                'target_latents': None,
                'edit_directions': None,
                'unedited_frames': 0 if with_editing else 40 * (len(generators) - 1)
                }

        args = Namespace(**args)
        with tempfile.TemporaryDirectory() as dirpath:

            generate_frames(args, np_latent, generators, dirpath)
            video_from_interpolations(args.fps, dirpath)
            
            gen_path = Path(dirpath) / "out.mp4"
            out_path = out_dir / f"out.{video_format}"

            if video_format == 'gif':
                vid_to_gif(gen_path, out_dir, scale=256, fps=args.fps)
            else:
                shutil.copy2(gen_path, out_path) 

        return str(out_path)

    def run_alignment(self, image_path):
        aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
        print("Aligned image has shape: {}".format(aligned_image.size))
        return aligned_image

    def run_on_batch(self, inputs):
        images, latents = self.e4e_net(
            inputs.to(self.device).float(), randomize_noise=False, return_latents=True
        )
        return images, latents

editor = ImageEditor()

# def change_component_visibility(component_types, invert_choices):

#     def visibility_impl(visible):
#         return [component_types[idx].update(visible=visible ^ invert_choices[idx]) for idx in range(len(component_types))]
    
#     return visibility_impl

def group_visibility(visible):
    print("visible: ", visible)

    return gr.Group.update(visibile=visible)

blocks = gr.Blocks()

with blocks:
    gr.Markdown("<h1><center>StyleGAN-NADA</center></h1>")
    gr.Markdown(
        "Demo for StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators (SIGGRAPH 2022)."
    )
    gr.Markdown(
        "For more information about the paper and code for training your own models (with examples OR text), see below."
    )

    # with gr.Row():

    #     with gr.Column():
    #         input_img    = gr.inputs.Image(type="filepath", label="Input image")
    #         style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")

    #         video_choice = gr.inputs.Checkbox(default=False, label="Generate Video?", optional=False)
            
    #         video_options_group = gr.Group()
            
    #         with video_options_group:
    #             edit_choice       = gr.inputs.Checkbox(default=False, label="With Editing?", optional=False)
    #             vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
        
    #         # img_button = gr.Button("Edit Image")
    #         # vid_button = gr.Button("Generate Video")
    #         img_button = gr.Button("Edit Image")
    #         vid_button = gr.Button("Generate Video")

    #     with gr.Column():
    #         img_output = gr.outputs.Image(type="file")
    #         vid_output = gr.outputs.Video()

    # # visibility_fn = change_component_visibility(component_types=[gr.Checkbox, gr.Radio, gr.Video, gr.Button, gr.Image, gr.Button],
    # #                                             invert_choices=[False, False, False, False, True, True])

    # # video_choice.change(fn=visibility_fn, inputs=video_choice, outputs=[edit_choice, vid_format_choice, vid_output, vid_button, img_output, img_button])
    # video_choice.change(fn=group_visibility, inputs=video_choice, outputs=video_options_group)
    # img_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=img_output)
    # vid_button.click(fn=editor.predict, inputs=[input_img, style_choice, video_choice, edit_choice, vid_format_choice], outputs=vid_output)
    
    with gr.Row():
        input_img = gr.inputs.Image(type="filepath", label="Input image")    
        style_choice = gr.inputs.CheckboxGroup(choices=editor.get_style_list(), type="value", label="Choose your styles!")

    with gr.Tabs():
        with gr.TabItem("Edit Images"):
            with gr.Column():
                img_button = gr.Button("Edit Image")
            with gr.Column():
                img_output = gr.outputs.Image(type="file", label="Output Image")
            
        with gr.TabItem("Create Video"):
            with gr.Column():
                with gr.Row():
                    vid_button = gr.Button("Generate Video")
                    loop_styles       = gr.inputs.Checkbox(default=True, label="Loop video back to the initial style?")
                    edit_choice       = gr.inputs.Checkbox(default=False, label="With latent space editing?")
                    vid_format_choice = gr.inputs.Radio(choices=["gif", "mp4"], type="value", default='mp4', label="Video Format")
            
            with gr.Column():
                vid_output = gr.outputs.Video(label="Output Video")

    img_button.click(fn=editor.edit_image, inputs=[input_img, style_choice], outputs=img_output)
    vid_button.click(fn=editor.edit_video, inputs=[input_img, style_choice, edit_choice, vid_format_choice, loop_styles], outputs=vid_output)

    article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.00946' target='_blank'>StyleGAN-NADA: CLIP-Guided Domain Adaptation of Image Generators</a> | <a href='https://stylegan-nada.github.io/' target='_blank'>Project Page</a> | <a href='https://github.com/rinongal/StyleGAN-nada' target='_blank'>Code</a></p> <center><img src='https://visitor-badge.glitch.me/badge?page_id=rinong_sgnada' alt='visitor badge'></center>"
    gr.Markdown(article)

blocks.launch(enable_queue=True)