import argparse
import cv2
import cv2
import torch
import gradio as gr
from transformers import SamModel, SamProcessor

import spaces
import numpy as np
from PIL import Image
from tqdm import tqdm
from torchvision.transforms import v2

from rynnec import disable_torch_init, model_init, mm_infer, mm_infer_segmentation
from rynnec.mm_utils import annToMask, load_video, load_images

from PIL import Image
from tqdm import tqdm
import numpy as np
import colorsys
import argparse


def get_hsv_palette(n_colors):
    hues = np.linspace(0, 1, int(n_colors) + 1)[1:-1]
    s = 0.8
    v = 0.9
    palette = [(0.0, 0.0, 0.0)] + [
        colorsys.hsv_to_rgb(h_i, s, v) for h_i in hues
    ]
    return (255 * np.asarray(palette)).astype("uint8")


def colorize_masks(images, index_masks, fac: float = 0.8, draw_contour=True, edge_thickness=20):
    max_idx = max([m.max() for m in index_masks])
    palette = get_hsv_palette(max_idx + 1)
    color_masks = []
    out_frames = []
    for img, mask in tqdm(zip(images, index_masks), desc='Visualize masks ...'):
        clr_mask = palette[mask.astype("int")]
        blended_img = img

        blended_img = compose_img_mask(blended_img, clr_mask, fac)

        if draw_contour:
            blended_img = draw_contours_on_image(blended_img, mask, clr_mask,
                                                 brightness_factor=1.8,
                                                 alpha=0.6,
                                                 thickness=edge_thickness)
        out_frames.append(blended_img)

    return out_frames, color_masks


def compose_img_mask(img, color_mask, fac: float = 0.5):
    mask_region = (color_mask.sum(axis=-1) > 0)[..., None]
    out_f = img.copy() / 255
    out_f[mask_region[:, :, 0]] = fac * img[mask_region[:, :, 0]] / 255 + (1 - fac) * color_mask[mask_region[:, :, 0]] / 255
    out_u = (255 * out_f).astype("uint8")
    return out_u


def draw_contours_on_image(img, index_mask, color_mask, brightness_factor=1.6, alpha=0.5, thickness=2, ignore_index=0):
    img = img.astype("float32")
    overlay = img.copy()

    unique_indices = np.unique(index_mask)
    if ignore_index is not None:
        unique_indices = [idx for idx in unique_indices if idx != ignore_index]

    for i in unique_indices:
        bin_mask = (index_mask == i).astype("uint8") * 255
        if bin_mask.sum() == 0:
            continue

        contours, _ = cv2.findContours(bin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        color = color_mask[index_mask == i][0].astype("float32")
        bright_color = np.clip(color * brightness_factor, 0, 255).tolist()

        cv2.drawContours(overlay, contours, -1, bright_color, thickness)

    blended = (1 - alpha) * img + alpha * overlay
    return np.clip(blended, 0, 255).astype("uint8")


def extract_first_frame_from_video(video):
    cap = cv2.VideoCapture(video)
    success, frame = cap.read()
    cap.release()
    if success:
        return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    return None


def extract_points_from_mask(mask_pil):
    mask = np.asarray(mask_pil)[..., 0]
    coords = np.nonzero(mask)
    coords = np.stack((coords[1], coords[0]), axis=1)

    return coords

def add_contour(img, mask, color=(1., 1., 1.)):
    img = img.copy()

    mask = mask.astype(np.uint8) * 255
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(img, contours, -1, color, thickness=8)

    return img


def load_first_frame(video_path):
    cap = cv2.VideoCapture(video_path)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        raise gr.Error("Could not read the video file.")
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image = Image.fromarray(frame)  
    return image


def clear_masks():
    return [], [], [], []

def clear_all():
    return [], [], [], [], None, "", ""


@spaces.GPU(duration=120)
def apply_sam(image, input_points):
    inputs = sam_processor(image, input_points=input_points, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = sam_model(**inputs)

    masks = sam_processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())[0][0]
    scores = outputs.iou_scores[0, 0]

    mask_selection_index = scores.argmax()

    mask_np = masks[mask_selection_index].numpy()

    return mask_np


@spaces.GPU(duration=120)
def run(mode, images, timestamps, masks, mask_ids, instruction, mask_output_video):
    if mode == "QA":
        response = run_text_inference(images, timestamps, masks, mask_ids, instruction)
    else:
        response, mask_output_video = run_seg_inference(images, timestamps, instruction)
    return response, mask_output_video


def run_text_inference(images, timestamps, masks, mask_ids, instruction):
    masks = torch.from_numpy(np.stack(masks, axis=0))

    if "<video>" not in instruction:
        instruction = "<video>\n" + instruction

    if len(masks) >= 2:
        obj_str = f"<video>\nThere are {len(masks)} objects in the video: " + ", ".join([f"<object{i}> [<REGION>]" for i in range(len(masks))])
        instruction = instruction.replace("<video>\n", obj_str)
    else:
        instruction = instruction.replace("<object0>", '[<REGION>]')

    output = mm_infer(
        (images, timestamps),
        processor,
        instruction,
        model=model,
        tokenizer=processor.tokenizer,
        do_sample=False,
        modal='video',
        masks=masks.cuda() if masks is not None else None,
        mask_ids=mask_ids
    )

    return output


def run_seg_inference(images, timestamps, instruction):
    output, masks = mm_infer_segmentation(
        (images, timestamps),
        seg_processor,
        instruction,
        model=seg_model,
        tokenizer=processor.tokenizer,
        do_sample=False,
        modal='video',
    )

    w, h = images[0].size
    masks = v2.Resize([h, w])(masks).cpu().numpy()

    mask_list_video = []

    images = [np.array(image) for image in images]
    masks = [mask[0] for mask in masks]
    show_images, _ = colorize_masks(images, masks)
    for i, image in enumerate(show_images):
        if masks[i].sum() > 1000:
            mask_list_video.append((Image.fromarray(image), f"Frame {i}"))        

    return output, mask_list_video


def generate_masks_video(image, mask_list_video, mask_raw_list_video, mask_ids, frame_idx):
    image['image'] = image['background'].convert('RGB')
    # del image['background'], image['composite']
    assert len(image['layers']) == 1, f"Expected 1 layer, got {len(image['layers'])}"

    mask = Image.fromarray((np.asarray(image['layers'][0])[..., 3] > 0).astype(np.uint8) * 255).convert('RGB')
    points = extract_points_from_mask(mask)
    np.random.seed(0)
    if points.shape[0] == 0:
        raise gr.Error("No points selected")

    points_selected_indices = np.random.choice(points.shape[0], size=min(points.shape[0], 8), replace=False)
    points = points[points_selected_indices]
    coords = [points.tolist()]
    mask_np = apply_sam(image['image'], coords)

    mask_raw_list_video.append(mask_np)
    mask_image = Image.fromarray((mask_np[:,:,np.newaxis] * np.array(image['image'])).astype(np.uint8))
    
    mask_list_video.append((mask_image, f"<object{len(mask_list_video)}>"))
    # Return a list containing the mask image.
    image['layers'] = []
    image['composite'] = image['background']
    mask_ids.append(frame_idx)
    return mask_list_video, image, mask_list_video, mask_raw_list_video, mask_ids


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VideoRefer gradio demo")
    parser.add_argument("--model-path", type=str, default="Alibaba-DAMO-Academy/RynnEC-2B", help="Path to the model checkpoint")
    parser.add_argument("--seg-model-path", type=str, default="Alibaba-DAMO-Academy/RynnEC-2B-Cog", help="Path to the model checkpoint")
    parser.add_argument("--port", type=int, default=4001)

    args_cli = parser.parse_args()

    with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber")) as demo:

        mask_list = gr.State([])  
        mask_raw_list = gr.State([])  
        mask_list_video = gr.State([])  
        mask_raw_list_video = gr.State([])  


        HEADER = ("""
            <div>
                <h1>RynnEC Demo</h1>
                <h5 style="margin: 0;">Feel free to click on anything that grabs your interest!</h5>
                <h5 style="margin: 0;">If this demo please you, please give us a star ⭐ on Github or 💖 on this space.</h5>
            </div>
            </div>
            <div style="display: flex; justify-content: left; margin-top: 10px;">
            <a href="https://arxiv.org/pdf/2501.00599"><img src="https://img.shields.io/badge/Arxiv-2501.00599-ECA8A7" style="margin-right: 5px;"></a>
            <a href="https://github.com/DAMO-NLP-SG/VideoRefer"><img src='https://img.shields.io/badge/Github-VideoRefer-F7C97E' style="margin-right: 5px;"></a>
            <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3"><img src='https://img.shields.io/badge/Github-VideoLLaMA3-9DC3E6' style="margin-right: 5px;"></a>
            </div>
            """)


        image_tips = """
                ### 💡 Tips:

                🧸 Upload an image, and you can use the drawing tool✍️ to highlight the areas you're interested in.
            
                🔖 For single-object caption mode, simply select the area and click the 'Generate Caption' button to receive a caption for the object.
                
                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
                
                📌 Click the button 'Clear Masks' to clear the current generated masks.
                
                """
        
        video_tips = """
                ### 💡 Tips:
                🧸 Upload an video, and you can use the drawing tool✍️ to highlight the areas you're interested in the first frame.
                
                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
                
                📌 Click the button 'Clear Masks' to clear the current generated masks.
                
                """


        with gr.TabItem("Video"):
            with gr.Row():
                with gr.Column():
                    video_input = gr.Video(label="Video", interactive=True)
                    frame_idx = gr.Slider(minimum=0, maximum=0, value=0, step=1, label="Select Frame", interactive=False)
                    selected_frame = gr.ImageEditor(
                        label="Annotate Frame",
                        type="pil", 
                        sources=[], 
                        interactive=True,
                    )
                    generate_mask_btn_video = gr.Button("1️⃣ Generate Mask", visible=True, variant="primary")
                    gr.Examples([f"./demo/videos/{i+1}.mp4" for i in range(4)], inputs=video_input, label="Examples")

                with gr.Column():
                    mode_video = gr.Radio(label="Mode", choices=["QA", "Seg"], value="QA")
                    mask_output_video = gr.Gallery(label="Referred Masks", object_fit='scale-down')

                    query_video = gr.Textbox(label="Question", value="Please describe <object0>.", interactive=True, visible=True)
                    response_video = gr.Textbox(label="Answer", interactive=False)

                    submit_btn_video = gr.Button("Generate Caption", variant="primary", visible=False)
                    submit_btn_video1 = gr.Button("2️⃣ Generate Answer", variant="primary", visible=True)
                    description_video = gr.Textbox(label="Output", visible=False)
                    
                    clear_masks_btn_video = gr.Button("Clear Masks", variant="secondary")

            gr.Markdown(video_tips)

            frames = gr.State(value=[])
            timestamps = gr.State(value=[])
            mask_ids = gr.State(value=[])

        def on_video_upload(video_path):
            frames, timestamps = load_video(video_path, fps=1, max_frames=128)
            frames = [Image.fromarray(x.transpose(1, 2, 0)) for x in frames]
            return frames, timestamps, frames[0], gr.update(value=0, maximum=len(frames) - 1, interactive=True)

        def on_frame_idx_change(frame_idx, frames):
            return frames[frame_idx]

        def to_seg_mode():
            return (
                *[gr.update(visible=False) for _ in range(4)],
                []
            )

        def to_qa_mode():
            return (
                *[gr.update(visible=True) for _ in range(4)],
                []
            )

        def on_mode_change(mode):
            if mode == "QA":
                return to_qa_mode()
            return to_seg_mode()

        mode_video.change(on_mode_change, inputs=[mode_video], outputs=[frame_idx, selected_frame, generate_mask_btn_video, response_video, mask_output_video])
        video_input.change(on_video_upload, inputs=[video_input], outputs=[frames, timestamps, selected_frame, frame_idx])
        frame_idx.change(on_frame_idx_change, inputs=[frame_idx, frames], outputs=[selected_frame])

        generate_mask_btn_video.click(
            fn=generate_masks_video,
            inputs=[selected_frame, mask_list_video, mask_raw_list_video, mask_ids, frame_idx],
            outputs=[mask_output_video, selected_frame, mask_list_video, mask_raw_list_video, mask_ids]
        )

        submit_btn_video1.click(
            fn=run,
            inputs=[mode_video, frames, timestamps, mask_raw_list_video, mask_ids, query_video, mask_output_video],
            outputs=[response_video, mask_output_video],
            api_name="describe_video"
        )

        video_input.clear(
            fn=clear_all,
            outputs=[mask_output_video, mask_list_video, mask_raw_list_video, mask_ids, selected_frame, query_video, response_video]
        )

        clear_masks_btn_video.click(
            fn=clear_masks,
            outputs=[mask_output_video, mask_list_video, mask_raw_list_video, mask_ids]
        )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
    sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
    # sam_model = sam_processor = None
    disable_torch_init()
    model, processor = model_init(args_cli.model_path)
    seg_model, seg_processor = model_init(args_cli.seg_model_path)
    # model = processor = None

    # demo.launch()
    demo.launch(
        share=False,
    )