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")
    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()
    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 = '/globalscratch/mridul/ldm/butterflies/model_runs/2024-06-18T21-37-12_HLE_lr1e-6_custom_NEW/checkpoints/epoch=000233.ckpt'
config_path = '/globalscratch/mridul/ldm/butterflies/model_runs/2024-06-18T21-37-12_HLE_lr1e-6_custom_NEW/configs/2024-06-18T21-37-12-project.yaml'
config = OmegaConf.load(config_path)  # TODO: Optionally download from same location as ckpt and chnage this logic
model = load_model_from_config(config, ckpt_path)  # TODO: check path

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

class_to_node = '/projects/ml4science/mridul/data/cambridge_butterfly/level_encodings/butterflies_hle_4levels_custom_NEW.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()}
species_name_to_class = {'_'.join(x.split('_')[2:]):x for x in class_to_node_dict.keys()}

species_names = list(species_name_to_class.keys())

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


    if opt.plms:
        sampler = PLMSSampler(model)
    else:
        sampler = DDIMSampler(model)


    prompt = class_to_node_dict[species_name_to_class[fish_name]]

    ### Trait Swapping
    if swap_fish_name!='None':
        # swap_fish_name = swap_fish_name.lower()
        swap_level = int(swap_level_input.split(" ")[-1]) - 1
        swap_fish = class_to_node_dict[species_name_to_class[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))
    # final_image.save(os.path.join(sample_path, f'{class_name.replace(" ", "-")}.png'))

    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(
    #     "--outdir",
    #     type=str,
    #     nargs="?",
    #     help="dir to write results to",
    #     default="outputs/txt2img-samples"
    # )
    parser.add_argument(
        "--ddim_steps",
        type=int,
        default=200,
        help="number of ddim sampling steps",
    )

    parser.add_argument(
        "--plms",
        action='store_true',
        help="use plms sampling",
    )

    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(
    #     "--H",
    #     type=int,
    #     default=256,
    #     help="image height, in pixel space",
    # )

    # parser.add_argument(
    #     "--W",
    #     type=int,
    #     default=256,
    #     help="image width, in pixel space",
    # )

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

    # parser.add_argument(
    #     "--output_dir_name",
    #     type=str,
    #     default='default_file',
    #     help="name of folder",
    # )

    # parser.add_argument(
    #     "--postfix",
    #     type=str,
    #     default='',
    #     help="name of folder",
    # )

    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 Butterfly Images Tool"
    description = "Write the Species name to generate an image for.\n For Trait Masking: Specify the Level information as well"


    def load_example(prompt, level, option, components):
        components['prompt_input'].value = prompt
        components['masking_level_input'].value = level
        # components['option'].value = option

    def setup_interface():
        with gr.Blocks() as demo:

            gr.Markdown("# Phylo Diffusion - Generating Butterfly Images Tool")
            gr.Markdown("### Write the Species name to generate a butterfly image")
            gr.Markdown("### 1. Trait Masking: Specify the Level information as well")
            gr.Markdown("### 2. Trait Swapping: Specify the species name to swap trait with at also at what level")

            with gr.Row():
                with gr.Column():
                    gr.Markdown("## Generate Images Based on Prompts")
                    gr.Markdown("Select a species to generate an image:")
                    # prompt_input = gr.Textbox(label="Species Name")
                    prompt_input = gr.Dropdown(label="Select Butterfly", choices=species_names, value="None")
                    gr.Markdown("Trait Masking")
                    with gr.Row():
                        masking_level_input = gr.Dropdown(label="Select Ancestral Level", choices=["None", "Level 3", "Level 2"], value="None")
                        # masking_node_input = gr.Dropdown(label="Select Internal", choices=["0", "1", "2", "3", "4", "5", "6", "7", "8"], value="0")

                    gr.Markdown("Trait Swapping")
                    with gr.Row():
                        swap_fish_name = gr.Dropdown(label="Select species Name to swap trait with:", choices=species_names, value="None")
                        swap_level_input = gr.Dropdown(label="Level of swapping", choices=["Level 3", "Level 2"], value="Level 3")
                    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=768, height=256)


                    # # Place to put example buttons
                    # gr.Markdown("## Select an example:")
                    # examples = [
                    # ("Gambusia Affinis", "None", "", "Level 3"), 
                    # ("Lepomis Auritus", "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}")
                    #     elif level != "None":
                    #         button = gr.Button(f"Species: {text} | Masking: {level}")
                    #     elif swap_text != "":
                    #         button = gr.Button(f"Species: {text} | Swapping with {swap_text} at {swap_level} ")
                    #     button.click(
                    #         fn=lambda text=text, level=level, swap_text=swap_text, swap_level=swap_level: (text, level, swap_text, swap_level),
                    #         inputs=[],
                    #         outputs=[prompt_input, masking_level_input, swap_fish_name, swap_level_input]
                    #     )
                

            # Display an image of the architecture


            submit_button.click(
                fn=generate_image,
                inputs=[prompt_input, masking_level_input,
                        swap_fish_name, swap_level_input],
                outputs=output_image
            )

        return demo

    # # Launch the interface
    # iface = setup_interface()
    
        # iface = gr.Interface(
    # fn=generate_image,
    # inputs=gr.Textbox(label="Prompt"),
    # outputs=[
    #     gr.Image(label="Generated Image"),
    # ] 
    # )

    iface = setup_interface()

    iface.launch(share=True)