import streamlit as st
import anthropic, openai, base64, cv2, glob, json, math, os, pytz, random, re, requests, textract, time, zipfile
import plotly.graph_objects as go
import streamlit.components.v1 as components
from datetime import datetime
from audio_recorder_streamlit import audio_recorder
from bs4 import BeautifulSoup
from collections import defaultdict, deque, Counter
from dotenv import load_dotenv
from gradio_client import Client
from huggingface_hub import InferenceClient
from io import BytesIO
from PIL import Image
from PyPDF2 import PdfReader
from urllib.parse import quote
from xml.etree import ElementTree as ET
from openai import OpenAI
import extra_streamlit_components as stx
from streamlit.runtime.scriptrunner import get_script_run_ctx
import asyncio
import edge_tts

# 🎯 1. Core Configuration & Setup
st.set_page_config(
    page_title="🚲TalkingAIResearcher🏆",
    page_icon="🚲🏆",
    layout="wide",
    initial_sidebar_state="auto",
    menu_items={
        'Get Help': 'https://huggingface.co/awacke1',
        'Report a bug': 'https://huggingface.co/spaces/awacke1',
        'About': "🚲TalkingAIResearcher🏆"
    }
)
load_dotenv()

# Add available English voices for Edge TTS
EDGE_TTS_VOICES = [
    "en-US-AriaNeural",  # Default voice
    "en-US-GuyNeural", 
    "en-US-JennyNeural",
    "en-GB-SoniaNeural",
    "en-GB-RyanNeural",
    "en-AU-NatashaNeural",
    "en-AU-WilliamNeural",
    "en-CA-ClaraNeural",
    "en-CA-LiamNeural"
]

# Initialize session state variables
if 'tts_voice' not in st.session_state:
    st.session_state['tts_voice'] = EDGE_TTS_VOICES[0]  # Default voice
if 'audio_format' not in st.session_state:
    st.session_state['audio_format'] = 'mp3'  # 🆕 Default audio format

# 🔑 2. API Setup & Clients
openai_api_key = os.getenv('OPENAI_API_KEY', "")
anthropic_key = os.getenv('ANTHROPIC_API_KEY_3', "")
xai_key = os.getenv('xai',"")
if 'OPENAI_API_KEY' in st.secrets:
    openai_api_key = st.secrets['OPENAI_API_KEY']
if 'ANTHROPIC_API_KEY' in st.secrets:
    anthropic_key = st.secrets["ANTHROPIC_API_KEY"]

openai.api_key = openai_api_key
claude_client = anthropic.Anthropic(api_key=anthropic_key)
openai_client = OpenAI(api_key=openai.api_key, organization=os.getenv('OPENAI_ORG_ID'))
HF_KEY = os.getenv('HF_KEY')
API_URL = os.getenv('API_URL')

# 📝 3. Session State Management
if 'transcript_history' not in st.session_state:
    st.session_state['transcript_history'] = []
if 'chat_history' not in st.session_state:
    st.session_state['chat_history'] = []
if 'openai_model' not in st.session_state:
    st.session_state['openai_model'] = "gpt-4o-2024-05-13"
if 'messages' not in st.session_state:
    st.session_state['messages'] = []
if 'last_voice_input' not in st.session_state:
    st.session_state['last_voice_input'] = ""
if 'editing_file' not in st.session_state:
    st.session_state['editing_file'] = None
if 'edit_new_name' not in st.session_state:
    st.session_state['edit_new_name'] = ""
if 'edit_new_content' not in st.session_state:
    st.session_state['edit_new_content'] = ""
if 'viewing_prefix' not in st.session_state:
    st.session_state['viewing_prefix'] = None
if 'should_rerun' not in st.session_state:
    st.session_state['should_rerun'] = False
if 'old_val' not in st.session_state:
    st.session_state['old_val'] = None
if 'last_query' not in st.session_state:
    st.session_state['last_query'] = ""  # 🆕 Store the last query for zip naming

# 🎨 4. Custom CSS
st.markdown("""
<style>
    .main { background: linear-gradient(to right, #1a1a1a, #2d2d2d); color: #fff; }
    .stMarkdown { font-family: 'Helvetica Neue', sans-serif; }
    .stButton>button {
        margin-right: 0.5rem;
    }
</style>
""", unsafe_allow_html=True)

FILE_EMOJIS = {
    "md": "📝",
    "mp3": "🎵",
    "wav": "🔊"  # 🆕 Add emoji for WAV
}

# 🧠 5. High-Information Content Extraction
def get_high_info_terms(text: str, top_n=10) -> list:
    """Extract high-information terms from text, including key phrases."""
    stop_words = set([
        'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with',
        'by', 'from', 'up', 'about', 'into', 'over', 'after', 'is', 'are', 'was', 'were',
        'be', 'been', 'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would',
        'should', 'could', 'might', 'must', 'shall', 'can', 'may', 'this', 'that', 'these',
        'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they', 'what', 'which', 'who',
        'when', 'where', 'why', 'how', 'all', 'any', 'both', 'each', 'few', 'more', 'most',
        'other', 'some', 'such', 'than', 'too', 'very', 'just', 'there'
    ])

    key_phrases = [
        'artificial intelligence', 'machine learning', 'deep learning', 'neural network',
        'personal assistant', 'natural language', 'computer vision', 'data science',
        'reinforcement learning', 'knowledge graph', 'semantic search', 'time series',
        'large language model', 'transformer model', 'attention mechanism',
        'autonomous system', 'edge computing', 'quantum computing', 'blockchain technology',
        'cognitive science', 'human computer', 'decision making', 'arxiv search',
        'research paper', 'scientific study', 'empirical analysis'
    ]

    # Extract bi-grams and uni-grams
    words = re.findall(r'\b\w+(?:-\w+)*\b', text.lower())
    bi_grams = [' '.join(pair) for pair in zip(words, words[1:])]
    combined = words + bi_grams

    # Filter out stop words and short words
    filtered = [
        term for term in combined
        if term not in stop_words
        and len(term.split()) <= 2  # Limit to uni-grams and bi-grams
        and any(c.isalpha() for c in term)
    ]

    # Count frequencies
    counter = Counter(filtered)
    most_common = [term for term, freq in counter.most_common(top_n)]
    return most_common

def clean_text_for_filename(text: str) -> str:
    """Remove punctuation and short filler words, return a compact string."""
    text = text.lower()
    text = re.sub(r'[^\w\s-]', '', text)
    words = text.split()
    stop_short = set(['the','and','for','with','this','that','from','just','very','then','been','only','also','about'])
    filtered = [w for w in words if len(w)>3 and w not in stop_short]
    return '_'.join(filtered)[:200]

# 📁 6. File Operations
def generate_filename(prompt, response, file_type="md"):
    """
    Generate filename with meaningful terms and short dense clips from prompt & response.
    The filename should be about 150 chars total, include high-info terms, and a clipped snippet.
    """
    prefix = datetime.now().strftime("%y%m_%H%M") + "_"
    combined = (prompt + " " + response).strip()
    info_terms = get_high_info_terms(combined, top_n=10)
    
    # Include a short snippet from prompt and response
    snippet = (prompt[:100] + " " + response[:100]).strip()
    snippet_cleaned = clean_text_for_filename(snippet)
    
    # Combine info terms and snippet
    name_parts = info_terms + [snippet_cleaned]
    full_name = '_'.join(name_parts)

    # Trim to ~150 chars
    if len(full_name) > 150:
        full_name = full_name[:150]
    
    filename = f"{prefix}{full_name}.{file_type}"
    return filename

def create_file(prompt, response, file_type="md"):
    """Create file with intelligent naming"""
    filename = generate_filename(prompt.strip(), response.strip(), file_type)
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(prompt + "\n\n" + response)
    return filename

def get_download_link(file, file_type="zip"):
    """Generate download link for file"""
    with open(file, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    if file_type == "zip":
        return f'<a href="data:application/zip;base64,{b64}" download="{os.path.basename(file)}">📂 Download {os.path.basename(file)}</a>'
    elif file_type == "mp3":
        return f'<a href="data:audio/mpeg;base64,{b64}" download="{os.path.basename(file)}">🎵 Download {os.path.basename(file)}</a>'
    elif file_type == "wav":
        return f'<a href="data:audio/wav;base64,{b64}" download="{os.path.basename(file)}">🔊 Download {os.path.basename(file)}</a>'  # 🆕 WAV download link
    elif file_type == "md":
        return f'<a href="data:text/markdown;base64,{b64}" download="{os.path.basename(file)}">📝 Download {os.path.basename(file)}</a>'
    else:
        return f'<a href="data:application/octet-stream;base64,{b64}" download="{os.path.basename(file)}">Download {os.path.basename(file)}</a>'

# 🔊 7. Audio Processing
def clean_for_speech(text: str) -> str:
    """Clean text for speech synthesis"""
    text = text.replace("\n", " ")
    text = text.replace("</s>", " ")
    text = text.replace("#", "")
    text = re.sub(r"\(https?:\/\/[^\)]+\)", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

@st.cache_resource
def speech_synthesis_html(result):
    """Create HTML for speech synthesis"""
    html_code = f"""
    <html><body>
    <script>
    var msg = new SpeechSynthesisUtterance("{result.replace('"', '')}");
    window.speechSynthesis.speak(msg);
    </script>
    </body></html>
    """
    components.html(html_code, height=0)

async def edge_tts_generate_audio(text, voice="en-US-AriaNeural", rate=0, pitch=0, file_format="mp3"):
    """Generate audio using Edge TTS"""
    text = clean_for_speech(text)
    if not text.strip():
        return None
    rate_str = f"{rate:+d}%"
    pitch_str = f"{pitch:+d}Hz"
    communicate = edge_tts.Communicate(text, voice, rate=rate_str, pitch=pitch_str)
    out_fn = generate_filename(text, text, file_type=file_format)
    await communicate.save(out_fn)
    return out_fn

def speak_with_edge_tts(text, voice="en-US-AriaNeural", rate=0, pitch=0, file_format="mp3"):
    """Wrapper for edge TTS generation"""
    return asyncio.run(edge_tts_generate_audio(text, voice, rate, pitch, file_format))

def play_and_download_audio(file_path, file_type="mp3"):
    """Play and provide download link for audio"""
    if file_path and os.path.exists(file_path):
        if file_type == "mp3":
            st.audio(file_path)
        elif file_type == "wav":
            st.audio(file_path)
        dl_link = get_download_link(file_path, file_type=file_type)
        st.markdown(dl_link, unsafe_allow_html=True)

# 🎬 8. Media Processing
def process_image(image_path, user_prompt):
    """Process image with GPT-4V"""
    with open(image_path, "rb") as imgf:
        image_data = imgf.read()
    b64img = base64.b64encode(image_data).decode("utf-8")
    resp = openai_client.chat.completions.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": [
                {"type": "text", "text": user_prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64img}"}}
            ]}
        ],
        temperature=0.0,
    )
    return resp.choices[0].message.content

def process_audio_file(audio_path):
    """Process audio with Whisper"""
    with open(audio_path, "rb") as f:
        transcription = openai_client.audio.transcriptions.create(model="whisper-1", file=f)
    st.session_state.messages.append({"role": "user", "content": transcription.text})
    return transcription.text

def process_video(video_path, seconds_per_frame=1):
    """Extract frames from video"""
    vid = cv2.VideoCapture(video_path)
    total = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vid.get(cv2.CAP_PROP_FPS)
    skip = int(fps*seconds_per_frame)
    frames_b64 = []
    for i in range(0, total, skip):
        vid.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = vid.read()
        if not ret: 
            break
        _, buf = cv2.imencode(".jpg", frame)
        frames_b64.append(base64.b64encode(buf).decode("utf-8"))
    vid.release()
    return frames_b64

def process_video_with_gpt(video_path, prompt):
    """Analyze video frames with GPT-4V"""
    frames = process_video(video_path)
    resp = openai_client.chat.completions.create(
        model=st.session_state["openai_model"],
        messages=[
            {"role":"system","content":"Analyze video frames."},
            {"role":"user","content":[
                {"type":"text","text":prompt},
                *[{"type":"image_url","image_url":{"url":f"data:image/jpeg;base64,{fr}"}} for fr in frames]
            ]}
        ]
    )
    return resp.choices[0].message.content

# 🤖 9. AI Model Integration

def save_full_transcript(query, text):
    """Save full transcript of Arxiv results as a file."""
    create_file(query, text, "md")

def parse_arxiv_refs(ref_text: str):
    """
    Parse papers by finding lines with two pipe characters as title lines.
    Returns list of paper dictionaries with audio files.
    """
    if not ref_text:
        return []

    results = []
    current_paper = {}
    lines = ref_text.split('\n')
    
    for i, line in enumerate(lines):
        # Check if this is a title line (contains exactly 2 pipe characters)
        if line.count('|') == 2:
            # If we have a previous paper, add it to results
            if current_paper:
                results.append(current_paper)
                if len(results) >= 20:  # Limit to 20 papers
                    break
            
            # Parse new paper header
            try:
                # Remove ** and split by |
                header_parts = line.strip('* ').split('|')
                date = header_parts[0].strip()
                title = header_parts[1].strip()
                # Extract arXiv URL if present
                url_match = re.search(r'(https://arxiv.org/\S+)', line)
                url = url_match.group(1) if url_match else f"paper_{len(results)}"
                
                current_paper = {
                    'date': date,
                    'title': title,
                    'url': url,
                    'authors': '',
                    'summary': '',
                    'content_start': i + 1  # Track where content begins
                }
            except Exception as e:
                st.warning(f"Error parsing paper header: {str(e)}")
                current_paper = {}
                continue
        
        # If we have a current paper and this isn't a title line, add to content
        elif current_paper:
            if not current_paper['authors']:  # First line after title is authors
                current_paper['authors'] = line.strip('* ')
            else:  # Rest is summary
                if current_paper['summary']:
                    current_paper['summary'] += ' ' + line.strip()
                else:
                    current_paper['summary'] = line.strip()
    
    # Don't forget the last paper
    if current_paper:
        results.append(current_paper)
    
    return results[:20]  # Ensure we return maximum 20 papers

def create_paper_audio_files(papers, input_question):
    """
    Create audio files for each paper's content and add file paths to paper dict.
    Also, display each audio as it's generated.
    """
    # Collect all content for combined summary
    combined_titles = []

    for paper in papers:
        try:
            # Generate audio for full content only
            full_text = f"{paper['title']} by {paper['authors']}. {paper['summary']}"
            full_text = clean_for_speech(full_text)
            # Determine file format based on user selection
            file_format = st.session_state['audio_format']
            full_file = speak_with_edge_tts(full_text, voice=st.session_state['tts_voice'], file_format=file_format)
            paper['full_audio'] = full_file

            # Display the audio immediately after generation
            st.write(f"### {FILE_EMOJIS.get(file_format, '')} {os.path.basename(full_file)}")
            play_and_download_audio(full_file, file_type=file_format)
            
            combined_titles.append(paper['title'])
        
        except Exception as e:
            st.warning(f"Error generating audio for paper {paper['title']}: {str(e)}")
            paper['full_audio'] = None

    # After all individual audios, create a combined summary audio
    if combined_titles:
        combined_text = f"Here are the titles of the papers related to your query: {'; '.join(combined_titles)}. Your original question was: {input_question}"
        file_format = st.session_state['audio_format']
        combined_file = speak_with_edge_tts(combined_text, voice=st.session_state['tts_voice'], file_format=file_format)
        st.write(f"### {FILE_EMOJIS.get(file_format, '')} Combined Summary Audio")
        play_and_download_audio(combined_file, file_type=file_format)
        papers.append({'title': 'Combined Summary', 'full_audio': combined_file})

def display_papers(papers):
    """
    Display papers with their audio controls using URLs as unique keys.
    """
    st.write("## Research Papers")
    papercount=0
    for idx, paper in enumerate(papers):
        papercount = papercount + 1
        if (papercount<=20):
            with st.expander(f"{papercount}. 📄 {paper['title']}", expanded=True):
                st.markdown(f"**{paper['date']} | {paper['title']} | ⬇️**")
                st.markdown(f"*{paper['authors']}*")
                st.markdown(paper['summary'])
                
                # Single audio control for full content
                if paper.get('full_audio'):
                    st.write("📚 Paper Audio")
                    file_ext = os.path.splitext(paper['full_audio'])[1].lower().strip('.')
                    if file_ext == "mp3":
                        st.audio(paper['full_audio'])
                    elif file_ext == "wav":
                        st.audio(paper['full_audio'])

def perform_ai_lookup(q, vocal_summary=True, extended_refs=False, 
                     titles_summary=True, full_audio=False):
    """Perform Arxiv search with audio generation per paper."""
    start = time.time()

    # Query the HF RAG pipeline
    client = Client("awacke1/Arxiv-Paper-Search-And-QA-RAG-Pattern")
    refs = client.predict(q, 20, "Semantic Search", 
                         "mistralai/Mixtral-8x7B-Instruct-v0.1",
                         api_name="/update_with_rag_md")[0]
    r2 = client.predict(q, "mistralai/Mixtral-8x7B-Instruct-v0.1", 
                       True, api_name="/ask_llm")

    # Combine for final text output
    result = f"### 🔎 {q}\n\n{r2}\n\n{refs}"
    st.markdown(result)

    # Parse and process papers
    papers = parse_arxiv_refs(refs)
    if papers:
        create_paper_audio_files(papers, input_question=q)
        display_papers(papers)
    else:
        st.warning("No papers found in the response.")

    elapsed = time.time()-start
    st.write(f"**Total Elapsed:** {elapsed:.2f} s")

    # Save full transcript
    create_file(q, result, "md")
    return result

def process_with_gpt(text):
    """Process text with GPT-4"""
    if not text: 
        return
    st.session_state.messages.append({"role":"user","content":text})
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        c = openai_client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=st.session_state.messages,
            stream=False
        )
        ans = c.choices[0].message.content
        st.write("GPT-4o: " + ans)
        create_file(text, ans, "md")
        st.session_state.messages.append({"role":"assistant","content":ans})
    return ans

def process_with_claude(text):
    """Process text with Claude"""
    if not text: 
        return
    with st.chat_message("user"):
        st.markdown(text)
    with st.chat_message("assistant"):
        r = claude_client.messages.create(
            model="claude-3-sonnet-20240229",
            max_tokens=1000,
            messages=[{"role":"user","content":text}]
        )
        ans = r.content[0].text
        st.write("Claude-3.5: " + ans)
        create_file(text, ans, "md")
        st.session_state.chat_history.append({"user":text,"claude":ans})
    return ans

# 📂 10. File Management
def create_zip_of_files(md_files, mp3_files, wav_files, input_question):
    """Create zip with intelligent naming based on top 10 common words."""
    # Exclude 'readme.md'
    md_files = [f for f in md_files if os.path.basename(f).lower() != 'readme.md']
    all_files = md_files + mp3_files + wav_files
    if not all_files:
        return None

    # Collect content for high-info term extraction
    all_content = []
    for f in all_files:
        if f.endswith('.md'):
            with open(f, 'r', encoding='utf-8') as file:
                all_content.append(file.read())
        elif f.endswith('.mp3') or f.endswith('.wav'):
            # Replace underscores with spaces and extract basename without extension
            basename = os.path.splitext(os.path.basename(f))[0]
            words = basename.replace('_', ' ')
            all_content.append(words)
    
    # Include the input question
    all_content.append(input_question)
    
    combined_content = " ".join(all_content)
    info_terms = get_high_info_terms(combined_content, top_n=10)
    
    timestamp = datetime.now().strftime("%y%m_%H%M")
    name_text = '_'.join(term.replace(' ', '-') for term in info_terms[:10])
    zip_name = f"{timestamp}_{name_text}.zip"
    
    with zipfile.ZipFile(zip_name,'w') as z:
        for f in all_files:
            z.write(f)
    
    return zip_name

def load_files_for_sidebar():
    """Load and group files for sidebar display based on first 9 characters of filename"""
    md_files = glob.glob("*.md")
    mp3_files = glob.glob("*.mp3")
    wav_files = glob.glob("*.wav")

    md_files = [f for f in md_files if os.path.basename(f).lower() != 'readme.md']
    all_files = md_files + mp3_files + wav_files

    groups = defaultdict(list)
    for f in all_files:
        # Get first 9 characters of filename (timestamp) as group name
        basename = os.path.basename(f)
        group_name = basename[:9] if len(basename) >= 9 else 'Other'
        groups[group_name].append(f)

    # Sort groups based on latest file modification time
    sorted_groups = sorted(groups.items(), key=lambda x: max(os.path.getmtime(f) for f in x[1]), reverse=True)
    return sorted_groups

def extract_keywords_from_md(files):
    """Extract keywords from markdown files"""
    text = ""
    for f in files:
        if f.endswith(".md"):
            c = open(f,'r',encoding='utf-8').read()
            text += " " + c
    return get_high_info_terms(text, top_n=5)

def display_file_manager_sidebar(groups_sorted):
    """Display file manager in sidebar with timestamp-based groups"""
    st.sidebar.title("🎵 Audio & Docs Manager")

    all_md = []
    all_mp3 = []
    all_wav = []
    for group_name, files in groups_sorted:
        for f in files:
            if f.endswith(".md"):
                all_md.append(f)
            elif f.endswith(".mp3"):
                all_mp3.append(f)
            elif f.endswith(".wav"):
                all_wav.append(f)

    top_bar = st.sidebar.columns(4)
    with top_bar[0]:
        if st.button("🗑 DelAllMD"):
            for f in all_md:
                os.remove(f)
            st.session_state.should_rerun = True
    with top_bar[1]:
        if st.button("🗑 DelAllMP3"):
            for f in all_mp3:
                os.remove(f)
            st.session_state.should_rerun = True
    with top_bar[2]:
        if st.button("🗑 DelAllWAV"):
            for f in all_wav:
                os.remove(f)
            st.session_state.should_rerun = True
    with top_bar[3]:
        if st.button("⬇️ ZipAll"):
            zip_name = create_zip_of_files(all_md, all_mp3, all_wav, input_question=st.session_state.get('last_query', ''))
            if zip_name:
                st.sidebar.markdown(get_download_link(zip_name, file_type="zip"), unsafe_allow_html=True)

    for group_name, files in groups_sorted:
        timestamp_dt = datetime.strptime(group_name, "%y%m_%H%M") if len(group_name) == 9 else None
        group_label = timestamp_dt.strftime("%Y-%m-%d %H:%M") if timestamp_dt else group_name
        
        with st.sidebar.expander(f"📁 {group_label} ({len(files)})", expanded=True):
            c1,c2 = st.columns(2)
            with c1:
                if st.button("👀ViewGrp", key="view_group_"+group_name):
                    st.session_state.viewing_prefix = group_name
            with c2:
                if st.button("🗑DelGrp", key="del_group_"+group_name):
                    for f in files:
                        os.remove(f)
                    st.success(f"Deleted group {group_name}!")
                    st.session_state.should_rerun = True

            for f in files:
                fname = os.path.basename(f)
                ext = os.path.splitext(fname)[1].lower()
                emoji = FILE_EMOJIS.get(ext.strip('.'), '')
                ctime = datetime.fromtimestamp(os.path.getmtime(f)).strftime("%H:%M:%S")
                st.write(f"{emoji} **{fname}** - {ctime}")

# 🎯 11. Main Application
def main():
    st.sidebar.markdown("### 🚲BikeAI🏆 Multi-Agent Research")
    
    # Add voice selector to sidebar
    st.sidebar.markdown("### 🎤 Voice Settings")
    selected_voice = st.sidebar.selectbox(
        "Select TTS Voice:",
        options=EDGE_TTS_VOICES,
        index=EDGE_TTS_VOICES.index(st.session_state['tts_voice'])
    )
    
    # Add audio format selector to sidebar
    st.sidebar.markdown("### 🔊 Audio Format")
    selected_format = st.sidebar.radio(
        "Choose Audio Format:",
        options=["MP3", "WAV"],
        index=0  # Default to MP3
    )
    
    # Update session state if voice or format changes
    if selected_voice != st.session_state['tts_voice']:
        st.session_state['tts_voice'] = selected_voice
        st.rerun()
    if selected_format.lower() != st.session_state['audio_format']:
        st.session_state['audio_format'] = selected_format.lower()
        st.rerun()

    tab_main = st.radio("Action:",["🎤 Voice","📸 Media","🔍 ArXiv","📝 Editor"],horizontal=True)

    mycomponent = components.declare_component("mycomponent", path="mycomponent")
    val = mycomponent(my_input_value="Hello")

    # Show input in a text box for editing if detected
    if val:
        val_stripped = val.replace('\\n', ' ')
        edited_input = st.text_area("✏️ Edit Input:", value=val_stripped, height=100)
        #edited_input = edited_input.replace('\n', ' ')
        
        run_option = st.selectbox("Model:", ["Arxiv", "GPT-4o", "Claude-3.5"])
        col1, col2 = st.columns(2)
        with col1:
            autorun = st.checkbox("⚙ AutoRun", value=True)
        with col2:
            full_audio = st.checkbox("📚FullAudio", value=False, 
                                     help="Generate full audio response")

        input_changed = (val != st.session_state.old_val)

        if autorun and input_changed:
            st.session_state.old_val = val
            st.session_state.last_query = edited_input  # Store the last query for zip naming
            if run_option == "Arxiv":
                perform_ai_lookup(edited_input, vocal_summary=True, extended_refs=False, 
                                  titles_summary=True, full_audio=full_audio)
            else:
                if run_option == "GPT-4o":
                    process_with_gpt(edited_input)
                elif run_option == "Claude-3.5":
                    process_with_claude(edited_input)
        else:
            if st.button("▶ Run"):
                st.session_state.old_val = val
                st.session_state.last_query = edited_input  # Store the last query for zip naming
                if run_option == "Arxiv":
                    perform_ai_lookup(edited_input, vocal_summary=True, extended_refs=False, 
                                      titles_summary=True, full_audio=full_audio)
                else:
                    if run_option == "GPT-4o":
                        process_with_gpt(edited_input)
                    elif run_option == "Claude-3.5":
                        process_with_claude(edited_input)

    if tab_main == "🔍 ArXiv":
        st.subheader("🔍 Query ArXiv")
        q = st.text_input("🔍 Query:")

        st.markdown("### 🎛 Options")
        vocal_summary = st.checkbox("🎙ShortAudio", value=True)
        extended_refs = st.checkbox("📜LongRefs", value=False)
        titles_summary = st.checkbox("🔖TitlesOnly", value=True)
        full_audio = st.checkbox("📚FullAudio", value=False,
                                 help="Full audio of results")
        full_transcript = st.checkbox("🧾FullTranscript", value=False,
                                      help="Generate a full transcript file")

        if q and st.button("🔍Run"):
            st.session_state.last_query = q  # Store the last query for zip naming
            result = perform_ai_lookup(q, vocal_summary=vocal_summary, extended_refs=extended_refs, 
                                       titles_summary=titles_summary, full_audio=full_audio)
            if full_transcript:
                save_full_transcript(q, result)

        st.markdown("### Change Prompt & Re-Run")
        q_new = st.text_input("🔄 Modify Query:")
        if q_new and st.button("🔄 Re-Run with Modified Query"):
            st.session_state.last_query = q_new  # Update last query
            result = perform_ai_lookup(q_new, vocal_summary=vocal_summary, extended_refs=extended_refs, 
                                       titles_summary=titles_summary, full_audio=full_audio)
            if full_transcript:
                save_full_transcript(q_new, result)

    elif tab_main == "🎤 Voice":
        st.subheader("🎤 Voice Input")
        user_text = st.text_area("💬 Message:", height=100)
        user_text = user_text.strip().replace('\n', ' ')
        if st.button("📨 Send"):
            process_with_gpt(user_text)
        st.subheader("📜 Chat History")
        t1,t2=st.tabs(["Claude History","GPT-4o History"])
        with t1:
            for c in st.session_state.chat_history:
                st.write("**You:**", c["user"])
                st.write("**Claude:**", c["claude"])
        with t2:
            for m in st.session_state.messages:
                with st.chat_message(m["role"]):
                    st.markdown(m["content"])

    elif tab_main == "📸 Media":
        st.header("📸 Images & 🎥 Videos")
        tabs = st.tabs(["🖼 Images", "🎥 Video"])
        with tabs[0]:
            imgs = glob.glob("*.png")+glob.glob("*.jpg")
            if imgs:
                c = st.slider("Cols",1,5,3)
                cols = st.columns(c)
                for i,f in enumerate(imgs):
                    with cols[i%c]:
                        st.image(Image.open(f),use_container_width=True)
                        if st.button(f"👀 Analyze {os.path.basename(f)}", key=f"analyze_{f}"):
                            a = process_image(f,"Describe this image.")
                            st.markdown(a)
            else:
                st.write("No images found.")
        with tabs[1]:
            vids = glob.glob("*.mp4")
            if vids:
                for v in vids:
                    with st.expander(f"🎥 {os.path.basename(v)}"):
                        st.video(v)
                        if st.button(f"Analyze {os.path.basename(v)}", key=f"analyze_{v}"):
                            a = process_video_with_gpt(v,"Describe video.")
                            st.markdown(a)
            else:
                st.write("No videos found.")

    elif tab_main == "📝 Editor":
        if getattr(st.session_state,'current_file',None):
            st.subheader(f"Editing: {st.session_state.current_file}")
            new_text = st.text_area("✏️ Content:", st.session_state.file_content, height=300)
            if st.button("💾 Save"):
                with open(st.session_state.current_file,'w',encoding='utf-8') as f:
                    f.write(new_text)
                st.success("Updated!")
                st.session_state.should_rerun = True
        else:
            st.write("Select a file from the sidebar to edit.")

    # Load and display files in the sidebar
    groups_sorted = load_files_for_sidebar()
    display_file_manager_sidebar(groups_sorted)

    if st.session_state.viewing_prefix and any(st.session_state.viewing_prefix == group for group, _ in groups_sorted):
        st.write("---")
        st.write(f"**Viewing Group:** {st.session_state.viewing_prefix}")
        for group_name, files in groups_sorted:
            if group_name == st.session_state.viewing_prefix:
                for f in files:
                    fname = os.path.basename(f)
                    ext = os.path.splitext(fname)[1].lower().strip('.')
                    st.write(f"### {fname}")
                    if ext == "md":
                        content = open(f,'r',encoding='utf-8').read()
                        st.markdown(content)
                    elif ext == "mp3":
                        st.audio(f)
                    elif ext == "wav":
                        st.audio(f)  # 🆕 Handle WAV files
                    else:
                        st.markdown(get_download_link(f), unsafe_allow_html=True)
                break
        if st.button("❌ Close"):
            st.session_state.viewing_prefix = None

    markdownPapers = """
    
    # Levels of AGI

## 1. Performance (rows) x Generality (columns)
- **Narrow**  
  - *clearly scoped or set of tasks*  
- **General**  
  - *wide range of non-physical tasks, including metacognitive abilities like learning new skills*  

## 2. Levels of AGI

### 2.1 Level 0: No AI
- **Narrow Non-AI**  
  - Calculator software; compiler  
- **General Non-AI**  
  - Human-in-the-loop computing, e.g., Amazon Mechanical Turk  

### 2.2 Level 1: Emerging  
*equal to or somewhat better than an unskilled human*  
- **Emerging Narrow AI**  
  - GOFAI; simple rule-based systems  
  - Example: SHRDLU  
    - *Reference:* Winograd, T. (1971). **Procedures as a Representation for Data in a Computer Program for Understanding Natural Language**. MIT AI Technical Report. [Link](https://dspace.mit.edu/handle/1721.1/7095)  
- **Emerging AGI**  
  - ChatGPT (OpenAI, 2023)  
  - Bard (Anil et al., 2023)  
    - *Reference:* Anil, R., et al. (2023). **Bard: Google’s AI Chatbot**. [arXiv](https://arxiv.org/abs/2303.12712)  
  - LLaMA 2 (Touvron et al., 2023)  
    - *Reference:* Touvron, H., et al. (2023). **LLaMA 2: Open and Efficient Foundation Language Models**. [arXiv](https://arxiv.org/abs/2307.09288)  

### 2.3 Level 2: Competent  
*at least 50th percentile of skilled adults*  
- **Competent Narrow AI**  
  - Toxicity detectors such as Jigsaw  
    - *Reference:* Das, S., et al. (2022). **Toxicity Detection at Scale with Jigsaw**. [arXiv](https://arxiv.org/abs/2204.06905)  
  - Smart Speakers (Apple, Amazon, Google)  
  - VQA systems (PaLI)  
    - *Reference:* Chen, T., et al. (2023). **PaLI: Pathways Language and Image model**. [arXiv](https://arxiv.org/abs/2301.01298)  
  - Watson (IBM)  
  - SOTA LLMs for subsets of tasks  
- **Competent AGI**  
  - Not yet achieved  

### 2.4 Level 3: Expert  
*at least 90th percentile of skilled adults*  
- **Expert Narrow AI**  
  - Spelling & grammar checkers (Grammarly, 2023)  
  - Generative image models  
    - Example: Imagen  
      - *Reference:* Saharia, C., et al. (2022). **Imagen: Photorealistic Text-to-Image Diffusion Models**. [arXiv](https://arxiv.org/abs/2205.11487)  
    - Example: DALL·E 2  
      - *Reference:* Ramesh, A., et al. (2022). **Hierarchical Text-Conditional Image Generation with CLIP Latents**. [arXiv](https://arxiv.org/abs/2204.06125)  
- **Expert AGI**  
  - Not yet achieved  

### 2.5 Level 4: Virtuoso  
*at least 99th percentile of skilled adults*  
- **Virtuoso Narrow AI**  
  - Deep Blue  
    - *Reference:* Campbell, M., et al. (2002). **Deep Blue**. IBM Journal of Research and Development. [Link](https://research.ibm.com/publications/deep-blue)  
  - AlphaGo  
    - *Reference:* Silver, D., et al. (2016, 2017). **Mastering the Game of Go with Deep Neural Networks and Tree Search**. [Nature](https://www.nature.com/articles/nature16961)  
- **Virtuoso AGI**  
  - Not yet achieved  

### 2.6 Level 5: Superhuman  
*outperforms 100% of humans*  
- **Superhuman Narrow AI**  
  - AlphaFold  
    - *Reference:* Jumper, J., et al. (2021). **Highly Accurate Protein Structure Prediction with AlphaFold**. [Nature](https://www.nature.com/articles/s41586-021-03819-2)  
  - AlphaZero  
    - *Reference:* Silver, D., et al. (2018). **A General Reinforcement Learning Algorithm that Masters Chess, Shogi, and Go through Self-Play**. [Science](https://www.science.org/doi/10.1126/science.aar6404)  
  - StockFish  
    - *Reference:* Stockfish (2023). **Stockfish Chess Engine**. [Website](https://stockfishchess.org)  
- **Artificial Superintelligence (ASI)**  
  - Not yet achieved  


# 🧬 Innovative Architecture of AlphaFold2: A Hybrid System

## 1. 🔢 Input Sequence  
- The process starts with an **input sequence** (protein sequence).  

## 2. 🗄️ Database Searches  
- **Genetic database search** 🔍  
  - Searches genetic databases to retrieve related sequences.  
- **Structure database search** 🔍  
  - Searches structural databases for template structures.  
- **Pairing** 🤝  
  - Aligns sequences and structures for further analysis.  

## 3. 🧩 MSA (Multiple Sequence Alignment)  
- **MSA representation** 📊 (r,c)  
  - Representation of multiple aligned sequences used as input.  

## 4. 📑 Templates  
- Template structures are paired to assist the model.  

## 5. 🔄 Evoformer (48 blocks)  
- A **deep learning module** that refines representations:  
  - **MSA representation** 🧱  
  - **Pair representation** 🧱 (r,c)  

## 6. 🧱 Structure Module (8 blocks)  
- Converts the representations into:  
  - **Single representation** (r,c)  
  - **Pair representation** (r,c)  

## 7. 🧬 3D Structure Prediction  
- The structure module predicts the **3D protein structure**.  
- **Confidence levels**:  
  - 🔵 *High confidence*  
  - 🟠 *Low confidence*  

## 8. ♻️ Recycling (Three Times)  
- The model **recycles** its output up to three times to refine the prediction.  

## 9. 📚 Reference  
**Jumper, J., et al. (2021).** Highly Accurate Protein Structure Prediction with AlphaFold. *Nature.*  
🔗 [Nature Publication Link](https://www.nature.com/articles/s41586-021-03819-2)  

    """
    st.sidebar.markdown(markdownPapers)
    
    if st.session_state.should_rerun:
        st.session_state.should_rerun = False
        st.rerun()

if __name__=="__main__":
    main()