import gradio as gr
import os
from typing import Tuple, Optional
import os
import shutil
import sys
from pathlib import Path
import cv2
import gradio as gr
import numpy as np
import spaces
# import supervision as sv
import torch
from PIL import Image
from tqdm import tqdm
import sys
from pathlib import Path
from huggingface_hub import login
# from dotenv import load_dotenv

# For Hugging Face Spaces, secrets are automatically loaded as environment variables
token = os.getenv("HF_TOKEN")
if token:
    login(token=token)
# Clear Hugging Face cache
# cache_dirs = [
#     "/home/user/.cache/huggingface/",
#     "/home/user/.cache/torch/",
#     "/home/user/.cache/pip/"
# ]

# for cache_dir in cache_dirs:
#     if os.path.exists(cache_dir):
#         print(f"Clearing cache: {cache_dir}")
#         shutil.rmtree(cache_dir, ignore_errors=True)
# Add the current directory to Python path
current_dir = Path(__file__).parent
sys.path.append(str(current_dir))
# sys.path.append("./BiomedParse/")
# BIOMEDPARSE_PATH = Path(__file__).parent / "BiomedParse"
# sys.path.append(str(BIOMEDPARSE_PATH))
# sys.path.append(str(BIOMEDPARSE_PATH / "BiomedParse"))  # Add the inner BiomedParse directory
from modeling.BaseModel import BaseModel
from modeling import build_model
from utilities.arguments import load_opt_from_config_files
from utilities.constants import BIOMED_CLASSES
from inference_utils.inference import interactive_infer_image
from inference_utils.output_processing import check_mask_stats
from inference_utils.processing_utils import read_rgb

import spaces

# breakpoint()
MARKDOWN = """
<div align="center" style="padding: 20px 0;">
    <h1 style="font-size: 3em; margin: 0;">
        ሀ<span style="color: #32CD32;">A</span>ኪ<span style="color: #FFD700;">i</span>ም
        <sup style="font-size: 0.5em;">AI</sup>
    </h1>

    <div style="display: flex; justify-content: center; align-items: center; gap: 15px; margin: 15px 0;">
        <a href="https://cyberbrainai.com/">
            <img src="https://cyberbrainai.com/assets/logo.svg" alt="CyberBrain AI" style="width:40px; height:40px; vertical-align: middle;">
        </a>
        <a href="https://colab.research.google.com/drive/1p3Yf_6xdZPMz5RUtt_NyxrDjrbSgvTDy#scrollTo=t30NqIrCKdAI">
            <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="ድinቅneሽ" style="vertical-align: middle;">
        </a>
        <a href="https://www.youtube.com/watch?v=Dv003fTyO-Y">
            <img src="https://badges.aleen42.com/src/youtube.svg" alt="YouTube" style="vertical-align: middle;">
        </a>
    </div>
</div>
    <div>
        <p style="font-size: 1.4em; line-height: 1.5; margin: 15px 0; text-align: left;">
            This demo integrates BiomedParse, a foundation model for joint segmentation, detection, and recognition across 9 biomedical imaging modalities. 
            The model supports <span style="color: #FF4500;">CT</span>, <span style="color: #4169E1;">MRI</span>, <span style="color: #32CD32;">X-Ray</span>, <span style="color: #9370DB;">Pathology</span>, <span style="color: #FFD700;">Ultrasound</span>, <span style="color: #FF69B4;">Endoscope</span>, <span style="color: #20B2AA;">Fundus</span>, <span style="color: #FF8C00;">Dermoscopy</span>, and <span style="color: #8B008B;">OCT</span>.
        </p>
    </div>

"""

IMAGE_PROCESSING_EXAMPLES = [
    ["BiomedParse Segmentation", 
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/T0011.jpg",
     "Optic disc in retinal Fundus"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/Part_3_226_pathology_breast.png",
     "optic disc, optic cup"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/covid_1585.png",
     "COVID-19 infection in chest X-Ray"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/TCGA_HT_7856_19950831_8_MRI-FLAIR_brain.png",
     "Lower-grade glioma in brain MRI"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/LIDC-IDRI-0140_143_280_CT_lung.png",
     "COVID-19 infection in chest CT"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/144DME_as_F.jpeg",
     "Cystoid macular edema in retinal OCT"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/Part_1_516_pathology_breast.png",
     "Glandular structure in colon Pathology"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/ISIC_0015551.jpg",
     "Melanoma in skin Dermoscopy"],
    ["BiomedParse Segmentation",
     "https://raw.githubusercontent.com/microsoft/BiomedParse/main/examples/C3_EndoCV2021_00462.jpg",
     "Neoplastic polyp in colon Endoscope"]
]

BIOMEDPARSE_MODES = {
    "CT": ["abdomen", "colon", "liver", "lung", "pelvis"],
    "MRI": ["brain", "heart", "prostate", "abdomen"],
    "MRI-FLAIR": ["brain"],
    "MRI-T1-Gd": ["brain"],
    "MRI-T2": ["prostate"],
    "OCT": ["retinal"],
    "X-Ray": ["chest"],
    "Dermoscopy": ["skin"],
    "Endoscope": ["colon"],
    "Fundus": ["retinal"],
    "Pathology": ["bladder", "breast", "cervix", "colon", "esophagus", "kidney", 
                 "liver", "ovarian", "prostate", "stomach", "testis", "thyroid", "uterus"],
    "Ultrasound": ["breast", "heart", "transperineal"]
}

IMAGE_INFERENCE_MODES = [
    "BIOMED SEGMENTATION",
    "BIOMED DETECTION", 
    "BIOMED RECOGNITION",
    "BIOMED SEGMENTATION + DETECTION",
    "BIOMED SEGMENTATION + RECOGNITION",
    "BIOMED DETECTION + RECOGNITION",
    "BIOMED SEGMENTATION + DETECTION + RECOGNITION"
]


def on_mode_dropdown_change(selected_mode):
    if selected_mode in IMAGE_INFERENCE_MODES:
        # Show modality dropdown and hide other inputs initially
        return [
            gr.Dropdown(visible=True, choices=list(BIOMEDPARSE_MODES.keys()), label="Modality"),
            gr.Dropdown(visible=True, label="Anatomical Site"),
            gr.Textbox(visible=False),
            gr.Textbox(visible=False)
        ]
    else:
        # Original behavior for other modes
        return [
            gr.Dropdown(visible=False),
            gr.Dropdown(visible=False),
            gr.Textbox(visible=True),
            gr.Textbox(visible=(selected_mode == None))
        ]

def on_modality_change(modality):
    if modality:
        return gr.Dropdown(choices=BIOMEDPARSE_MODES[modality], visible=True)
    return gr.Dropdown(visible=False)


def initialize_model():
    opt = load_opt_from_config_files(["configs/biomedparse_inference.yaml"])
    pretrained_pth = 'hf_hub:microsoft/BiomedParse'
    opt['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = BaseModel(opt, build_model(opt)).from_pretrained(pretrained_pth).eval()
    with torch.no_grad():
        model.model.sem_seg_head.predictor.lang_encoder.get_text_embeddings(
            BIOMED_CLASSES + ["background"], is_eval=True
        )
    return model


model = initialize_model()


# Utility functions
@spaces.GPU
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process_image(image_path, text_prompts, modality):
    image = read_rgb(image_path)
    text_prompts = [prompt.strip() for prompt in text_prompts.split(',')]
    
    # Run inference
    pred_masks = interactive_infer_image(model, Image.fromarray(image), text_prompts)
    
    # Prepare outputs
    results = []
    dice_scores = []
    p_values = []

    for i, prompt in enumerate(text_prompts):
        # Calculate p-value for the selected modality
        print("PROMPT: ", prompt, flush=True)   
        p_value = check_mask_stats(image, pred_masks[i] * 255, modality, prompt)
        p_values.append(f"P-value for '{prompt}' ({modality}): {p_value:.4f}")

        # Overlay predictions on the image
        overlay_image = image.copy()
        overlay_image[pred_masks[i] > 0.5] = [255, 0, 0]  # Highlight predictions in red
        results.append(overlay_image)

    return results, p_values

# Define Gradio interface
with gr.Blocks() as demo:
    gr.HTML(MARKDOWN)    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Input Image")
            prompts_input = gr.Textbox(lines=2, placeholder="Enter prompts separated by commas...", label="Prompts")
            modality_dropdown = gr.Dropdown(
                choices=list(BIOMEDPARSE_MODES.keys()),
                value=list(BIOMEDPARSE_MODES.keys())[0],
                label="Modality"
            )
            submit_btn = gr.Button("Submit")
        with gr.Column():
            output_gallery = gr.Gallery(label="Findings")
            pvalue_output = gr.Textbox(label="Confidence (P-values)", interactive=False)
    
    submit_btn.click(
        process_image,
        inputs=[image_input, prompts_input, modality_dropdown],
        outputs=[output_gallery, pvalue_output]
    )
    # with gr.Row():
    #         gr.Examples(
    #             fn=process_image,
    #             examples=IMAGE_PROCESSING_EXAMPLES,
    #             inputs=[
    #                 image_processing_mode_dropdown_component,
    #                 image_processing_image_input_component,
    #                 image_processing_text_input_component
    #             ],
    #             outputs=[
    #                 image_processing_image_output_component,
    #                 image_processing_text_output_component
    #             ],
    #             run_on_click=True
    #         )

# Launch the app
demo.launch()