import streamlit as st
import pandas as pd
import ast
import json
import re
import time
from sentence_transformers import SentenceTransformer, util
from google.genai import Client, types
from typing import List, Dict
import os

os.environ["HF_HOME"] = "/app/hf_cache"

# Page configuration
st.set_page_config(
    page_title="SAOKE Problem Solver",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 1rem;
        background: linear-gradient(45deg, #1f77b4, #ff7f0e);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    }
    
    .step-header {
        font-size: 1.5rem;
        font-weight: bold;
        color: #2e7d32;
        margin: 1rem 0;
        padding: 0.5rem;
        border-left: 4px solid #4caf50;
        background-color: #e8f5e8;
    }
    
    .effect-card {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 1rem;
        margin: 0.5rem 0;
        background-color: #f9f9f9;
    }
    
    .mechanism-card {
        border: 1px solid #2196f3;
        border-radius: 8px;
        padding: 1rem;
        margin: 0.5rem 0;
        background-color: #e3f2fd;
    }
    
    .solution-container {
        border: 2px solid #4caf50;
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
        background-color: #f1f8e9;
    }
    
    .progress-container {
        margin: 1rem 0;
    }
    
    .status-success {
        color: #4caf50;
        font-weight: bold;
    }
    
    .status-processing {
        color: #ff9800;
        font-weight: bold;
    }
    
    .status-error {
        color: #f44336;
        font-weight: bold;
    }
</style>
""", unsafe_allow_html=True)

class SAOKEWebApp:
    def __init__(self):
        # Initialize session state
        if 'step' not in st.session_state:
            st.session_state.step = 1
        if 'problem' not in st.session_state:
            st.session_state.problem = ""
        if 'effects' not in st.session_state:
            st.session_state.effects = []
        if 'mechanisms' not in st.session_state:
            st.session_state.mechanisms = []
        if 'solution' not in st.session_state:
            st.session_state.solution = ""
        if 'models_loaded' not in st.session_state:
            st.session_state.models_loaded = False
        if 'data_loaded' not in st.session_state:
            st.session_state.data_loaded = False
        if 'top_50_mechanisms' not in st.session_state:
            st.session_state.top_50_mechanisms = {}
        if 'mechanism_indices' not in st.session_state:
            st.session_state.mechanism_indices = {}
        if 'original_mechanisms' not in st.session_state:
            st.session_state.original_mechanisms = []
        if 'original_llm_indices' not in st.session_state:
            st.session_state.original_llm_indices = {}
    
    @st.cache_resource
    def load_models(_self):
        """Load sentence transformer model"""
        try:
            model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
            return model
        except Exception as e:
            st.error(f"Error loading model: {str(e)}")
            return None
    
    @st.cache_data
    def load_data(_self):
        """Load and process SAOKE mechanisms data"""
        try:
            df = pd.read_excel("SAOKE_technologies.xlsx")
            
            list_mechanisms = []
            for row in df.iterrows():
                output = {
                    "technology": row[1]["technology"],
                    "mechanisms": []
                }
                
                for idx, saoke in enumerate(ast.literal_eval(row[1]["mechanisms"])):
                    mechanism = {
                        "id": idx + 1,
                        "dependency": None if idx == 0 else idx,
                        "SAOKE": saoke
                    }
                    output["mechanisms"].append(mechanism)
                list_mechanisms.append(output)
            
            # Filter mechanisms with effects
            list_mechanisms = [d for d in list_mechanisms 
                             if all(m.get("SAOKE", {}).get("effect") is not None 
                                   for m in d["mechanisms"])]
            
            return list_mechanisms
        except Exception as e:
            st.error(f"Error loading data: {str(e)}")
            return []
    
    def get_gemini_client(self):
        """Initialize Gemini client"""
        return Client(api_key=os.getenv("GOOGLE_API_KEY"))
    
    def ask_gemini(self, prompt: str):
        """Send request to Gemini API"""
        try:
            client = self.get_gemini_client()
            grounding_tool = types.Tool(google_search=types.GoogleSearch())
            config = types.GenerateContentConfig(tools=[grounding_tool])
            
            response = client.models.generate_content(
                model="gemini-2.5-pro",
                contents=prompt,
                config=config,
            )
            return response
        except Exception as e:
            st.error(f"Gemini API error: {e}")
            time.sleep(2)
            return self.ask_gemini(prompt)
    
    def extract_effects(self, problem: str) -> List[Dict]:
        """Extract effects from problem using Gemini"""
        prompt = f"""
You are extracting EFFECTS from a technical PROBLEM to support SAOKE-based mechanism retrieval, where cosine similarity will match technical functions between effects and mechanisms.

**MANDATORY RULES FOR GOOD MATCHING TEXT:**
1. Each effect must be self-contained, technically precise, and reflect a distinct capability or outcome.
2. Avoid vague terms like "optimize", "improve", "enhance" without specifying what is optimized/how.
3. Include explicit Network/AI/system-related technologies, protocols, or methods if implied in the problem.
4. Use exact functional terminology that would appear in patent claims.
5. ONLY OUTPUT A LIST OF JSON OBJECTS, NOTHING ELSE.

Whats SAOKE ? S-A-O-K-E decomposition methodology applied to patents:
**Subject (S):** The entity that performs the action (e.g., device, user, system).
**Action (A):** This represents the specific intervention, process, or method that the invention performs. It describes what the invention *does* to or with specific objects or systems (e.g., transmits, applies, mixes).
**Object (O):** The entity or target that the action is performed upon (e.g., signal, data, mixture).
**Knowledge (K):** This is the body of technical and scientific information that underpins the invention. It is the knowledge that is necessary to design, implement, and operate the action successfully.
**Effect (E):** This refers to the outcome, result, or consequence of the action. It describes the benefit, improvement, or new capability that the invention provides.

TASK:
Given the problem identify the root causes and output a list of effect following the SAOKE concept :

Output as a JSON list:
[
 {{"effect-name": "Short technical label",
  "description": "One-sentence precise technical outcome",
 }}
]

<<<PROBLEM>>>
{problem}
"""
        
        response = self.ask_gemini(prompt)
        result_text = response.text.replace("```json", "").replace("```", "").replace("\n", "")
        result_text = re.sub(r',\s*([}\]])', r'\1', result_text)
        
        # response = self.ask_gemini(prompt)
        # result_text = response.text
        # result_text = result_text[result_text.find('{'):result_text.find('}')+1].replace("```json", "").replace("```", "").replace("\n", "")
        # result_text = re.sub(r',\s*([}\]])', r'\1', result_text)

        try:
            effects = json.loads(result_text)
            return effects
        except json.JSONDecodeError as e:
            st.error(f"Error parsing effects: {e}")
            return []
    
    def select_mechanism_with_llm(self, effect_description: str, top_50_mechanisms: List[Dict]) -> Dict:
        """Use LLM to select best mechanism from top 50 for given effect"""
        
        # Prepare the mechanisms list for the prompt
        mechanisms_list = []
        for i, pair in enumerate(top_50_mechanisms):
            mechanism = pair['mechanism']
            mechanisms_list.append({
                "index": i,
                "technology": mechanism['technology'],
                "subject": mechanism['subject'],
                "action": mechanism['action'],
                "object": mechanism['object'],
                "knowledge": mechanism['knowledge'],
                "effect": mechanism['effect'],
            })
        
        prompt = f"""
###TASK###
Whats SAOKE ? S-A-O-K-E decomposition methodology applied to patents:
**Subject (S):** The entity that performs the action (e.g., device, user, system).
**Action (A):** This represents the specific intervention, process, or method that the invention performs. It describes what the invention *does* to or with specific objects or systems (e.g., transmits, applies, mixes).
**Object (O):** The entity or target that the action is performed upon (e.g., signal, data, mixture).
**Knowledge (K):** This is the body of technical and scientific information that underpins the invention. It is the knowledge that is necessary to design, implement, and operate the action successfully.
**Effect (E):** This refers to the outcome, result, or consequence of the action. It describes the benefit, improvement, or new capability that the invention provides.

the entire invention can be mapped as a linked set of S-A–O-K–E units. For example:
Step 1: (S₁, A₁, O₁, K₁) → E₁
Step 2: (S₂,A₂, O₂, K₂=E₁+...) → E₂
Step 3: (S₃, A₃, O₃, K₃=E₂+...) → E₃
...and so on.

Using this concept, I have identified one single effect, you will choose from a list of mechanisms which one suits the best the described effect.
Output only ONE mechanism in the same format as provided, such as :
{{"subject": "the subject of the mechanism",
"action": "...",
"object": "...",
"knowledge": "...",
"effect": "...",
"technology": "..."
}}

###List of effect and mechanisms ###
Effect: {effect_description}

Mechanisms to choose from:
{json.dumps(mechanisms_list, indent=2)}
"""
        
        response = self.ask_gemini(prompt)
        
        try:
            result_text = response.text
            # Extract JSON from response
            if "```json" in result_text:
                result_text = result_text[result_text.find("```json"):].replace("```json", "").replace("```", "").replace("\n", "")
            else:
                # Look for JSON-like structure
                start = result_text.find("{")
                end = result_text.rfind("}") + 1
                if start != -1 and end != -1:
                    result_text = result_text[start:end]
            
            result_text = re.sub(r',\s*([}\]])', r'\1', result_text)
            selected_mechanism = json.loads(result_text)
            return selected_mechanism
            
        except Exception as e:
            st.warning(f"LLM selection failed for effect, using top similarity match: {e}")
            # Fallback to highest similarity if LLM selection fails
            return top_50_mechanisms[0]['mechanism']
    
    def match_mechanisms(self, effects: List[Dict], list_mechanisms: List[Dict], model) -> List[Dict]:
        """Match effects to mechanisms using semantic similarity and LLM selection"""
        try:
            # Prepare mechanism effects list
            mechanism_effects_list = []
            mechanism_details = []
            for row in list_mechanisms:
                for m in row["mechanisms"]:
                    mechanism_effects_list.append(m["SAOKE"]["effect"])
                    mechanism_details.append((row, m))
            
            mechanism_embeddings = model.encode(mechanism_effects_list)
            effect_mechanism_pairs = []
            
            for effect_idx, effect in enumerate(effects):
                effect_description = f"{effect['effect-name']}. Description: {effect['description']}"
                effect_embedding = model.encode([effect_description])
                
                # Calculate similarities
                similarities = []
                for i, mech_embedding in enumerate(mechanism_embeddings):
                    similarity = util.cos_sim(effect_embedding[0], mech_embedding).item()
                    similarities.append((similarity, i))
                
                # Sort and get top 50
                similarities.sort(reverse=True)
                top_50 = similarities[:50]
                
                # Store top 50 for this effect
                effect_key = f"effect_{effect_idx}"
                st.session_state.top_50_mechanisms[effect_key] = []
                
                for similarity, mech_idx in top_50:
                    row, m = mechanism_details[mech_idx]
                    mechanism = m["SAOKE"].copy()
                    mechanism["technology"] = row["technology"]
                    mechanism["similarity"] = similarity
                    st.session_state.top_50_mechanisms[effect_key].append({
                        "effect": effect_description,
                        "mechanism": mechanism
                    })
                
                # Use LLM to select best mechanism from top 50
                with st.spinner(f"🤖 LLM selecting best mechanism for effect {effect_idx + 1}..."):
                    selected_mechanism = self.select_mechanism_with_llm(
                        effect_description, 
                        st.session_state.top_50_mechanisms[effect_key]
                    )
                
                # Find the index of the selected mechanism in our top 50 list
                selected_index = 0  # Default to first if not found
                for i, pair in enumerate(st.session_state.top_50_mechanisms[effect_key]):
                    if (pair['mechanism']['technology'] == selected_mechanism.get('technology', '') and
                        pair['mechanism']['effect'] == selected_mechanism.get('effect', '')):
                        selected_index = i
                        break
                
                # Initialize current index to LLM selection
                st.session_state.mechanism_indices[effect_key] = selected_index
                # Store original LLM-selected index for reset functionality
                st.session_state.original_llm_indices[effect_key] = selected_index
                
                # Add LLM-selected mechanism to pairs
                selected_pair = st.session_state.top_50_mechanisms[effect_key][selected_index]
                effect_mechanism_pairs.append(selected_pair)
            
            # Store original mechanisms for reset functionality
            st.session_state.original_mechanisms = effect_mechanism_pairs.copy()
            
            return effect_mechanism_pairs
        except Exception as e:
            st.error(f"Error matching mechanisms: {e}")
            return []
    
    def generate_solution(self, problem: str, mechanisms: List[Dict]) -> str:
        """Generate solution using Gemini"""
        prompt = f"""
TASK
Using SAOKE concept:
Whats SAOKE ? S-A-O-K-E decomposition methodology applied to patents: Subject (S): The entity that performs the action (e.g., device, user, system). Action (A): This represents the specific intervention, process, or method that the invention performs. It describes what the invention does to or with specific objects or systems (e.g., transmits, applies, mixes). Object (O): The entity or target that the action is performed upon (e.g., signal, data, mixture). Knowledge (K): This is the body of technical and scientific information that underpins the invention. It is the knowledge that is necessary to design, implement, and operate the action successfully. Effect (E): This refers to the outcome, result, or consequence of the action. It describes the benefit, improvement, or new capability that the invention provides.

The entire invention can be mapped as a linked set of S-A–O-K–E units. For example: Step 1: (S₁, A₁, O₁, K₁) → E₁ Step 2: (S₂,A₂, O₂, K₂=E₁+...) → E₂ Step 3: (S₃, A₃, O₃, K₃=E₂+...) → E₃ ...and so on.

From a problem I've extracted all the effects in order to find from a list of mechanism which one would be suited the best to cover each effect and finally solve the initial problem.

Using the list of mechanism identified craft a solution which would use all of the mechanism in order to solve the initial problem.

Structure the solution following this plan : 
1. Scenario:
    -State the scenario in which we want to tailor a solution in a short sentence/
    
2. Context and goals:
    -State the current state, the problem, and the high-level objective in two to four sentences.
    -Define the success signal in plain language and why it matters to stakeholders.

3. Requirements and criteria:
    -Functional requirements (FR): enumerate capabilities and behaviors.
    -Non-functional requirements (NFR): security, performance, latency, availability, compliance, UX constraints.
    -Acceptance criteria: binary, testable statements tied to FR/NFR.

CONTEXT INFORMATION
Problem:
{problem}

List of mechanism identified per effect identified for the problem:
{mechanisms}
"""
        
        response = self.ask_gemini(prompt)
        return response.text
    
    def render_progress_bar(self, current_step: int, total_steps: int = 4):
        """Render progress bar"""
        progress = current_step / total_steps
        st.progress(progress)
        
        cols = st.columns(total_steps)
        steps = ["Problem Input", "Effects Extraction", "Mechanism Matching", "Solution Generation"]
        
        for i, (col, step_name) in enumerate(zip(cols, steps)):
            with col:
                if i + 1 < current_step:
                    st.markdown(f"✅ **{step_name}**")
                elif i + 1 == current_step:
                    st.markdown(f"🔄 **{step_name}**")
                else:
                    st.markdown(f"⏳ {step_name}")
    
    def render_step1(self):
        """Render Step 1: Problem Input"""
        st.markdown('<div class="step-header">Step 1: Problem Input</div>', unsafe_allow_html=True)
        
        st.markdown("**Enter your technical problem below:**")
        problem = st.text_area(
            "Problem Description",
            value=st.session_state.problem,
            height=200,
            placeholder="Describe your technical challenge in detail...",
            label_visibility="collapsed"
        )
        
        col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
        
        # Show back to effects button if we have effects
        if st.session_state.effects:
            with col1:
                if st.button("⬅️ Back to Effects", use_container_width=True):
                    st.session_state.step = 2
                    st.rerun()
        
        with col3 if st.session_state.effects else col2:
            if st.button("🔍 Analyze Problem", type="primary", use_container_width=True):
                if problem.strip():
                    st.session_state.problem = problem
                    # Reset subsequent steps when changing problem
                    st.session_state.effects = []
                    st.session_state.mechanisms = []
                    st.session_state.solution = ""
                    st.session_state.top_50_mechanisms = {}
                    st.session_state.mechanism_indices = {}
                    st.session_state.original_mechanisms = []
                    st.session_state.original_llm_indices = {}
                    st.session_state.step = 2
                    st.rerun()
                else:
                    st.error("Please enter a problem description")
    
    def render_step2(self):
        """Render Step 2: Effects Extraction"""
        st.markdown('<div class="step-header">Step 2: Effects Extraction</div>', unsafe_allow_html=True)
        
        if not st.session_state.effects:
            with st.spinner("🔬 Extracting effects from your problem..."):
                model = self.load_models()
                if model:
                    effects = self.extract_effects(st.session_state.problem)
                    st.session_state.effects = effects
                    st.rerun()
        
        if st.session_state.effects:
            st.markdown("**Extracted Effects:**")
            
            for i, effect in enumerate(st.session_state.effects, 1):
                with st.container():
                    st.markdown(f"""
                    <div class="effect-card">
                        <h4>Effect {i}: {effect['effect-name']}</h4>
                        <p><strong>Description:</strong> {effect['description']}</p>
                    </div>
                    """, unsafe_allow_html=True)
            
            col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
            
            with col1:
                if st.button("⬅️ Back to Problem", use_container_width=True):
                    st.session_state.step = 1
                    st.rerun()
            
            with col2:
                if st.button("🔄 Re-generate Effects", use_container_width=True):
                    st.session_state.effects = []
                    # Reset subsequent steps
                    st.session_state.mechanisms = []
                    st.session_state.solution = ""
                    st.session_state.top_50_mechanisms = {}
                    st.session_state.mechanism_indices = {}
                    st.session_state.original_mechanisms = []
                    st.session_state.original_llm_indices = {}
                    st.rerun()
            
            with col4:
                if st.button("➡️ Continue to Matching", type="primary", use_container_width=True):
                    st.session_state.step = 3
                    st.rerun()
    
    def render_step3(self):
        """Render Step 3: Mechanism Matching"""
        st.markdown('<div class="step-header">Step 3: Mechanism Matching</div>', unsafe_allow_html=True)
        
        if not st.session_state.mechanisms:
            with st.spinner("🔗 Matching effects to mechanisms..."):
                model = self.load_models()
                list_mechanisms = self.load_data()
                
                if model and list_mechanisms:
                    mechanisms = self.match_mechanisms(st.session_state.effects, list_mechanisms, model)
                    st.session_state.mechanisms = mechanisms
                    st.rerun()
        
        if st.session_state.mechanisms:
            st.markdown("**Effect-Mechanism Pairs:**")
            
            for i, pair in enumerate(st.session_state.mechanisms, 1):
                effect_key = f"effect_{i-1}"
                
                with st.container():
                    col1, col2 = st.columns([4, 2])
                    
                    with col1:
                        # Show similarity score if available
                        similarity_score = pair['mechanism'].get('similarity', 0)
                        st.markdown(f"""
                        <div class="mechanism-card">
                            <h4>Pair {i} - 🤖 LLM Selected (Similarity: {similarity_score:.3f})</h4>
                            <p><strong>Effect:</strong> {pair['effect']}</p>
                            <p><strong>Technology:</strong> {pair['mechanism']['technology']}</p>
                            <p><strong>Subject:</strong> {pair['mechanism']['subject']}</p>
                            <p><strong>Action:</strong> {pair['mechanism']['action']}</p>
                            <p><strong>Object:</strong> {pair['mechanism']['object']}</p>
                            <p><strong>Knowledge:</strong> {pair['mechanism']['knowledge']}</p>
                            <p><strong>Effect:</strong> {pair['mechanism']['effect']}</p>
                        </div>
                        """, unsafe_allow_html=True)
                    
                    with col2:
                        # Create two columns for navigation buttons
                        nav_col1, nav_col2 = st.columns(2)
                        
                        with nav_col1:
                            # Previous mechanism button
                            if st.button(f"⏮️ Prev #{i}", key=f"prev_{i}", use_container_width=True):
                                if effect_key in st.session_state.mechanism_indices and effect_key in st.session_state.top_50_mechanisms:
                                    current_idx = st.session_state.mechanism_indices[effect_key]
                                    max_idx = len(st.session_state.top_50_mechanisms[effect_key]) - 1
                                    
                                    # Circular navigation: if at first (0), go to last
                                    if current_idx <= 0:
                                        st.session_state.mechanism_indices[effect_key] = max_idx
                                    else:
                                        st.session_state.mechanism_indices[effect_key] = current_idx - 1
                                    
                                    new_idx = st.session_state.mechanism_indices[effect_key]
                                    st.session_state.mechanisms[i-1] = st.session_state.top_50_mechanisms[effect_key][new_idx]
                                    # Reset solution when mechanism changes
                                    st.session_state.solution = ""
                                    st.rerun()
                        
                        with nav_col2:
                            # Next mechanism button
                            if st.button(f"⏭️ Next #{i}", key=f"next_{i}", use_container_width=True):
                                if effect_key in st.session_state.mechanism_indices and effect_key in st.session_state.top_50_mechanisms:
                                    current_idx = st.session_state.mechanism_indices[effect_key]
                                    max_idx = len(st.session_state.top_50_mechanisms[effect_key]) - 1
                                    
                                    # Circular navigation: if at last, go to first (0)
                                    if current_idx >= max_idx:
                                        st.session_state.mechanism_indices[effect_key] = 0
                                    else:
                                        st.session_state.mechanism_indices[effect_key] = current_idx + 1
                                    
                                    new_idx = st.session_state.mechanism_indices[effect_key]
                                    st.session_state.mechanisms[i-1] = st.session_state.top_50_mechanisms[effect_key][new_idx]
                                    # Reset solution when mechanism changes
                                    st.session_state.solution = ""
                                    st.rerun()
                        
                        # Reset button - now resets to LLM-selected mechanism
                        if st.button(f"🔄 Reset #{i}", key=f"reset_{i}", use_container_width=True):
                            if effect_key in st.session_state.original_llm_indices:
                                # Reset to original LLM-selected index instead of 0
                                original_idx = st.session_state.original_llm_indices[effect_key]
                                st.session_state.mechanism_indices[effect_key] = original_idx
                                if effect_key in st.session_state.top_50_mechanisms:
                                    st.session_state.mechanisms[i-1] = st.session_state.top_50_mechanisms[effect_key][original_idx]
                                    # Reset solution when mechanism changes
                                    st.session_state.solution = ""
                                    st.rerun()
                        
                        # Show current position
                        if effect_key in st.session_state.mechanism_indices and effect_key in st.session_state.top_50_mechanisms:
                            current = st.session_state.mechanism_indices[effect_key] + 1
                            total = len(st.session_state.top_50_mechanisms[effect_key])
                            st.caption(f"📍 {current}/{total}")
                            
                            # Show navigation hint
                            if total > 1:
                                st.caption("🔄 Circular navigation enabled")
            
            col1, col2, col3 = st.columns([1, 1, 1])
            
            with col1:
                if st.button("⬅️ Back to Effects", use_container_width=True):
                    st.session_state.step = 2
                    st.rerun()
            
            with col3:
                if st.button("🚀 Generate Solution", type="primary", use_container_width=True):
                    st.session_state.step = 4
                    st.rerun()
    
    def render_step4(self):
        """Render Step 4: Solution Generation"""
        st.markdown('<div class="step-header">Step 4: Solution Generation</div>', unsafe_allow_html=True)
        
        if not st.session_state.solution:
            with st.spinner("✨ Generating your solution..."):
                solution = self.generate_solution(st.session_state.problem, st.session_state.mechanisms)
                st.session_state.solution = solution
                st.rerun()
        
        if st.session_state.solution:
            st.markdown("**Generated Solution:**")
            
            # Render the solution as markdown with custom styling
            st.markdown(
                f"""
                <div class="solution-container">
                    {st.session_state.solution}
                </div>
                """,
                unsafe_allow_html=True
            )
            
            # Also render as proper markdown for better formatting
            with st.expander("📖 View Formatted Solution", expanded=False):
                st.markdown(st.session_state.solution)
            
            col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
            
            with col1:
                if st.button("⬅️ Back to Mechanisms", use_container_width=True):
                    st.session_state.step = 3
                    st.rerun()
            
            with col2:
                if st.button("🔄 Re-generate Solution", use_container_width=True):
                    st.session_state.solution = ""
                    st.rerun()
            
            with col3:
                if st.button("🔄 Start New Analysis", use_container_width=True):
                    # Reset all session state
                    for key in ['step', 'problem', 'effects', 'mechanisms', 'solution', 'top_50_mechanisms', 'mechanism_indices', 'original_mechanisms', 'original_llm_indices']:
                        if key in st.session_state:
                            del st.session_state[key]
                    st.session_state.step = 1
                    st.rerun()
            
            with col4:
                # Download solution as text file
                st.download_button(
                    label="📥 Download Solution",
                    data=st.session_state.solution,
                    file_name="saoke_solution.txt",
                    mime="text/plain",
                    use_container_width=True
                )
    
    def run(self):
        """Main application runner"""
        # Header
        st.markdown('<h1 class="main-header">🔬 SAOKE Problem Solver</h1>', unsafe_allow_html=True)
        st.markdown("---")
        
        # Progress bar
        self.render_progress_bar(st.session_state.step)
        st.markdown("---")
        
        # Check data availability
        if not st.session_state.data_loaded:
            try:
                list_mechanisms = self.load_data()
                if list_mechanisms:
                    st.session_state.data_loaded = True
                    st.success("✅ Data loaded successfully!")
                else:
                    st.error("❌ Failed to load SAOKE_technologies.xlsx. Please ensure the file exists.")
                    return
            except Exception as e:
                st.error(f"❌ Error loading data: {e}")
                return
        
        # Render current step
        if st.session_state.step == 1:
            self.render_step1()
        elif st.session_state.step == 2:
            self.render_step2()
        elif st.session_state.step == 3:
            self.render_step3()
        elif st.session_state.step == 4:
            self.render_step4()
        
        # Sidebar with current state
        with st.sidebar:
            st.markdown("## Current State")
            st.markdown(f"**Current Step:** {st.session_state.step}/4")
            
            if st.session_state.problem:
                st.markdown("**Problem:** ✅ Entered")
            if st.session_state.effects:
                st.markdown(f"**Effects:** ✅ {len(st.session_state.effects)} extracted")
            if st.session_state.mechanisms:
                st.markdown(f"**Mechanisms:** ✅ {len(st.session_state.mechanisms)} matched")
            if st.session_state.solution:
                st.markdown("**Solution:** ✅ Generated")

def main():
    app = SAOKEWebApp()
    app.run()

if __name__ == "__main__":
    main()