import cProfile
import pstats
import io
import gc
import random
import time
import gradio as gr
import spaces
import imageio
from huggingface_hub import HfApi
import torch
from PIL import Image
from diffusers import (
    ControlNetModel,
    DPMSolverMultistepScheduler,
    StableDiffusionControlNetPipeline,
)
from preprocess_anime import Preprocessor, conditionally_manage_memory
from settings import API_KEY, MAX_NUM_IMAGES, MAX_SEED

preprocessor = None
controlnet = None
scheduler = None
pipe = None
compiled = False
api = HfApi()

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

def get_additional_prompt():
    prompt = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
    top = ["tank top", "blouse", "button up shirt", "sweater", "corset top"]
    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt", "leggings", "high-waisted shorts"]
    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker", "necklace", "headband", "headphones"]
    return f"{prompt}, {random.choice(top)}, {random.choice(bottom)}, {random.choice(accessory)}, score_9"
    # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]

def get_prompt(prompt, additional_prompt):
    default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
    randomize = get_additional_prompt()
    nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
    bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
    ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
    abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
    shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari"
        
    if prompt == "":
        prompts = [randomize, nude, bodypaint, pet_play, bondage, ahegao2, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari]
        prompts_nsfw = [nude, bodypaint, abg, ahegao2, shibari]
        preset = random.choice(prompts)
        prompt = f"{preset}"
        # print(f"-------------{preset}-------------")
    else:
        # prompt = f"{prompt}, {randomize}"
        prompt = f"{default},{prompt}"
    print(f"{prompt}")
    return prompt

def initialize_models():
    global preprocessor, controlnet, scheduler, pipe
    if preprocessor is None:
        preprocessor = Preprocessor()

    if controlnet is None:
        model_id = "lllyasviel/control_v11p_sd15_normalbae"
        print("initializing controlnet")
        controlnet = ControlNetModel.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            attn_implementation="flash_attention_2",
        ).to("cuda")

    if scheduler is None:
        scheduler = DPMSolverMultistepScheduler.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            solver_order=2,
            subfolder="scheduler",
            use_karras_sigmas=True,
            final_sigmas_type="sigma_min",
            algorithm_type="sde-dpmsolver++",
            prediction_type="epsilon",
            thresholding=False,
            denoise_final=True,
            device_map="cuda",
        )

    if pipe is None:
        base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
        pipe = StableDiffusionControlNetPipeline.from_single_file(
            base_model_url,
            safety_checker=None,
            controlnet=controlnet,
            scheduler=scheduler,
            torch_dtype=torch.float16,
        )
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
        pipe.to("cuda")
        print("---------------Loaded controlnet pipeline---------------")

@spaces.GPU(duration=11)
@torch.inference_mode()
def process_image(
    image,
    prompt,
    a_prompt,
    n_prompt,
    num_images,
    image_resolution,
    preprocess_resolution,
    num_steps,
    guidance_scale,
    seed,
):
    initialize_models()
    preprocessor.load("NormalBae")
    control_image = preprocessor(
        image=image,
        image_resolution=image_resolution,
        detect_resolution=preprocess_resolution,
    )
    custom_prompt = str(get_prompt(prompt, a_prompt))
    negative_prompt = str(n_prompt)
    global compiled
    generator = torch.cuda.manual_seed(seed)
    if not compiled:
        print("-----------------------------------Not Compiled-----------------------------------")
        compiled = True
    start = time.time()
    results = pipe(
        prompt=custom_prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_images_per_prompt=num_images,
        num_inference_steps=num_steps,
        generator=generator,
        image=control_image,
    ).images[0]
    print(f"Inference done in: {time.time() - start:.2f} seconds")
    
    timestamp = int(time.time())
    img_path = f"{timestamp}.jpg"
    results_path = f"{timestamp}_out.jpg"
    imageio.imsave(img_path, image)
    results.save(results_path)
    
    api.upload_file(
        path_or_fileobj=img_path,
        path_in_repo=img_path,
        repo_id="broyang/anime-ai-outputs",
        repo_type="dataset",
        token=API_KEY,
        run_as_future=True,
    )
    api.upload_file(
        path_or_fileobj=results_path,
        path_in_repo=results_path,
        repo_id="broyang/anime-ai-outputs",
        repo_type="dataset",
        token=API_KEY,
        run_as_future=True,
    )

    conditionally_manage_memory()

    results.save("temp_image.png")
    return results

def main():
    prod = True
    show_options = True
    if prod:
        show_options = False

    print("CUDA version:", torch.version.cuda)
    print("loading pipe")

    css = """
    h1 {
        text-align: center;
        display:block;
    }
    h2 {
        text-align: center;
        display:block;
    }
    h3 {
        text-align: center;
        display:block;
    }
    footer {visibility: hidden}
    """
    with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
        with gr.Row():
            with gr.Accordion("Advanced options", open=show_options, visible=show_options):
                num_images = gr.Slider(
                    label="Images", minimum=1, maximum=MAX_NUM_IMAGES, value=1, step=1
                )
                image_resolution = gr.Slider(
                    label="Image resolution",
                    minimum=256,
                    maximum=1024,
                    value=768,
                    step=256,
                )
                preprocess_resolution = gr.Slider(
                    label="Preprocess resolution",
                    minimum=128,
                    maximum=1024,
                    value=768,
                    step=1,
                )
                num_steps = gr.Slider(
                    label="Number of steps", minimum=1, maximum=100, value=12, step=1
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
                )
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                a_prompt = gr.Textbox(
                    label="Additional prompt",
                    value = ""
                )
                n_prompt = gr.Textbox(
                    label="Negative prompt",
                    value="EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
                )
        with gr.Column():
            prompt = gr.Textbox(
                label="Description",
                placeholder="Leave empty for something spicy 👀",
            )
        with gr.Row():
            with gr.Column():
                image = gr.Image(
                    label="Input",
                    sources=["upload"],
                    show_label=True,
                    format="webp",
                )
                with gr.Column():
                    run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
            with gr.Column():
                result = gr.Image(
                    label="Anime AI",
                    interactive=False,
                    format="webp",
                    visible = True,
                    show_share_button= False,
                )
                with gr.Column():
                    use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
        config = [
            image,
            prompt,
            a_prompt,
            n_prompt,
            num_images,
            image_resolution,
            preprocess_resolution,
            num_steps,
            guidance_scale,
            seed,
        ]
        
        @spaces.GPU(duration=11)
        @torch.inference_mode()
        @gr.on(triggers=[image.upload], inputs=config, outputs=[result])
        def auto_process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed):
            return process_image(image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
    
        @gr.on(triggers=[image.upload], inputs=None, outputs=[use_ai_button, run_button])
        def turn_buttons_off():
            return gr.update(visible=False), gr.update(visible=False)
    
        @gr.on(triggers=[use_ai_button.click], inputs=None, outputs=[use_ai_button, run_button])
        def turn_buttons_off():
            return gr.update(visible=False), gr.update(visible=False)
    
        @gr.on(triggers=[run_button.click], inputs=None, outputs=[use_ai_button, run_button])
        def turn_buttons_off():
            return gr.update(visible=False), gr.update(visible=False)
    
        @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button])
        def turn_buttons_on():
            return gr.update(visible=True), gr.update(visible=True)
        
        with gr.Row():
            helper_text = gr.Markdown("## Tap and hold (on mobile) to save the image.", visible=True)
            
        prompt.submit(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
            show_progress="none",
        ).then(
            fn=auto_process_image,
            inputs=config,
            outputs=result,
            api_name=False,
            show_progress="minimal",
        )

        run_button.click(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
            show_progress="none",
        ).then(
            fn=auto_process_image,
            inputs=config,
            outputs=result,
            show_progress="minimal",
        )

        def update_config():
            try:
                print("Updating image to AI Temp Image")
                ai_temp_image = Image.open("temp_image.png")
                return ai_temp_image
            except FileNotFoundError:
                print("No AI Image Available")
                return None
        
        use_ai_button.click(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
            show_progress="none",
        ).then(
            fn=lambda _: update_config(),
            inputs=[image],
            outputs=image,
            show_progress="minimal",
        ).then(
            fn=auto_process_image,
            inputs=[image, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed],
            outputs=result,
            show_progress="minimal",
        )

    demo.launch()

if __name__ == "__main__":
    pr = cProfile.Profile()
    pr.enable()
    main()
    pr.disable()

    s = io.StringIO()
    sortby = 'cumulative'
    ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
    ps.print_stats()
    print(s.getvalue())