import os
import json
import gradio as gr
import torch
from PIL import Image, ImageDraw
from transformers import GroundingDinoProcessor
from hf_model import CountEX
from utils import post_process_grounded_object_detection, post_process_grounded_object_detection_with_queries
import google.generativeai as genai
# Global variables for model and processor
model = None
processor = None
device = None

# Configure Gemini
genai.configure(api_key='AIzaSyApqa65vVYTmw4FC4wP-6-_xpBLxXdctxE')
gemini_model = genai.GenerativeModel("gemini-2.0-flash")

PARSING_PROMPT = """Parse sentences of the form "Count A, not B" into two lists—A (include) and B (exclude)—splitting on "and", "or", and commas, and reattaching shared head nouns (e.g., "red and black beans" → "red beans", "black beans").

Rules:
- Remove from B items that are equivalent to items in A (synonyms/variants/abbreviations/regional terms)
- Keep B items that are more specific than A (for fine-grained exclusion)
- If B is more general than A but shares the head noun, remove B (contradictory)

Case 1 — Different head nouns → Keep B
Example 1: Count green apples and red beans, not yellow screws and white rice → A: ["green apples", "red beans"], B: ["yellow screws", "white rice"]
Example 2: Count black beans, not poker chips or nails → A: ["black beans"], B: ["poker chips", "nails"]

Case 2 — Equivalent items → Remove from B
Example 1: Count fries and TV, not chips and television → A: ["fries", "TV"], B: []
Example 2: Count garbanzo beans and couch, not chickpeas and sofa → A: ["garbanzo beans", "couch"], B: []

Case 3 — B more specific than A → Keep B (for fine-grained exclusion)
Example 1: Count apples and beans, not green apples and black beans → A: ["apples", "beans"], B: ["green apples", "black beans"]
Example 2: Count beans, not white beans or yellow beans → A: ["beans"], B: ["white beans", "yellow beans"]
Example 3: Count people, not women → A: ["people"], B: ["women"]

Case 4 — B more general than A → Remove B (contradictory)
Example 1: Count green apples, not apples → A: ["green apples"], B: []
Example 2: Count red beans and green apples, not beans and apples → A: ["red beans", "green apples"], B: []

User instruction: {instruction}

Respond ONLY with a JSON object in this exact format, no other text:
{{"A": ["item1", "item2"], "B": ["item3"]}}
"""

def parse_counting_instruction(instruction: str) -> tuple[str, str]:
    """
    Parse natural language counting instruction using Gemini 2.0 Flash.
    
    Args:
        instruction: Natural language instruction like "count apples, not green apples"
    
    Returns:
        tuple: (positive_caption, negative_caption)
    """
    try:
        prompt = PARSING_PROMPT.format(instruction=instruction)
        response = gemini_model.generate_content(prompt)
        response_text = response.text.strip()
        
        # Clean up response - remove markdown code blocks if present
        if response_text.startswith("```"):
            response_text = response_text.split("```")[1]
            if response_text.startswith("json"):
                response_text = response_text[4:]
        response_text = response_text.strip()
        
        result = json.loads(response_text)
        
        # Convert lists to caption strings
        pos_items = result.get("A", [])
        neg_items = result.get("B", [])
        
        # Join items with " and " and add period
        pos_caption = " and ".join(pos_items) + "." if pos_items else ""
        neg_caption = " and ".join(neg_items) + "." if neg_items else "None."
        
        return pos_caption, neg_caption
        
    except Exception as e:
        print(f"Error parsing instruction: {e}")
        # Fallback: treat entire instruction as positive caption
        return instruction.strip() + ".", "None."


def load_model():
    """Load model and processor once at startup"""
    global model, processor, device

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

    # Load model - change path for HF Spaces
    model_id = "yifehuang97/CountEx_KC_aug_v3_12131215"  # Change to your HF model repo
    model = CountEX.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
    model = model.to(torch.bfloat16)
    model = model.to(device)
    model.eval()

    # Load processor
    processor_id = "fushh7/llmdet_swin_tiny_hf"
    processor = GroundingDinoProcessor.from_pretrained(processor_id)

    return model, processor, device


import numpy as np


def discriminative_point_suppression(
    points, 
    neg_points, 
    pos_queries,      # (N, D) numpy array
    neg_queries,      # (M, D) numpy array
    image_size,
    pixel_threshold=5,
    similarity_threshold=0.3,
):
    """
    Discriminative Point Suppression (DPS):
    
    Step 1: Find spatially closest negative point for each positive point
    Step 2: If distance < pixel_threshold, check query similarity
    Step 3: Suppress only if query similarity > similarity_threshold
    
    This two-stage design ensures suppression only when predictions are
    both spatially overlapping AND semantically conflicting.
    
    Args:
        points: List of [x, y] positive points (normalized, 0-1)
        neg_points: List of [x, y] negative points (normalized, 0-1)
        pos_queries: (N, D) query embeddings for positive predictions
        neg_queries: (M, D) query embeddings for negative predictions
        image_size: (width, height) in pixels
        pixel_threshold: spatial distance threshold in pixels
        similarity_threshold: cosine similarity threshold for semantic conflict
    
    Returns:
        filtered_points: points after suppression
        filtered_indices: indices of kept points
        suppression_info: dict with detailed suppression decisions
    """
    if not neg_points or not points:
        return points, list(range(len(points))), {}
    
    width, height = image_size
    N, M = len(points), len(neg_points)
    
    # === Step 1: Spatial Matching ===
    points_arr = np.array(points) * np.array([width, height])  # (N, 2)
    neg_points_arr = np.array(neg_points) * np.array([width, height])  # (M, 2)
    
    # Compute pairwise distances
    spatial_dist = np.linalg.norm(
        points_arr[:, None, :] - neg_points_arr[None, :, :], axis=-1
    )  # (N, M)
    
    # Find nearest negative for each positive
    nearest_neg_idx = spatial_dist.argmin(axis=1)  # (N,)
    nearest_neg_dist = spatial_dist.min(axis=1)    # (N,)
    
    # Check spatial condition
    spatially_close = nearest_neg_dist < pixel_threshold  # (N,)
    
    # === Step 2: Query Similarity Check (only for spatially close pairs) ===
    # Normalize queries
    pos_q = pos_queries / (np.linalg.norm(pos_queries, axis=-1, keepdims=True) + 1e-8)
    neg_q = neg_queries / (np.linalg.norm(neg_queries, axis=-1, keepdims=True) + 1e-8)
    
    # Compute similarity only for matched pairs
    matched_neg_q = neg_q[nearest_neg_idx]  # (N, D)
    query_sim = (pos_q * matched_neg_q).sum(axis=-1)  # (N,) cosine similarity
    
    # Check semantic condition
    semantically_similar = query_sim > similarity_threshold  # (N,)
    
    # === Step 3: Joint Decision ===
    # Suppress only if BOTH conditions are met
    should_suppress = spatially_close & semantically_similar  # (N,)
    
    # === Filter ===
    keep_mask = ~should_suppress
    filtered_points = np.array(points)[keep_mask].tolist()
    filtered_indices = np.where(keep_mask)[0].tolist()
    
    # === Suppression Info ===
    suppression_info = {
        "nearest_neg_idx": nearest_neg_idx.tolist(),
        "nearest_neg_dist": nearest_neg_dist.tolist(),
        "query_similarity": query_sim.tolist(),
        "spatially_close": spatially_close.tolist(),
        "semantically_similar": semantically_similar.tolist(),
        "suppressed_indices": np.where(should_suppress)[0].tolist(),
    }
    
    return filtered_points, filtered_indices, suppression_info


def count_objects(image, instruction, box_threshold, point_radius, point_color):
    """
    Main inference function for counting objects

    Args:
        image: Input PIL Image
        instruction: Natural language instruction (e.g., "count apples, not green apples")
        box_threshold: Detection confidence threshold
        point_radius: Radius of visualization points
        point_color: Color of visualization points

    Returns:
        Annotated image, count, and parsed captions
    """
    global model, processor, device

    if model is None:
        load_model()

    # Parse instruction using Gemini
    pos_caption, neg_caption = parse_counting_instruction(instruction)
    parsed_info = f"Positive: {pos_caption}\nNegative: {neg_caption}"

    # Ensure image is RGB
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Process positive caption
    pos_inputs = processor(
        images=image,
        text=pos_caption,
        return_tensors="pt",
        padding=True
    )
    pos_inputs = pos_inputs.to(device)
    pos_inputs['pixel_values'] = pos_inputs['pixel_values'].to(torch.bfloat16)

    # Process negative caption
    use_neg = bool(neg_caption and neg_caption.strip() and neg_caption != '.' and neg_caption != 'None.')

    if not use_neg:
        neg_caption = "None."
    neg_inputs = processor(
        images=image,
        text=neg_caption,
        return_tensors="pt",
        padding=True
    )
    neg_inputs = {k: v.to(device) for k, v in neg_inputs.items()}
    neg_inputs['pixel_values'] = neg_inputs['pixel_values'].to(torch.bfloat16)

    # Add negative inputs to positive inputs dict
    pos_inputs['neg_token_type_ids'] = neg_inputs['token_type_ids']
    pos_inputs['neg_attention_mask'] = neg_inputs['attention_mask']
    pos_inputs['neg_pixel_mask'] = neg_inputs['pixel_mask']
    pos_inputs['neg_pixel_values'] = neg_inputs['pixel_values']
    pos_inputs['neg_input_ids'] = neg_inputs['input_ids']
    pos_inputs['use_neg'] = True

    # Run inference
    with torch.no_grad():
        outputs = model(**pos_inputs)

    # Post-process outputs
    outputs["pred_points"] = outputs["pred_boxes"][:, :, :2]
    outputs["pred_logits"] = outputs["logits"]

    threshold = box_threshold if box_threshold > 0 else model.box_threshold
    pos_queries = outputs["pos_queries"].squeeze(0).float()
    neg_queries = outputs["neg_queries"].squeeze(0).float()
    pos_queries = pos_queries[-1].squeeze(0)
    neg_queries = neg_queries[-1].squeeze(0)
    pos_queries = pos_queries.unsqueeze(0)
    neg_queries = neg_queries.unsqueeze(0)
    results = post_process_grounded_object_detection_with_queries(outputs, pos_queries, box_threshold=threshold)[0]

    boxes = results["boxes"]
    boxes = [box.tolist() for box in boxes]
    points = [[box[0], box[1]] for box in boxes]

    # Negative prediction 
    neg_points = []
    neg_results = None
    if "neg_pred_boxes" in outputs and "neg_logits" in outputs:
        neg_outputs = outputs.copy()
        neg_outputs["pred_boxes"] = outputs["neg_pred_boxes"]
        neg_outputs["logits"] = outputs["neg_logits"]
        neg_outputs["pred_points"] = outputs["neg_pred_boxes"][:, :, :2]
        neg_outputs["pred_logits"] = outputs["neg_logits"]

        neg_results = post_process_grounded_object_detection_with_queries(neg_outputs, neg_queries, box_threshold=threshold)[0]
        neg_boxes = neg_results["boxes"]
        neg_boxes = [box.tolist() for box in neg_boxes]
        neg_points = [[box[0], box[1]] for box in neg_boxes]
    
    pos_queries_np = results["queries"].cpu().numpy()
    neg_queries_np = neg_results["queries"].cpu().numpy() if neg_results else np.array([])
    
    img_size = image.size
    
    if len(neg_points) > 0 and len(neg_queries_np) > 0:
        filtered_points, kept_indices, suppression_info = discriminative_point_suppression(
            points, 
            neg_points, 
            pos_queries_np,
            neg_queries_np,
            image_size=img_size,
            pixel_threshold=5,
            similarity_threshold=0.3,
        )
        filtered_boxes = [boxes[i] for i in kept_indices]
    else:
        filtered_points = points
        filtered_boxes = boxes
    
    points = filtered_points
    boxes = filtered_boxes

    # Visualize results
    img_w, img_h = image.size
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)

    for point in points:
        x = point[0] * img_w
        y = point[1] * img_h
        draw.ellipse(
            [x - point_radius, y - point_radius, x + point_radius, y + point_radius],
            fill=point_color
        )

    count = len(points)

    return img_draw, f"Count: {count}", parsed_info


def count_objects_manual(image, pos_caption, neg_caption, box_threshold, point_radius, point_color):
    """
    Manual mode: directly use provided positive and negative captions.
    """
    global model, processor, device

    if model is None:
        load_model()

    # Ensure captions end with period
    if pos_caption and not pos_caption.endswith('.'):
        pos_caption = pos_caption + '.'
    if neg_caption and not neg_caption.endswith('.'):
        neg_caption = neg_caption + '.'
    
    if not neg_caption or neg_caption.strip() == '':
        neg_caption = "None."

    parsed_info = f"Positive: {pos_caption}\nNegative: {neg_caption}"

    # Ensure image is RGB
    if image.mode != "RGB":
        image = image.convert("RGB")

    # Process positive caption
    pos_inputs = processor(
        images=image,
        text=pos_caption,
        return_tensors="pt",
        padding=True
    )
    pos_inputs = pos_inputs.to(device)
    pos_inputs['pixel_values'] = pos_inputs['pixel_values'].to(torch.bfloat16)

    # Process negative caption
    use_neg = bool(neg_caption and neg_caption.strip() and neg_caption != '.' and neg_caption != 'None.')

    if not use_neg:
        neg_caption = "None."
    neg_inputs = processor(
        images=image,
        text=neg_caption,
        return_tensors="pt",
        padding=True
    )
    neg_inputs = {k: v.to(device) for k, v in neg_inputs.items()}
    neg_inputs['pixel_values'] = neg_inputs['pixel_values'].to(torch.bfloat16)

    # Add negative inputs to positive inputs dict
    pos_inputs['neg_token_type_ids'] = neg_inputs['token_type_ids']
    pos_inputs['neg_attention_mask'] = neg_inputs['attention_mask']
    pos_inputs['neg_pixel_mask'] = neg_inputs['pixel_mask']
    pos_inputs['neg_pixel_values'] = neg_inputs['pixel_values']
    pos_inputs['neg_input_ids'] = neg_inputs['input_ids']
    pos_inputs['use_neg'] = True

    # Run inference
    with torch.no_grad():
        outputs = model(**pos_inputs)

    # Post-process outputs
    outputs["pred_points"] = outputs["pred_boxes"][:, :, :2]
    outputs["pred_logits"] = outputs["logits"]

    threshold = box_threshold if box_threshold > 0 else model.box_threshold
    pos_queries = outputs["pos_queries"].squeeze(0).float()
    neg_queries = outputs["neg_queries"].squeeze(0).float()
    pos_queries = pos_queries[-1].squeeze(0)
    neg_queries = neg_queries[-1].squeeze(0)
    pos_queries = pos_queries.unsqueeze(0)
    neg_queries = neg_queries.unsqueeze(0)
    results = post_process_grounded_object_detection_with_queries(outputs, pos_queries, box_threshold=threshold)[0]

    boxes = results["boxes"]
    boxes = [box.tolist() for box in boxes]
    points = [[box[0], box[1]] for box in boxes]

    # Negative prediction 
    neg_points = []
    neg_results = None
    if "neg_pred_boxes" in outputs and "neg_logits" in outputs:
        neg_outputs = outputs.copy()
        neg_outputs["pred_boxes"] = outputs["neg_pred_boxes"]
        neg_outputs["logits"] = outputs["neg_logits"]
        neg_outputs["pred_points"] = outputs["neg_pred_boxes"][:, :, :2]
        neg_outputs["pred_logits"] = outputs["neg_logits"]

        neg_results = post_process_grounded_object_detection_with_queries(neg_outputs, neg_queries, box_threshold=threshold)[0]
        neg_boxes = neg_results["boxes"]
        neg_boxes = [box.tolist() for box in neg_boxes]
        neg_points = [[box[0], box[1]] for box in neg_boxes]
    
    pos_queries_np = results["queries"].cpu().numpy()
    neg_queries_np = neg_results["queries"].cpu().numpy() if neg_results else np.array([])
    
    img_size = image.size
    
    if len(neg_points) > 0 and len(neg_queries_np) > 0:
        filtered_points, kept_indices, suppression_info = discriminative_point_suppression(
            points, 
            neg_points, 
            pos_queries_np,
            neg_queries_np,
            image_size=img_size,
            pixel_threshold=5,
            similarity_threshold=0.3,
        )
        filtered_boxes = [boxes[i] for i in kept_indices]
    else:
        filtered_points = points
        filtered_boxes = boxes
    
    points = filtered_points
    boxes = filtered_boxes

    # Visualize results
    img_w, img_h = image.size
    img_draw = image.copy()
    draw = ImageDraw.Draw(img_draw)

    for point in points:
        x = point[0] * img_w
        y = point[1] * img_h
        draw.ellipse(
            [x - point_radius, y - point_radius, x + point_radius, y + point_radius],
            fill=point_color
        )

    count = len(points)

    return img_draw, f"Count: {count}", parsed_info


# Create Gradio interface
def create_demo():
    with gr.Blocks(title="CountEx: Discriminative Visual Counting") as demo:
        gr.Markdown("""
        # CountEx: Fine-Grained Counting via Exemplars and Exclusion
        Count specific objects in images using text prompts with exclusion capability.
        """)
        
        # State to track current input mode
        current_mode = gr.State(value="natural_language")

        with gr.Row():
            # Left column - Input
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil", label="Input Image")
                
                with gr.Tabs() as input_tabs:
                    # Tab 1: Natural Language Input
                    with gr.TabItem("Natural Language", id=0) as tab_nl:
                        instruction = gr.Textbox(
                            label="Counting Instruction",
                            placeholder="e.g., Count apples, not green apples",
                            value="Count apples, not green apples",
                            lines=2
                        )
                        gr.Markdown("""
                        **Examples:**
                        - "Count apples, not green apples"
                        - "Count red and black beans, exclude white beans"  
                        - "Count people, not women"
                        """)
                    
                    # Tab 2: Manual Input
                    with gr.TabItem("Manual Input", id=1) as tab_manual:
                        pos_caption = gr.Textbox(
                            label="Positive Prompt (objects to count)",
                            placeholder="e.g., apple",
                            value="apple."
                        )
                        neg_caption = gr.Textbox(
                            label="Negative Prompt (objects to exclude)",
                            placeholder="e.g., green apple",
                            value="None."
                        )
                
                # Single submit button outside tabs
                submit_btn = gr.Button("Count Objects", variant="primary", size="lg")
                
                # Shared settings
                with gr.Accordion("Advanced Settings", open=False):
                    box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.42,
                        step=0.01,
                        label="Detection Threshold"
                    )
                    point_radius = gr.Slider(
                        minimum=1,
                        maximum=20,
                        value=5,
                        step=1,
                        label="Point Radius"
                    )
                    point_color = gr.Dropdown(
                        choices=["blue", "red", "green", "yellow", "cyan", "magenta", "white"],
                        value="blue",
                        label="Point Color"
                    )

            # Right column - Output
            with gr.Column(scale=1):
                output_image = gr.Image(type="pil", label="Result")
                count_output = gr.Textbox(label="Count Result")
                parsed_output = gr.Textbox(label="Parsed Captions", lines=2)

        # Examples for Natural Language mode
        gr.Markdown("### Examples (Natural Language)")
        gr.Examples(
            examples=[
                ["examples/apples.png", "Count apples, not green apples"],
                ["examples/apples.png", "Count apples, exclude red apples"],
                ["examples/apple.jpg", "Count green apples"],
                ["examples/apple.jpg", "Count apples, exclude red apples"],
                ["examples/apple.jpg", "Count apples, exclude green apples"],
                ["examples/black_beans.jpg", "Count black beans and soy beans"],
                ["examples/candy.jpg", "Count brown coffee candy, exclude black coffee candy"],
                ["examples/strawberry.jpg", "Count blueberries and strawberry"],
                ["examples/strawberry2.jpg", "Count blueberries, exclude strawberry"],
                ["examples/women.jpg", "Count people, not women"],
                ["examples/women.jpg", "Count people, not man"],
                ["examples/boat-1.jpg", "Count boats, exclude blue boats"],
                ["examples/boat-1.jpg", "Count boats, exclude red boats"],
            ],
            inputs=[input_image, instruction],
            outputs=[output_image, count_output, parsed_output],
            fn=count_objects,
            cache_examples=False,
        )
        
        # Update mode when tab changes
        def set_mode_nl():
            return "natural_language"
        
        def set_mode_manual():
            return "manual"
        
        tab_nl.select(fn=set_mode_nl, outputs=[current_mode])
        tab_manual.select(fn=set_mode_manual, outputs=[current_mode])
        
        # Unified handler that routes based on mode
        def handle_submit(mode, image, instr, pos_cap, neg_cap, threshold, radius, color):
            if mode == "natural_language":
                return count_objects(image, instr, threshold, radius, color)
            else:
                return count_objects_manual(image, pos_cap, neg_cap, threshold, radius, color)

        # Single button click handler
        submit_btn.click(
            fn=handle_submit,
            inputs=[current_mode, input_image, instruction, pos_caption, neg_caption, 
                    box_threshold, point_radius, point_color],
            outputs=[output_image, count_output, parsed_output]
        )

    return demo


if __name__ == "__main__":
    # Load model at startup
    print("Loading model...")
    load_model()
    print("Model loaded!")

    # Create and launch demo
    demo = create_demo()
    demo.launch()