import pandas as pd
import json
import plotly.express as px
import gradio as gr

def load_data():
    file_path = "causalquest_embeddings_tsne.json"
    with open(file_path, 'r') as f:
        data = json.load(f)
    
    df = pd.DataFrame({
        'query_id': data['query_id'],
        'query': data['query'],
        'source': data['source'],
        'is_causal': data['is_causal'],
        'domain_class': data['domain_class'],
        'action_class': data['action_class'],
        'is_subjective': data['is_subjective'],
        'x': data['tsne_x'],
        'y': data['tsne_y']
    })

    # subsample data for faster visualization
    df = df.sample(1000, random_state=42)
    
    return df

def create_plot(df):
    fig = px.scatter(df, x='x', y='y', hover_data=['query', 'source', 'is_causal', 'domain_class', 'action_class', 'is_subjective'])
    
    # Remove the hover box
    fig.update_traces(hoverinfo='none', hovertemplate=None)
    
    # Function to wrap long text
    def wrap_text(text, max_length=50):
        if len(text) <= max_length:
            return text
        words = text.split()
        lines = []
        current_line = []
        current_length = 0
        for word in words:
            if current_length + len(word) + 1 <= max_length:
                current_line.append(word)
                current_length += len(word) + 1
            else:
                lines.append(' '.join(current_line))
                current_line = [word]
                current_length = len(word)
        if current_line:
            lines.append(' '.join(current_line))
        return '<br>'.join(lines)
    
    # Apply text wrapping to queries
    df['wrapped_query'] = df['query'].apply(wrap_text)
    
    # Add custom hover text with wrapped query
    fig.update_traces(
        customdata=df[['wrapped_query', 'source', 'is_causal', 'domain_class', 'action_class', 'is_subjective']],
        hovertemplate='Query: %{customdata[0]}<br>Source: %{customdata[1]}<br>Is Causal: %{customdata[2]}'
    )
    
    return fig

def visualize_data(processed_file):
    df = load_data()
    fig = create_plot(df)
    return fig

# Create Gradio interface
iface = gr.Interface(
    fn=visualize_data,
    inputs=None,
    outputs=gr.Plot(label="tsne Visualization"),
    title="CausalQuest Visualization",
    description="Visualize causalquest data using pre-computed embeddings and tsne coordinates",
    allow_flagging="never",
    fill_width=True
)

# Launch the app
iface.launch()