import torch
import gradio as gr


import argparse, os, sys, glob
import torch
import pickle
import numpy as np
from omegaconf import OmegaConf
from PIL import Image
from tqdm import tqdm, trange
from einops import rearrange
from torchvision.utils import make_grid

from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler


def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu") # TODO: change for GPU resources
    # pl_sd = torch.load(ckpt)#, map_location="cpu")
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    if len(m) > 0 and verbose:
        print("missing keys:")
        print(m)
    if len(u) > 0 and verbose:
        print("unexpected keys:")
        print(u)

    # model.cuda() # TODO: change for GPU resources
    model.eval()
    return model


def masking_embed(embedding, levels=1):
    """
    size of embedding - nx1xd, n: number of samples, d - 512
    replacing the last 128*levels from the embedding
    """
    replace_size = 128*levels
    random_noise = torch.randn(embedding.shape[0], embedding.shape[1], replace_size)
    embedding[:, :, -replace_size:] = random_noise
    return embedding


# LOAD MODEL GLOBALLY
ckpt_path = './model_files/fishes/epoch=000119.ckpt'
config_path = './model_files/fishes/2024-03-01T23-15-36-project.yaml'
config = OmegaConf.load(config_path)
model = load_model_from_config(config, ckpt_path) 

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = model.to(device)

def generate_image(fish_name, masking_level_input,
                        swap_fish_name, swap_level_input):
    
    fish_name = fish_name.lower()


    label_to_class_mapping = {0: 'Alosa-chrysochloris', 1: 'Carassius-auratus', 2: 'Cyprinus-carpio', 3: 'Esox-americanus', 
    4: 'Gambusia-affinis', 5: 'Lepisosteus-osseus', 6: 'Lepisosteus-platostomus', 7: 'Lepomis-auritus', 8: 'Lepomis-cyanellus', 
    9: 'Lepomis-gibbosus', 10: 'Lepomis-gulosus', 11: 'Lepomis-humilis', 12: 'Lepomis-macrochirus', 13: 'Lepomis-megalotis', 
    14: 'Lepomis-microlophus', 15: 'Morone-chrysops', 16: 'Morone-mississippiensis', 17: 'Notropis-atherinoides', 
    18: 'Notropis-blennius', 19: 'Notropis-boops', 20: 'Notropis-buccatus', 21: 'Notropis-buchanani', 22: 'Notropis-dorsalis', 
    23: 'Notropis-hudsonius', 24: 'Notropis-leuciodus', 25: 'Notropis-nubilus', 26: 'Notropis-percobromus', 
    27: 'Notropis-stramineus', 28: 'Notropis-telescopus', 29: 'Notropis-texanus', 30: 'Notropis-volucellus', 
    31: 'Notropis-wickliffi', 32: 'Noturus-exilis', 33: 'Noturus-flavus', 34: 'Noturus-gyrinus', 35: 'Noturus-miurus', 
    36: 'Noturus-nocturnus', 37: 'Phenacobius-mirabilis'}

    def get_label_from_class(class_name):
        for key, value in label_to_class_mapping.items():
            if value == class_name:
                return key


    sampler = DDIMSampler(model)

    prompt = opt.prompt
    all_images = []
    labels = []

    class_to_node = './model_files/fishes/class_to_ancestral_label.pkl'
    with open(class_to_node, 'rb') as pickle_file:
        class_to_node_dict = pickle.load(pickle_file)

    class_to_node_dict =  {key.lower(): value for key, value in class_to_node_dict.items()}
    

    prompt = class_to_node_dict[fish_name]

    ### Trait Swapping
    if swap_fish_name:
        swap_fish_name = swap_fish_name.lower()
        swap_level = int(swap_level_input.split(" ")[-1]) - 1
        swap_fish = class_to_node_dict[swap_fish_name]

        swap_fish_split = swap_fish[0].split(',')
        fish_name_split = prompt[0].split(',')
        fish_name_split[swap_level] = swap_fish_split[swap_level]

        prompt = [','.join(fish_name_split)]

    all_samples=list()
    with torch.no_grad():
        with model.ema_scope():
            uc = None
            for n in trange(opt.n_iter, desc="Sampling"):

                all_prompts = opt.n_samples * (prompt)
                all_prompts = [tuple(all_prompts)]
                c = model.get_learned_conditioning({'class_to_node': all_prompts})
                if masking_level_input != "None":
                    masked_level = int(masking_level_input.split(" ")[-1])
                    masked_level = 4-masked_level
                    c = masking_embed(c, levels=masked_level)
                shape = [3, 64, 64]
                samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                conditioning=c,
                                                batch_size=opt.n_samples,
                                                shape=shape,
                                                verbose=False,
                                                unconditional_guidance_scale=opt.scale,
                                                unconditional_conditioning=uc,
                                                eta=opt.ddim_eta)

                x_samples_ddim = model.decode_first_stage(samples_ddim)
                x_samples_ddim = torch.clamp((x_samples_ddim+1.0)/2.0, min=0.0, max=1.0)

                all_samples.append(x_samples_ddim)

    ###### to make grid
    # additionally, save as grid
    grid = torch.stack(all_samples, 0)
    grid = rearrange(grid, 'n b c h w -> (n b) c h w')
    grid = make_grid(grid, nrow=opt.n_samples)

    # to image
    grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
    final_image = Image.fromarray(grid.astype(np.uint8))

    return final_image


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--prompt",
        type=str,
        nargs="?",
        default="a painting of a virus monster playing guitar",
        help="the prompt to render"
    )

    parser.add_argument(
        "--ddim_steps",
        type=int,
        default=200,
        help="number of ddim sampling steps",
    )

    parser.add_argument(
        "--ddim_eta",
        type=float,
        default=1.0,
        help="ddim eta (eta=0.0 corresponds to deterministic sampling",
    )
    parser.add_argument(
        "--n_iter",
        type=int,
        default=1,
        help="sample this often",
    )

    parser.add_argument(
        "--n_samples",
        type=int,
        default=1,
        help="how many samples to produce for the given prompt",
    )

    parser.add_argument(
        "--scale",
        type=float,
        # default=5.0,
        default=1.0,
        help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
    )
    opt = parser.parse_args()

    title = "🎞️ Phylo Diffusion - Generating Fish Images Tool"

    def load_example(prompt, level, option, components):
        components['prompt_input'].value = prompt
        components['masking_level_input'].value = level
    
    def setup_interface():
        with gr.Blocks() as demo:

            gr.Markdown("# Phylo-Diffusion: Generating Fish Images Tool")
            gr.Markdown("### Write the Species name to generate a fish image")
            gr.Markdown("### 1. Trait Masking: Specify the Level information to mask")
            gr.Markdown("### 2. Trait Swapping: Specify the species name to swap trait with and at what level")
    
            with gr.Row():
                with gr.Column():
                    # gr.Markdown("## Generate Images Based on Prompts")
                    gr.Markdown("**NOTE:** The demo is currently running on free CPU resources provided by Hugging Face, so it may take up to 10 minutes to generate an image. We're working on securing additional resources to speed up the process. Thank you for your patience!")
                    prompt_input = gr.Textbox(label="Species Name")
    
                    # Radio button to select experiment type, with no default selection
                    experiment_choice = gr.Radio(label="Select Experiment", choices=["Trait Masking", "Trait Swapping"], value=None)
    
                    # Trait Masking Inputs (hidden initially)
                    masking_level_input = gr.Dropdown(label="Select Ancestral Level", choices=["None", "Level 3", "Level 2"], value="None", visible=False)
    
                    # Trait Swapping Inputs (hidden initially)
                    swap_fish_name = gr.Textbox(label="Species Name to swap trait with:", visible=False)
                    swap_level_input = gr.Dropdown(label="Level of swapping", choices=["Level 3", "Level 2"], value="Level 3", visible=False)
    
                    submit_button = gr.Button("Generate")
    
                    gr.Markdown("## Phylogeny Tree")
                    architecture_image = "phylogeny_tree.jpg"  # Update this with the actual path
                    gr.Image(value=architecture_image, label="Phylogeny Tree")
    
                with gr.Column():
                    gr.Markdown("## Generated Image")
                    output_image = gr.Image(label="Generated Image", width=256, height=256)
    
                    # Place to put example buttons
                    gr.Markdown("## Select an example:")
                    examples = [
                        ("Gambusia Affinis", "None", "", "Level 3"),
                        ("Lepomis Auritus", "Level 3", "", "Level 3"),
                        ("Noturus nocturnus", "None", "Notropis dorsalis", "Level 2")
                    ]
    
                    for text, level, swap_text, swap_level in examples:
                        if level == "None" and swap_text == "":
                            button = gr.Button(f"Species: {text}")
                            experiment_type = "None"
                        elif level != "None":
                            button = gr.Button(f"Species: {text} | Masking: {level}")
                            experiment_type = "Trait Masking"
                        elif swap_text != "":
                            button = gr.Button(f"Species: {text} | Swapping with {swap_text} at {swap_level} ")
                            experiment_type = "Trait Swapping"
    
                        # Update radio button, fields and auto-trigger the "Generate" action
                        button.click(
                            fn=lambda text=text, level=level, swap_text=swap_text, swap_level=swap_level, experiment_type=experiment_type: (
                                text, level, swap_text, swap_level, experiment_type
                            ),
                            inputs=[],
                            outputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input, experiment_choice]
                        ).then(
                            fn=generate_image, 
                            inputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input], 
                            outputs=output_image
                        )
    
            # Update visibility of inputs based on experiment selection
            def update_inputs(experiment_type):
                if experiment_type == "Trait Masking":
                    return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
                elif experiment_type == "Trait Swapping":
                    return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
                else:
                    # No experiment selected
                    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    
            experiment_choice.change(
                fn=update_inputs,
                inputs=[experiment_choice],
                outputs=[masking_level_input, swap_fish_name, swap_level_input]
            )
    
            # Submit button functionality
            submit_button.click(
                fn=generate_image,
                inputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input],
                outputs=output_image
            )
    
        return demo
        

    iface = setup_interface()
    iface.launch(share=True)