import requests
import gradio as gr

def parse_api_response(data):
    features_by_token = {}
    for result in data['results']:
        token = result['token']
        if token == '<bos>':
            continue
        features_by_token[token] = []
        for feature in result['top_features']:
            features_by_token[token].append({
                'id': feature['feature_index'],
                'activation': feature['activation_value']
            })
    return features_by_token

def analyze_text(text: str):
    if not text:
        return None, ""

    try:
        response = requests.post(
            "https://www.neuronpedia.org/api/search-with-topk",
            headers={"Content-Type": "application/json"},
            json={
                "modelId": "gemma-2-2b",
                "text": text,
                "layer": "20-gemmascope-res-16k"
            }
        )
        response.raise_for_status()
        features = parse_api_response(response.json())
        first_feature = next(iter(next(iter(features.values()))))
        dashboard = create_dashboard(first_feature['id'])
        return features, dashboard
    except Exception as e:
        return None, str(e)

def create_dashboard(feature_id: int):
    return f"""
        <div class="dashboard-container p-4">
            <h3 class="text-lg font-semibold mb-4">Feature {feature_id} Dashboard</h3>
            <iframe
                src="https://www.neuronpedia.org/gemma-2-2b/20-gemmascope-res-16k/{feature_id}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300"
                width="100%"
                height="600"
                frameborder="0"
                class="rounded-lg"
            ></iframe>
        </div>
    """

css = """
.token-section { margin-bottom: 1.5rem; }
.token-header { 
    font-size: 1.25rem;
    font-weight: 600;
    margin-bottom: 0.75rem;
}
.feature-button {
    background: #f8fafc;
    border: 1px solid #e2e8f0;
    border-radius: 0.5rem;
    padding: 0.5rem 1rem;
    margin: 0.25rem;
    font-size: 0.875rem;
    cursor: pointer;
}
.feature-button:hover { background: #f1f5f9; }
.feature-button.selected {
    background: #e0e7ff;
    border-color: #6366f1;
}
"""

def update_dashboard(feature_id, features_state):
    return create_dashboard(feature_id)

with gr.Blocks(css=css) as demo:
    gr.Markdown("# Brand Analyzer")
    gr.Markdown("*Analyze text using Gemma's interpretable neural features*")
    
    features_state = gr.State()
    selected_feature = gr.State()
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(lines=5, label="Input Text")
            analyze_btn = gr.Button("Analyze Features", variant="primary")
            gr.Examples(["WordLift", "Think Different", "Just Do It"], inputs=input_text)
        
        with gr.Column(scale=2):
            features_container = gr.HTML()
            dashboard = gr.HTML()

    def render_features(features_dict, selected=None):
        if not features_dict:
            return ""
        
        html = ""
        for token, features in features_dict.items():
            html += f'<div class="token-section">'
            html += f'<div class="token-header">{token}</div>'
            for feature in features[:3]:
                selected_class = "selected" if selected == feature['id'] else ""
                html += f"""
                <button 
                    onclick='updateDashboard({feature["id"]})'
                    class='feature-button {selected_class}'>
                    Feature {feature["id"]} (Activation: {feature["activation"]:.2f})
                </button>
                """
            html += '</div>'
        return html

    analyze_btn.click(
        analyze_text,
        inputs=[input_text],
        outputs=[features_state, dashboard]
    ).then(
        render_features,
        inputs=[features_state, selected_feature],
        outputs=[features_container]
    )

if __name__ == "__main__":
    demo.launch(share=False)