import streamlit as st
import numpy as np
from sentence_transformers import SentenceTransformer, util
import plotly.graph_objects as go
import plotly.express as px
from typing import List, Tuple
import pandas as pd

# Page configuration
st.set_page_config(layout="wide", page_title="🎯 Sentence Transformer Explorer")

# Initialize session state for sentences if not exists
if 'sentences' not in st.session_state:
    st.session_state.sentences = [
        "I love programming in Python",
        "Coding with Python is my favorite",
        "The weather is beautiful today"
    ]

# Load model
@st.cache_resource
def load_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

model = load_model()

# Example templates
EXAMPLE_TEMPLATES = {
    "Similar Meanings": [
        "I love programming in Python",
        "Coding with Python is my favorite",
        "I enjoy developing software using Python"
    ],
    "Different Topics": [
        "The cat sleeps on the mat",
        "Python is a programming language",
        "The weather is beautiful today"
    ],
    "Semantic Relations": [
        "Paris is the capital of France",
        "Berlin is the capital of Germany",
        "London is the capital of England"
    ]
}

def load_example_sentences():
    selected_template = st.session_state.get('template_selection')
    st.session_state.sentences = EXAMPLE_TEMPLATES[selected_template]
    # Force update of text inputs
    for i, sentence in enumerate(st.session_state.sentences):
        st.session_state[f'sentence_{i}'] = sentence

def main():
    st.title("🎯 Interactive Sentence Transformer Explorer")
    
    with st.expander("ℹ️ How it works", expanded=True):
        st.markdown("""
        This interactive tool helps you understand how Sentence Transformers work:
        
        1. **Sentence Embedding**: Convert sentences into numerical vectors
        2. **Word Importance**: See how each word contributes to the final embedding
        3. **Similarity Analysis**: Compare how similar sentences are to each other
        4. **Interactive Examples**: Try different sentences and see the results
        """)
    
    # Example selection
    col1, col2 = st.columns(2)
    
    with col1:
        st.selectbox(
            "Choose an example template:",
            options=list(EXAMPLE_TEMPLATES.keys()),
            key='template_selection'
        )
    
    with col2:
        st.button(
            "Load Example",
            on_click=load_example_sentences,
            type="primary"
        )
    
    # Dynamic sentence input
    num_sentences = st.slider("Number of sentences:", 2, 5, 3)
    sentences = []
    
    # Create text inputs with keys
    for i in range(num_sentences):
        default_value = st.session_state.sentences[i] if i < len(st.session_state.sentences) else ""
        sentence = st.text_input(
            f"Sentence {i+1}",
            value=default_value,
            key=f'sentence_{i}'
        )
        sentences.append(sentence)
    
    if st.button("Analyze Sentences", type="primary"):
        if all(sentences):
            # Your existing analysis code here...
            embeddings = model.encode(sentences)
            similarity_matrix = util.cos_sim(embeddings, embeddings).numpy()
            
            st.subheader("📊 Analysis Results")
            
            tab1, tab2 = st.tabs(["Sentence Similarity", "Embedding Analysis"])
            
            with tab1:
                # Create similarity heatmap
                fig = go.Figure(data=go.Heatmap(
                    z=similarity_matrix,
                    x=sentences,
                    y=sentences,
                    colorscale='RdBu',
                    text=np.round(similarity_matrix, 3),
                    texttemplate='%{text}',
                    textfont={"size": 10},
                    hoverongaps=False
                ))
                
                fig.update_layout(
                    title="Sentence Similarity Matrix",
                    height=400
                )
                st.plotly_chart(fig, use_container_width=True)
                
                # Add similarity interpretation
                st.markdown("#### 💡 Interpretation")
                for i in range(len(sentences)):
                    for j in range(i+1, len(sentences)):
                        similarity = similarity_matrix[i][j]
                        interpretation = (
                            "Very similar" if similarity > 0.8 
                            else "Moderately similar" if similarity > 0.5 
                            else "Different"
                        )
                        st.write(f"Sentences {i+1} & {j+1}: {interpretation} ({similarity:.3f})")
            
            with tab2:
                # Create embedding statistics
                embedding_stats = pd.DataFrame({
                    'Sentence': sentences,
                    'Embedding_Length': [np.linalg.norm(emb) for emb in embeddings],
                    'Mean_Value': [np.mean(emb) for emb in embeddings],
                    'Std_Dev': [np.std(emb) for emb in embeddings]
                })
                
                st.dataframe(embedding_stats)
                
                st.markdown("""
                #### 📝 Understanding Embeddings
                - **Embedding Length**: Represents the magnitude of the vector
                - **Mean Value**: Average of all dimensions
                - **Standard Deviation**: Spread of values across dimensions
                """)
        else:
            st.warning("Please enter all sentences before analyzing.")

if __name__ == "__main__":
    main()