import requests
import gradio as gr

def get_features(text: str):
    url = "https://www.neuronpedia.org/api/search-with-topk"
    payload = {
        "modelId": "gemma-2-2b", 
        "text": text,
        "layer": "20-gemmascope-res-16k"
    }
    try:
        response = requests.post(url, headers={"Content-Type": "application/json"}, json=payload)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return None

def create_dashboard(feature_id: int) -> str:
    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>
    """

def handle_feature_click(feature_id):
    return create_dashboard(feature_id)

def analyze_text(text: str):
    if not text:
        return [], ""
        
    features_data = get_features(text)
    if not features_data:
        return [], ""

    features = []
    first_feature_id = None
    
    for result in features_data['results']:
        if result['token'] == '<bos>':
            continue
            
        token = result['token']
        token_features = []
        
        for feature in result['top_features'][:3]:
            feature_id = feature['feature_index']
            if first_feature_id is None:
                first_feature_id = feature_id
                
            token_features.append({
                "token": token,
                "id": feature_id,
                "activation": feature['activation_value']
            })
        
        features.append({"token": token, "features": token_features})
    
    return features, create_dashboard(first_feature_id) if first_feature_id else ""

css = """
@import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');

body { font-family: 'Open Sans', sans-serif !important; }

.dashboard-container {
    border: 1px solid #e0e5ff;
    border-radius: 8px;
    background-color: #ffffff;
}

.token-header {
    font-size: 1.25rem;
    font-weight: 600;
    margin-top: 1rem;
    margin-bottom: 0.5rem;
}

.feature-button {
    display: inline-block;
    margin: 0.25rem;
    padding: 0.5rem 1rem;
    background-color: #f3f4f6;
    border: 1px solid #e5e7eb;
    border-radius: 0.375rem;
    font-size: 0.875rem;
}

.feature-button:hover {
    background-color: #e5e7eb;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.Markdown("# Brand Analyzer", elem_classes="text-2xl font-bold mb-2")
    gr.Markdown("*Analyze text using interpretable neural features*", elem_classes="text-gray-600 mb-6")
    
    features_state = gr.State([])
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                lines=5,
                placeholder="Enter text to analyze...",
                label="Input Text"
            )
            analyze_btn = gr.Button("Analyze Features", variant="primary")
            gr.Examples(
                examples=["WordLift", "Think Different", "Just Do It"],
                inputs=input_text
            )
        
        with gr.Column(scale=2):
            @gr.render(inputs=features_state)
            def render_features(features):
                if not features:
                    return
                
                for token_group in features:
                    gr.Markdown(f"### {token_group['token']}")
                    with gr.Row():
                        for feature in token_group['features']:
                            btn = gr.Button(
                                f"Feature {feature['id']} (Activation: {feature['activation']:.2f})",
                                elem_classes=["feature-button"]
                            )
                            btn.click(
                                fn=lambda fid=feature['id']: handle_feature_click(fid),
                                outputs=dashboard
                            )
            
            dashboard = gr.HTML()
    
    analyze_btn.click(
        fn=analyze_text,
        inputs=[input_text],
        outputs=[features_state, dashboard]
    )

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