import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
import os
import traceback
from safetensors import safe_open

# Get the current directory
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(BASE_DIR, "clip_seg")

# Model paths
MODEL_PATH = os.path.join(MODELS_DIR, "model.safetensors")

# Global variables for models
PROCESSOR = None
MODEL = None
DEVICE = None

def initialize_models():
    """Initialize the CLIPSeg model and processor."""
    global PROCESSOR, MODEL, DEVICE

    try:
        # Create models directory if it doesn't exist
        os.makedirs(MODELS_DIR, exist_ok=True)

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

        # Initialize processor
        PROCESSOR = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")

        # Initialize model in FP16
        MODEL = CLIPSegForImageSegmentation.from_pretrained(
            "CIDAS/clipseg-rd64-refined",
            cache_dir=MODELS_DIR,
            torch_dtype=torch.float16  # Load in FP16
        )

        # Check if the model file exists before attempting to load
        if not os.path.exists(MODEL_PATH):
            print(f"Error: Model file not found at {MODEL_PATH}")
            return False
            
        # Load the state dict from the .safetensors file using safe_open
        state_dict = {}
        with safe_open(MODEL_PATH, framework="pt", device=str(DEVICE)) as f:
            for key in f.keys():
                state_dict[key] = f.get_tensor(key)

        MODEL.load_state_dict(state_dict, strict=False)
        MODEL.to(DEVICE)  # Model will be on CUDA in FP16
        MODEL.eval()  # Set the model to evaluation mode

        print("Models loaded successfully!")
        
        # Add a GPU availability check
        if torch.cuda.is_available():
            print(f"CUDA is available. Using device: {DEVICE}")
            print(f"CUDA device count: {torch.cuda.device_count()}")
            print(f"Current CUDA device: {torch.cuda.current_device()}")
            print(f"CUDA device name: {torch.cuda.get_device_name(torch.cuda.current_device())}")
        else:
            print("CUDA is not available. Using CPU.")
            
        return True

    except Exception as e:
        print(f"Error initializing models: {str(e)}")
        traceback.print_exc()
        return False

def segment_image(input_image, text_prompt, threshold=0.5):
    """Segment objects in an image based on a text prompt using CLIPSeg."""
    global PROCESSOR, MODEL, DEVICE

    # Check if models are initialized
    if MODEL is None or PROCESSOR is None:
        if not initialize_models():
            return input_image

    if input_image is None:
        return None

    # Convert input_image to PIL Image if it's a NumPy array
    if isinstance(input_image, np.ndarray):
        input_image = Image.fromarray(input_image)

    try:
        # Prepare inputs for the model
        inputs = PROCESSOR(
            text=[text_prompt], images=[input_image], return_tensors="pt"
        )
        # Convert inputs to FP16 and move to device
        inputs = {
            k: (
                v.to(DEVICE).half()
                if isinstance(v, torch.Tensor) and v.dtype == torch.float32
                else v.to(DEVICE)
            )
            for k, v in inputs.items()
        }

        # Get segmentation logits
        with torch.no_grad():
            outputs = MODEL(**inputs)

        # Resize and process logits
        logits = outputs.logits
        upsampled_logits = torch.nn.functional.interpolate(
            logits.unsqueeze(0).unsqueeze(0)
            if logits.dim() == 2
            else logits.unsqueeze(0),
            size=input_image.size[::-1],
            mode="bilinear",
            align_corners=False,
        )

        # Convert logits to probabilities and create a binary mask
        probs = torch.sigmoid(upsampled_logits).squeeze(0).squeeze(0)
        mask = (probs > threshold).cpu().numpy().astype(np.uint8)
        mask = mask.astype(bool)

        # Create output image
        output = np.array(input_image).copy()

        # Create masked version (keep original colors for masked area)
        mask_rgb = np.zeros_like(output)
        mask_rgb[mask] = output[mask]

        # Create background (darkened version of original)
        background = output.copy()
        background = cv2.addWeighted(background, 0.3, np.zeros_like(background), 0.7, 0)

        # Combine background and masked area
        output = background.copy()
        output[mask] = mask_rgb[mask]

        # Add subtle highlight to the masked area
        highlight = np.zeros_like(output)
        highlight[mask] = [20, 20, 20]  # Subtle brightness increase
        output = cv2.add(output, highlight)

        # Optional: Add a subtle border around the masked area
        mask_border = cv2.dilate(mask.astype(np.uint8), None, iterations=1) & ~mask
        output[mask_border] = [255, 255, 255]  # White border

        return output

    except Exception as e:
        print(f"CLIPSeg error: {str(e)}")
        traceback.print_exc()
        return input_image

def create_gradio_interface():
    """Create the Gradio interface."""
    # Initialize models before creating the interface
    initialize_models()

    return gr.Interface(
        fn=segment_image,
        inputs=[
            gr.Image(type="numpy", label="Input Image"),
            gr.Textbox(
                label="Text Prompt",
                placeholder="Enter object to segment (e.g., 'face', 'person')",
            ),
            gr.Slider(minimum=0.0, maximum=0.9, step=0.01, value=0.15, label="Threshold"),
        ],
        outputs=gr.Image(type="numpy", label="Segmented Image"),
        title="Text-Prompted Image Segmentation with CLIPSeg",
        description="Upload an image and enter a text prompt to segment specific objects using CLIPSeg.",
        theme="default",
        css="""
            .gradio-container {max-width: 95%; margin: auto}
            .image-preview {height: auto !important; max-height: 70vh !important; width: 100% !important; object-fit: contain}
            @media (max-width: 768px) {
                .image-preview {max-height: 50vh !important;}
            }
        """,
    )

if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=False, inbrowser=True)