import gradio as gr
import json
import matplotlib.pyplot as plt
import pandas as pd
import io
import base64
import math
import ast
import logging
from matplotlib.widgets import Cursor

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Function to safely parse JSON or Python dictionary input
def parse_input(json_input):
    logger.debug("Attempting to parse input: %s", json_input)
    try:
        # Try to parse as JSON first
        data = json.loads(json_input)
        logger.debug("Successfully parsed as JSON")
        return data
    except json.JSONDecodeError as e:
        logger.error("JSON parsing failed: %s", str(e))
        try:
            # If JSON fails, try to parse as Python literal (e.g., with single quotes)
            data = ast.literal_eval(json_input)
            logger.debug("Successfully parsed as Python literal")
            # Convert Python dictionary to JSON-compatible format (replace single quotes with double quotes)
            def dict_to_json(obj):
                if isinstance(obj, dict):
                    return {str(k): dict_to_json(v) for k, v in obj.items()}
                elif isinstance(obj, list):
                    return [dict_to_json(item) for item in obj]
                else:
                    return obj
            converted_data = dict_to_json(data)
            logger.debug("Converted to JSON-compatible format")
            return converted_data
        except (SyntaxError, ValueError) as e:
            logger.error("Python literal parsing failed: %s", str(e))
            raise ValueError(f"Malformed input: {str(e)}. Ensure property names are in double quotes (e.g., \"content\") or correct Python dictionary format.")

# Function to ensure a value is a float, converting from string if necessary
def ensure_float(value):
    if value is None:
        return None
    if isinstance(value, str):
        try:
            return float(value)
        except ValueError:
            logger.error("Failed to convert string '%s' to float", value)
            return None
    if isinstance(value, (int, float)):
        return float(value)
    return None

# Function to process and visualize log probs with hover and alternatives
def visualize_logprobs(json_input):
    try:
        # Parse the input (handles both JSON and Python dictionaries)
        data = parse_input(json_input)
        
        # Ensure data is a list or dictionary with 'content'
        if isinstance(data, dict) and "content" in data:
            content = data["content"]
        elif isinstance(data, list):
            content = data
        else:
            raise ValueError("Input must be a list or dictionary with 'content' key")

        # Extract tokens, log probs, and top alternatives, skipping None or non-finite values
        tokens = []
        logprobs = []
        top_alternatives = []  # List to store top 3 log probs (selected token + 2 alternatives)
        for entry in content:
            logprob = ensure_float(entry.get("logprob", None))
            if logprob is not None and math.isfinite(logprob):
                tokens.append(entry["token"])
                logprobs.append(logprob)
                # Get top_logprobs, default to empty dict if None
                top_probs = entry.get("top_logprobs", {})
                # Ensure all values in top_logprobs are floats
                finite_top_probs = {}
                for key, value in top_probs.items():
                    float_value = ensure_float(value)
                    if float_value is not None and math.isfinite(float_value):
                        finite_top_probs[key] = float_value
                # Get the top 3 log probs (including the selected token)
                all_probs = {entry["token"]: logprob}  # Add the selected token's logprob
                all_probs.update(finite_top_probs)  # Add alternatives
                sorted_probs = sorted(all_probs.items(), key=lambda x: x[1], reverse=True)
                top_3 = sorted_probs[:3]  # Top 3 log probs (highest to lowest)
                top_alternatives.append(top_3)
            else:
                logger.debug("Skipping entry with logprob: %s (type: %s)", entry.get("logprob"), type(entry.get("logprob", None)))

        # Create the plot with hover functionality
        if logprobs:
            fig, ax = plt.subplots(figsize=(10, 5))
            scatter = ax.plot(range(len(logprobs)), logprobs, marker="o", linestyle="-", color="b", label="Selected Token")[0]
            ax.set_title("Log Probabilities of Generated Tokens")
            ax.set_xlabel("Token Position")
            ax.set_ylabel("Log Probability")
            ax.grid(True)
            ax.set_xticks([])  # Hide X-axis labels by default

            # Add hover functionality using Matplotlib's Cursor for tooltips
            cursor = Cursor(ax, useblit=True, color='red', linewidth=1)
            token_annotations = []
            for i, (x, y) in enumerate(zip(range(len(logprobs)), logprobs)):
                annotation = ax.annotate('', (x, y), xytext=(10, 10), textcoords='offset points', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8), visible=False)
                token_annotations.append(annotation)

            def on_hover(event):
                if event.inaxes == ax:
                    for i, (x, y) in enumerate(zip(range(len(logprobs)), logprobs)):
                        contains, _ = scatter.contains(event)
                        if contains and abs(event.xdata - x) < 0.5 and abs(event.ydata - y) < 0.5:
                            token_annotations[i].set_text(tokens[i])
                            token_annotations[i].set_visible(True)
                            fig.canvas.draw_idle()
                        else:
                            token_annotations[i].set_visible(False)
                            fig.canvas.draw_idle()

            fig.canvas.mpl_connect('motion_notify_event', on_hover)

            # Save plot to a bytes buffer
            buf = io.BytesIO()
            plt.savefig(buf, format="png", bbox_inches="tight", dpi=100)
            buf.seek(0)
            plt.close()

            # Convert to base64 for Gradio
            img_bytes = buf.getvalue()
            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
            img_html = f'<img src="data:image/png;base64,{img_base64}" style="max-width: 100%; height: auto;">'
        else:
            img_html = "No finite log probabilities to plot."

        # Create DataFrame for the table
        table_data = []
        for i, entry in enumerate(content):
            logprob = ensure_float(entry.get("logprob", None))
            if logprob is not None and math.isfinite(logprob) and "top_logprobs" in entry and entry["top_logprobs"] is not None:
                token = entry["token"]
                top_logprobs = entry["top_logprobs"]
                # Ensure all values in top_logprobs are floats
                finite_top_logprobs = {}
                for key, value in top_logprobs.items():
                    float_value = ensure_float(value)
                    if float_value is not None and math.isfinite(float_value):
                        finite_top_logprobs[key] = float_value
                # Extract top 3 alternatives from top_logprobs
                top_3 = sorted(finite_top_logprobs.items(), key=lambda x: x[1], reverse=True)[:3]
                row = [token, f"{logprob:.4f}"]
                for alt_token, alt_logprob in top_3:
                    row.append(f"{alt_token}: {alt_logprob:.4f}")
                while len(row) < 5:
                    row.append("")
                table_data.append(row)

        df = (
            pd.DataFrame(
                table_data,
                columns=[
                    "Token",
                    "Log Prob",
                    "Top 1 Alternative",
                    "Top 2 Alternative",
                    "Top 3 Alternative",
                ],
            )
            if table_data
            else None
        )

        # Generate colored text
        if logprobs:
            min_logprob = min(logprobs)
            max_logprob = max(logprobs)
            if max_logprob == min_logprob:
                normalized_probs = [0.5] * len(logprobs)
            else:
                normalized_probs = [
                    (lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs
                ]

            colored_text = ""
            for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)):
                r = int(255 * (1 - norm_prob))  # Red for low confidence
                g = int(255 * norm_prob)        # Green for high confidence
                b = 0
                color = f"rgb({r}, {g}, {b})"
                colored_text += f'<span style="color: {color}; font-weight: bold;">{token}</span>'
                if i < len(tokens) - 1:
                    colored_text += " "
            colored_text_html = f"<p>{colored_text}</p>"
        else:
            colored_text_html = "No finite log probabilities to display."

        # Create an alternative visualization for top 3 tokens
        alt_viz_html = ""
        if logprobs and top_alternatives:
            alt_viz_html = "<h3>Top 3 Token Log Probabilities</h3><ul>"
            for i, (token, probs) in enumerate(zip(tokens, top_alternatives)):
                alt_viz_html += f"<li>Position {i} (Token: {token}):<br>"
                for tok, prob in probs:
                    alt_viz_html += f"{tok}: {prob:.4f}<br>"
                alt_viz_html += "</li>"
            alt_viz_html += "</ul>"

        return img_html, df, colored_text_html, alt_viz_html

    except Exception as e:
        logger.error("Visualization failed: %s", str(e))
        return f"Error: {str(e)}", None, None, None

# Gradio interface
with gr.Blocks(title="Log Probability Visualizer") as app:
    gr.Markdown("# Log Probability Visualizer")
    gr.Markdown(
        "Paste your JSON or Python dictionary log prob data below to visualize the tokens and their probabilities. Ensure property names are in double quotes (e.g., \"content\") for JSON, or use correct Python dictionary format."
    )

    json_input = gr.Textbox(
        label="JSON Input",
        lines=10,
        placeholder="Paste your JSON (e.g., {\"content\": [...]}) or Python dict (e.g., {'content': [...]}) here...",
    )

    plot_output = gr.HTML(label="Log Probability Plot (Hover for Tokens)")
    table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
    text_output = gr.HTML(label="Colored Text (Confidence Visualization)")
    alt_viz_output = gr.HTML(label="Top 3 Token Log Probabilities")

    btn = gr.Button("Visualize")
    btn.click(
        fn=visualize_logprobs,
        inputs=json_input,
        outputs=[plot_output, table_output, text_output, alt_viz_output],
    )

app.launch()