import os
import cv2
import einops 
import gradio as gr
import numpy as np
import torch
import random
from huggingface_hub import hf_hub_download

from pytorch_lightning import seed_everything
from utils.resize import resize_image, HWC3
from cldm.model import create_model, load_state_dict
from cldm.ddim_lle import DDIMSampler as DDIMSampler_LLE
from cldm.ddim_hlg import DDIMSampler as DDIMSampler_HLG
from automation_pose_mask.openpose import OpenposeDetector
from automation_pose_mask.auto_mask import MaskDetector
from PIL import Image
from rembg import remove

from utils.config import (
    model_yaml,
    category_dict,
    attribute_dict
)

##########################################
# Download model files from HF Hub
##########################################

MODEL_REPO = "NguyenDinhHieu/EquiFashionModel"

openpose_body_model_path = hf_hub_download(MODEL_REPO, filename="body_pose_model.pth")
openpose_hand_model_path = hf_hub_download(MODEL_REPO, filename="hand_pose_model.pth")
sam_model_path = hf_hub_download(MODEL_REPO, filename="open_clip_pytorch_model.bin")
my_model_path = hf_hub_download(MODEL_REPO, filename="eqf_final.ckpt")

##########################################
# Initialize model on GPU once
##########################################

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

apply_openpose = OpenposeDetector(
    body_model_path=openpose_body_model_path,
    hand_model_path=openpose_hand_model_path
)

apply_mask = MaskDetector(sam_model_path=sam_model_path)

model = create_model(model_yaml).to(device)
model.load_state_dict(load_state_dict(my_model_path, location=device))
model.eval()

hlg_sampler = DDIMSampler_HLG(model)
lle_sampler = DDIMSampler_LLE(model)

##########################################
# Example images
##########################################

example_path = os.path.join(os.path.dirname(__file__), "preselected_images")
example_image_list = [os.path.join(example_path, x) for x in os.listdir(example_path)]

##########################################
# Utility functions
##########################################

def pil_to_binary_mask(pil_image, threshold=0):
    np_image = np.array(pil_image)
    grayscale_image = Image.fromarray(np_image).convert("L")
    binary_mask = (np.array(grayscale_image) > threshold).astype(np.uint8) * 255
    return Image.fromarray(binary_mask)

def add_white_background(image):
    image = image.convert("RGBA")
    white_bg = Image.new("RGBA", image.size, "WHITE")
    white_bg.paste(image, (0, 0), image)
    return white_bg.convert("RGB")

##########################################
# HLG PROCESS 
##########################################

def hlg_process(hlg_prompt, input_image, category, a_prompt, n_prompt,
                num_samples, image_resolution, detect_resolution, ddim_steps,
                guess_mode, strength, scale, seed, eta):

    with torch.no_grad():
        input_image = HWC3(input_image)
        detected_map, _ = apply_openpose(resize_image(input_image, detect_resolution))
        detected_map = HWC3(detected_map)
        img = resize_image(input_image, image_resolution)
        H, W, C = img.shape

        detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)

        control = torch.from_numpy(detected_map).float().to(device) / 255.0
        control = torch.stack([control for _ in range(num_samples)], dim=0)
        control = einops.rearrange(control, 'b h w c -> b c h w')

        if seed == -1:
            seed = random.randint(0, 4294967294)
        seed_everything(seed)

        cond = {
            "c_concat": [control],
            "c_crossattn": [model.get_learned_conditioning([hlg_prompt + ', ' + a_prompt] * num_samples)]
        }
        un_cond = {
            "c_concat": None if guess_mode else [control],
            "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]
        }
        shape = (4, H // 8, W // 8)

        model.control_scales = ([strength] * 13)

        samples, _ = hlg_sampler.sample(ddim_steps, num_samples, shape, cond,
                                        verbose=False, eta=eta,
                                        unconditional_guidance_scale=scale,
                                        unconditional_conditioning=un_cond)

        x_samples = model.decode_first_stage(samples)
        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c')
                     * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)

        results = [Image.fromarray(x_samples[i]) for i in range(num_samples)]
        return [add_white_background(remove(img)) for img in results]

##########################################
# LLE PROCESS 
##########################################

def lle_process(lle_prompt, dict_img_mask, category, a_prompt, n_prompt,
                num_samples, image_resolution, detect_resolution, ddim_steps,
                guess_mode, strength, scale, seed, eta, attribute, selection_mode):

    input_image = dict_img_mask["background"].convert("RGB")
    input_image = HWC3(np.array(input_image))

    detected_map, keypoints = apply_openpose(resize_image(input_image, detect_resolution))
    detected_map = HWC3(detected_map)

    if selection_mode == "Automatically recognize":
        mask = apply_mask(resize_image(input_image, detect_resolution), keypoints,
                          category=category, attribute=attribute, sam_mode=True)
    else:
        mask = pil_to_binary_mask(dict_img_mask['layers'][0].convert("RGB"))

    if mask is not None:
        mask = torch.from_numpy(np.array(mask.convert("L"))).float().to(device) / 255.0
        mask = mask.unsqueeze(0).unsqueeze(0)

    img = resize_image(input_image, image_resolution)
    H, W, C = img.shape

    init_img = torch.from_numpy(img).float().to(device) / 127.0 - 1.0
    init_img = einops.rearrange(init_img, 'h w c -> 1 c h w')
    init_img = torch.stack([init_img] * num_samples)

    detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST)
    control = torch.from_numpy(detected_map).float().to(device) / 255.0
    control = torch.stack([control]*num_samples)
    control = einops.rearrange(control, 'b h w c -> b c h w')

    if seed == -1:
        seed = random.randint(0, 4294967294)
    seed_everything(seed)

    cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning([lle_prompt + ', ' + a_prompt] * num_samples)]}
    un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
    shape = (4, H // 8, W // 8)

    samples, _ = lle_sampler.sample(ddim_steps, num_samples, shape, cond,
                                     verbose=False, eta=eta,
                                     unconditional_guidance_scale=scale,
                                     unconditional_conditioning=un_cond,
                                     init_img=init_img, mask=mask,
                                     english_attribute=attribute)

    x_samples = model.decode_first_stage(samples)
    x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c')
                 * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)

    return [Image.fromarray(x_samples[i]) for i in range(num_samples)]

##########################################
# Send result to attribute editor
##########################################

def result2input(images):
    return {"background": images[-1], "layers": None, "composite": None}

##########################################
# FULL UI
##########################################

def create_hfddm():
    with gr.Blocks().queue() as app:

        category = gr.Radio(list(category_dict.values()), value=list(category_dict.values())[0], label="Clothing Category")

        with gr.Row():
            with gr.Column():
                with gr.Tab("Draft Design"):
                    hlg_prompt = gr.Textbox(label="High-level design prompt")
                    hlg_input_image = gr.Image(sources=("upload", "webcam"), type="numpy", value=example_image_list[0], label="Reference pose")
                    gr.Examples(inputs=hlg_input_image, examples=example_image_list)
                    hlg_run = gr.Button("Generate")

                with gr.Tab("Attribute Editing"):
                    lle_prompt = gr.Textbox(label="Attribute prompt")
                    lle_input_image = gr.ImageEditor(sources='upload', type="pil", label="Edit regions", value=example_image_list[0])
                    gr.Examples(inputs=lle_input_image, examples=example_image_list)
                    selection_mode = gr.Radio(["Automatically recognize", "User interface"], label="Mask Selection", value="Automatically recognize")

                    current_tab = {}
                    lle_run = {}
                    for tab_elem in attribute_dict.values():
                        with gr.Tab(tab_elem):
                            current_tab[tab_elem] = gr.Label(value=tab_elem, visible=False)
                            lle_run[tab_elem] = gr.Button("Generate")

            with gr.Column():
                result_gallery = gr.Gallery(label="Result", show_label=False, elem_id="gallery", selected_index=0, interactive=False)
                send2llg = gr.Button("Send to Attribute Editing")

        with gr.Accordion("Advanced Options", open=False):
            num_samples = gr.Slider(label="Images", minimum=1, maximum=4, value=2, step=1)
            image_resolution = gr.Slider(label="Resolution", minimum=256, maximum=768, value=512, step=64)
            strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
            guess_mode = gr.Checkbox(label='Guess Mode', value=False)
            detect_resolution = gr.Slider(label="Pose Detection Resolution", minimum=128, maximum=1024, value=512, step=1)
            ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=10, step=1, visible=False)
            scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
            seed = gr.Slider(label="Seed", minimum=-1, maximum=4294967294, value=11, step=1)
            eta = gr.Number(label="ETA (DDIM)", value=0.0)
            a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed, masterpiece, 8k, white background')
            n_prompt = gr.Textbox(label="Negative Prompt", value='worst quality, low quality, bad anatomy, watermark, signature, blurry')

        hlg_run.click(fn=hlg_process, inputs=[hlg_prompt, hlg_input_image, category, a_prompt, n_prompt,
                                              num_samples, image_resolution, detect_resolution, ddim_steps,
                                              guess_mode, strength, scale, seed, eta], outputs=[result_gallery])

        for tab_elem in attribute_dict.values():
            lle_run[tab_elem].click(fn=lle_process, inputs=[lle_prompt, lle_input_image, category, a_prompt, n_prompt,
                                                            num_samples, image_resolution, detect_resolution,
                                                            ddim_steps, guess_mode, strength, scale, seed, eta,
                                                            current_tab[tab_elem], selection_mode], outputs=[result_gallery])

        send2llg.click(fn=result2input, inputs=result_gallery, outputs=lle_input_image)

    return app

hfddm_block = create_hfddm()

demo = gr.Blocks(title="AI Fashion Design", theme=gr.themes.Monochrome(secondary_hue="orange", neutral_hue="gray")).queue()

with demo:
    gr.Markdown("# **AI Fashion Design** 👗")
    with gr.Tab("Fashion Design"):
        hfddm_block.render()

demo.launch()