import gradio as gr
import json
import matplotlib.pyplot as plt
import pandas as pd
import io
import base64
import ast
import math

# Function to process and visualize log probs
def visualize_logprobs(json_input):
    try:
        # Try to parse as JSON first
        try:
            data = json.loads(json_input)
        except json.JSONDecodeError:
            # If JSON fails, try to parse as Python literal (e.g., with single quotes)
            data = ast.literal_eval(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 and log probs, skipping None values and handling non-finite values
        tokens = []
        logprobs = []
        for entry in content:
            if entry['logprob'] is not None and math.isfinite(entry['logprob']):
                tokens.append(entry['token'])
                logprobs.append(entry['logprob'])
        
        # Prepare data for the table
        table_data = []
        for entry in content:
            if entry['logprob'] is not None and math.isfinite(entry['logprob']):
                token = entry['token']
                logprob = entry['logprob']
                top_logprobs = entry['top_logprobs'] or {}
                
                # Filter out non-finite (e.g., -inf, inf, nan) log probs from top_logprobs
                finite_top_logprobs = {k: v for k, v in top_logprobs.items() if math.isfinite(v)}
                # Extract top 3 finite alternatives, sorted by log prob (most probable first)
                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}")
                # Pad with empty strings if fewer than 3 alternatives
                while len(row) < 5:
                    row.append("")
                table_data.append(row)
        
        # Create the plot (only for finite log probs)
        if logprobs:
            plt.figure(figsize=(10, 5))
            plt.plot(range(len(logprobs)), logprobs, marker='o', linestyle='-', color='b')
            plt.title("Log Probabilities of Generated Tokens")
            plt.xlabel("Token Position")
            plt.ylabel("Log Probability")
            plt.grid(True)
            plt.xticks(range(len(logprobs)), tokens, rotation=45, ha='right')
            plt.tight_layout()
            
            # Save plot to a bytes buffer
            buf = io.BytesIO()
            plt.savefig(buf, format='png', bbox_inches='tight')
            buf.seek(0)
            plt.close()
            
            # Convert buffer 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 a DataFrame for the table
        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 based on log probabilities
        if logprobs:
            # Normalize log probs to [0, 1] for color scaling (0 = most uncertain, 1 = most confident)
            min_logprob = min(logprobs)
            max_logprob = max(logprobs)
            if max_logprob == min_logprob:
                normalized_probs = [0.5] * len(logprobs)  # Avoid division by zero
            else:
                normalized_probs = [(lp - min_logprob) / (max_logprob - min_logprob) for lp in logprobs]
            
            # Create HTML for colored text
            colored_text = ""
            for i, (token, norm_prob) in enumerate(zip(tokens, normalized_probs)):
                # Map normalized probability to RGB color (green for high confidence, red for low)
                r = int(255 * (1 - norm_prob))  # Red increases as uncertainty increases
                g = int(255 * norm_prob)        # Green decreases as uncertainty increases
                b = 0                           # Blue stays 0 for simplicity
                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 += " "  # Add space between tokens
            
            colored_text_html = f'<p>{colored_text}</p>'
        else:
            colored_text_html = "No finite log probabilities to display."
        
        return img_html, df, colored_text_html
    
    except Exception as e:
        return f"Error: {str(e)}", 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.")
    
    # Input
    json_input = gr.Textbox(label="JSON Input", lines=10, placeholder="Paste your JSON or Python dict here...")
    
    # Outputs
    plot_output = gr.HTML(label="Log Probability Plot")
    table_output = gr.Dataframe(label="Token Log Probabilities and Top Alternatives")
    text_output = gr.HTML(label="Colored Text (Confidence Visualization)")
    
    # Button to trigger visualization
    btn = gr.Button("Visualize")
    btn.click(
        fn=visualize_logprobs,
        inputs=json_input,
        outputs=[plot_output, table_output, text_output]
    )

# Launch the app
app.launch()