import os
import random
import gradio as gr


import cv2
import torch
import numpy as np
from PIL import Image

from transformers import CLIPVisionModelWithProjection
from diffusers.utils import load_image
from diffusers.models import ControlNetModel
# from diffusers.image_processor import IPAdapterMaskProcessor
from insightface.app import FaceAnalysis
# import sys
# import glob
# import os
import io
import spaces

from pipeline_stable_diffusion_xl_instantid import StableDiffusionXLInstantIDPipeline, draw_kps

import pandas as pd
import json
import requests
from PIL import Image
from io import BytesIO


def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
               pad_to_max_side=False, mode=Image.BILINEAR, base_pixel_number=64):
    
    w, h = input_image.size
    if size is not None:
        w_resize_new, h_resize_new = size
    else:
        ratio = min_side / min(h, w)
        w, h = round(ratio*w), round(ratio*h)
        ratio = max_side / max(h, w)
        input_image = input_image.resize([round(ratio*w), round(ratio*h)], mode)
        w_resize_new = (round(ratio * w) // base_pixel_number) * base_pixel_number
        h_resize_new = (round(ratio * h) // base_pixel_number) * base_pixel_number
    input_image = input_image.resize([w_resize_new, h_resize_new], mode)

    if pad_to_max_side:
        res = np.ones([max_side, max_side, 3], dtype=np.uint8) * 255
        offset_x = (max_side - w_resize_new) // 2
        offset_y = (max_side - h_resize_new) // 2
        res[offset_y:offset_y+h_resize_new, offset_x:offset_x+w_resize_new] = np.array(input_image)
        input_image = Image.fromarray(res)
    return input_image

def process_image_by_bbox_larger(input_image, bbox_xyxy, min_bbox_ratio=0.2):
    """
    Process an image based on a bounding box, cropping and resizing as necessary.

    Parameters:
    - input_image: PIL Image object.
    - bbox_xyxy: Tuple (x1, y1, x2, y2) representing the bounding box coordinates.

    Returns:
    - A processed image cropped and resized to 1024x1024 if the bounding box is valid,
      or None if the bounding box does not meet the required size criteria.
    """
    # Constants
    target_size = 1024
    # min_bbox_ratio = 0.2  # Bounding box should be at least 20% of the crop

    # Extract bounding box coordinates
    x1, y1, x2, y2 = bbox_xyxy
    bbox_w = x2 - x1
    bbox_h = y2 - y1

    # Calculate the area of the bounding box
    bbox_area = bbox_w * bbox_h

    # Start with the smallest square crop that allows bbox to be at least 20% of the crop area
    crop_size = max(bbox_w, bbox_h)
    initial_crop_area = crop_size * crop_size
    while (bbox_area / initial_crop_area) < min_bbox_ratio:
        crop_size += 10  # Gradually increase until bbox is at least 20% of the area
        initial_crop_area = crop_size * crop_size

    # Once the minimum condition is satisfied, try to expand the crop further
    max_possible_crop_size = min(input_image.width, input_image.height)
    while crop_size < max_possible_crop_size:
        # Calculate a potential new area
        new_crop_size = crop_size + 10
        new_crop_area = new_crop_size * new_crop_size
        if (bbox_area / new_crop_area) < min_bbox_ratio:
            break  # Stop if expanding further violates the 20% rule
        crop_size = new_crop_size

    # Determine the center of the bounding box
    center_x = (x1 + x2) // 2
    center_y = (y1 + y2) // 2

    # Calculate the crop coordinates centered around the bounding box
    crop_x1 = max(0, center_x - crop_size // 2)
    crop_y1 = max(0, center_y - crop_size // 2)
    crop_x2 = min(input_image.width, crop_x1 + crop_size)
    crop_y2 = min(input_image.height, crop_y1 + crop_size)

    # Ensure the crop is square, adjust if it goes out of image bounds
    if crop_x2 - crop_x1 != crop_y2 - crop_y1:
        side_length = min(crop_x2 - crop_x1, crop_y2 - crop_y1)
        crop_x2 = crop_x1 + side_length
        crop_y2 = crop_y1 + side_length

    # Crop the image
    cropped_image = input_image.crop((crop_x1, crop_y1, crop_x2, crop_y2))

    # Resize the cropped image to 1024x1024
    resized_image = cropped_image.resize((target_size, target_size), Image.LANCZOS)

    return resized_image

def calc_emb_cropped(image, app):
    face_image = image.copy()

    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))

    face_info = face_info[0]

    cropped_face_image = process_image_by_bbox_larger(face_image, face_info["bbox"], min_bbox_ratio=0.2)

    return cropped_face_image

def process_benchmark_csv(banchmark_csv_path):
    # Reading the first CSV file into a DataFrame
    df = pd.read_csv(banchmark_csv_path)

    # Drop any unnamed columns
    df = df.loc[:, ~df.columns.str.contains('^Unnamed')]

    # Drop columns with all NaN values
    df.dropna(axis=1, how='all', inplace=True)

    # Drop rows with all NaN values
    df.dropna(axis=0, how='all', inplace=True)

    df = df.loc[df['High resolution'] == 1]

    df.reset_index(drop=True, inplace=True)

    return df

def make_canny_condition(image, min_val=100, max_val=200, w_bilateral=True):
    if w_bilateral:
        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
        bilateral_filtered_image = cv2.bilateralFilter(image, d=9, sigmaColor=75, sigmaSpace=75)
        image = cv2.Canny(bilateral_filtered_image, min_val, max_val)
    else:
        image = np.array(image)
        image = cv2.Canny(image, min_val, max_val)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    image = Image.fromarray(image)
    return image


default_negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"

# Load face detection and recognition package
app = FaceAnalysis(name='antelopev2', root='./', providers=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))

base_dir = "./instantID_ckpt/checkpoint_174000"
face_adapter = f'{base_dir}/pytorch_model.bin'
controlnet_path = f'{base_dir}/controlnet'
base_model_path = f'briaai/BRIA-2.3'
resolution = 1024

controlnet_lnmks = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)

controlnet_canny = ControlNetModel.from_pretrained("briaai/BRIA-2.3-ControlNet-Canny",
                                                   torch_dtype=torch.float16)

controlnet = [controlnet_lnmks, controlnet_canny]

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

image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        '/home/ubuntu/BRIA-2.3-InstantID/ip_adapter/image_encoder',
        torch_dtype=torch.float16,
    )

pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
        base_model_path,
        controlnet=controlnet,
        torch_dtype=torch.float16,
        image_encoder=image_encoder # For compatibility issues - needs to be there
    )

pipe = pipe.to(device)

use_native_ip_adapter = True
pipe.use_native_ip_adapter=use_native_ip_adapter

pipe.load_ip_adapter_instantid(face_adapter)

clip_embeds=None


Loras_dict = {
    "":"",
    "Vangogh_Vanilla": "bold, dramatic brush strokes, vibrant colors, swirling patterns, intense, emotionally charged paintings of",
    "Avatar_internlm": "2d anime sketch avatar of",
    # "Tomer_Hanuka_V3": "Fluid lines",
    "Storyboards": "Illustration style for storyboarding",
    "3D_illustration": "3D object illustration, abstract",
    # "beetl_general_death_style_v2": "a pale, dead, unnatural color face with dark circles around the eyes",
    "Characters": "gaming vector Art"
    }

lora_names = Loras_dict.keys()

lora_base_path = "./LoRAs"

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


@spaces.GPU
def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale=0.8, kps_scale=0.6, canny_scale=0.4, lora_name="", lora_scale=0.7, progress=gr.Progress(track_tqdm=True)):
    if image_path is None:
        raise gr.Error(f"Cannot find any input face image! Please upload a face image.")
    
    # img = np.array(Image.open(image_path))[:,:,::-1]
    img = Image.open(image_path)

    face_image_orig = img #Image.open(BytesIO(response.content))
    face_image_cropped = calc_emb_cropped(face_image_orig, app)
    face_image = resize_img(face_image_cropped, max_side=resolution, min_side=resolution)
    # face_image_padded = resize_img(face_image_cropped, max_side=resolution, min_side=resolution, pad_to_max_side=True)
    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
    face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
    face_emb = face_info['embedding']
    face_kps = draw_kps(face_image, face_info['kps'])

    if canny_scale>0.0:
        # Convert PIL image to a file-like object
        image_file = io.BytesIO()
        face_image_cropped.save(image_file, format='JPEG')  # Save in the desired format (e.g., 'JPEG' or 'PNG')
        image_file.seek(0)  # Move to the start of the BytesIO stream

        url = "https://engine.prod.bria-api.com/v1/background/remove"

        payload = {}
        files = [
            ('file', ('image_name.jpeg', image_file, 'image/jpeg'))  # Specify file name, file-like object, and MIME type
        ]
        headers = {
        'api_token': 'a10d6386dd6a11ebba800242ac130004'
        }

        response = requests.request("POST", url, headers=headers, data=payload, files=files)

        print(response.text)

        response_json = json.loads(response.content.decode('utf-8'))

        img = requests.get(response_json['result_url'])

        processed_image = Image.open(io.BytesIO(img.content))

        # Assuming `processed_image` is the RGBA image returned
        if processed_image.mode == 'RGBA':
            # Create a white background image
            white_background = Image.new("RGB", processed_image.size, (255, 255, 255))
            # Composite the RGBA image over the white background
            face_image = Image.alpha_composite(white_background.convert('RGBA'), processed_image).convert('RGB')
        else:
            face_image = processed_image.convert('RGB')  # If already RGB, just ensure mode is correct

        canny_img = make_canny_condition(face_image, min_val=20, max_val=40, w_bilateral=True)
                    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    if lora_name != "":
        lora_path = os.path.join(lora_base_path, lora_name, "pytorch_lora_weights.safetensors")
        pipe.load_lora_weights(lora_path)
        pipe.fuse_lora(lora_scale)  
        pipe.enable_lora()

        lora_prefix = Loras_dict[lora_name]

        prompt = f"{lora_prefix} {prompt}"
        

    print("Start inference...")    
    images = pipe(
        prompt = prompt,
        negative_prompt = default_negative_prompt,
        image_embeds = face_emb,
        image = [face_kps, canny_img] if canny_scale>0.0 else face_kps,
        controlnet_conditioning_scale = [kps_scale, canny_scale] if canny_scale>0.0 else kps_scale,
        control_guidance_end = [1.0, 1.0] if canny_scale>0.0 else 1.0,
        ip_adapter_scale = ip_adapter_scale,
        num_inference_steps = num_steps,
        guidance_scale = guidance_scale,
        generator = generator,
        visual_prompt_embds = clip_embeds,
        cross_attention_kwargs = None,
        num_images_per_prompt=num_images,
    ).images #[0]

    if lora_name != "":
        pipe.disable_lora()
        pipe.unfuse_lora()
        pipe.unload_lora_weights()

    return images

### Description
title = r"""
<h1>Bria-2.3 ID preservation</h1>
"""

description = r"""
<b>🤗 Gradio demo</b> for bria ID preservation.<br>

Steps:<br>
1. Upload an image with a face. If multiple faces are detected, we use the largest one. For images with already tightly cropped faces, detection may fail, try images with a larger margin.
2. Click <b>Submit</b> to generate new images of the subject.
"""

Footer = r"""
Enjoy
"""

css = '''
.gradio-container {width: 85% !important}
'''
with gr.Blocks(css=css) as demo:

    # description
    gr.Markdown(title)
    gr.Markdown(description)

    with gr.Row():
        with gr.Column():
            
            # upload face image
            img_file = gr.Image(label="Upload a photo with a face", type="filepath")
            
            # Textbox for entering a prompt
            prompt = gr.Textbox(
            label="Prompt", 
            placeholder="Enter your prompt here", 
            info="Describe what you want to generate or modify in the image."
            )
            
            lora_name = gr.Dropdown(choices=lora_names, label="LoRA", value="",  info="Select a LoRA name from the list, not selecting any will disable LoRA.")

            submit = gr.Button("Submit", variant="primary")

            # use_lcm = gr.Checkbox(
            #     label="Use LCM-LoRA to accelerate sampling", value=False,
            #     info="Reduces sampling steps significantly, but may decrease quality.",
            # )
            
            with gr.Accordion(open=False, label="Advanced Options"):
                num_steps = gr.Slider( 
                    label="Number of sample steps",
                    minimum=1,
                    maximum=100,
                    step=1,
                    value=30,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.1,
                    maximum=10.0,
                    step=0.1,
                    value=5.0,
                )
                num_images = gr.Slider(
                    label="Number of output images",
                    minimum=1,
                    maximum=3,
                    step=1,
                    value=1,
                )
                ip_adapter_scale = gr.Slider(
                    label="ip adapter scale",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.01,
                    value=0.8,
                )
                kps_scale = gr.Slider(
                    label="kps control scale",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.01,
                    value=0.6,
                )
                canny_scale = gr.Slider(
                    label="canny control scale",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.01,
                    value=0.4,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=99999999,
                    step=1,
                    value=0,
                )
                seed = gr.Slider(
                    label="lora_scale",
                    minimum=0.0,
                    maximum=1.0,
                    step=0.01,
                    value=0.7,
                )
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        with gr.Column():
            gallery = gr.Gallery(label="Generated Images")

        submit.click(
            fn=randomize_seed_fn,
            inputs=[seed, randomize_seed],
            outputs=seed,
            queue=False,
            api_name=False,
        ).then(
            fn=generate_image,
            inputs=[img_file, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name],
            outputs=[gallery]
        )

    # use_lcm.input(
    #         fn=toggle_lcm_ui,
    #         inputs=[use_lcm],
    #         outputs=[num_steps, guidance_scale],
    #         queue=False,
    #     )       
    
    # gr.Examples(
    #     examples=get_example(),
    #     inputs=[img_file],
    #     run_on_click=True,
    #     fn=run_example,
    #     outputs=[gallery],
    # )
    
    gr.Markdown(Footer)

demo.launch()