import streamlit as st
import pandas as pd
import plotly.express as px

# Page configuration MUST be the first Streamlit command
st.set_page_config(
    page_title="🚀 NASA Bioscience Explorer",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Now import transformers and other libraries
try:
    from transformers import pipeline
    import trafilatura
    import textwrap
    TRANSFORMERS_AVAILABLE = True
except ImportError as e:
    TRANSFORMERS_AVAILABLE = False
    st.error(f"⚠️ Transformers library not available: {e}")
    st.info("The app will work without summarization features.")

# Load the summarizer
def load_summarizer():
    if not TRANSFORMERS_AVAILABLE:
        return None
    try:
        return pipeline("summarization", model="facebook/bart-large-cnn", device=-1)
    except Exception as e:
        st.error(f"Error loading summarizer: {e}")
        return None

# Cached summarization function
def summarize_entire_text(text):
    summarizer = load_summarizer()
    if summarizer is None:
        return "Summarization not available. Please check the logs for errors."
    
    # Limit input to 8000 characters
    text = text[:8000] if len(text) > 8000 else text
    
    # Split text into chunks of 4000 characters
    chunk_size = 4000
    chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
    
    # Summarize each chunk
    summaries = []
    for chunk in chunks:
        if chunk.strip():
            try:
                result = summarizer(chunk, max_length=150, min_length=40, do_sample=False)
                summaries.append(result[0]['summary_text'])
            except Exception as e:
                st.error(f"Error summarizing chunk: {e}")
                continue
    
    # Combine summaries if there are multiple chunks
    if len(summaries) > 1:
        combined_summary = " ".join(summaries)
        try:
            final_result = summarizer(combined_summary, max_length=250, min_length=100, do_sample=False)
            return final_result[0]['summary_text']
        except Exception as e:
            st.error(f"Error in final summarization: {e}")
            return " ".join(summaries)
    elif len(summaries) == 1:
        return summaries[0]
    else:
        return "No text to summarize."

# Load data
def load_data():
    try:
        df = pd.read_csv('./data/SB_publication_PMC.csv')
        if df.empty:
            return pd.DataFrame()
    except Exception as e:
        st.error(f"Error loading data: {str(e)}")
        return pd.DataFrame()
    
    # Simple categorization based on title keywords
    def categorize_topic(title):
        title_lower = title.lower()
        if any(word in title_lower for word in ['bone', 'skeletal', 'oste']):
            return 'Bone Health'
        elif any(word in title_lower for word in ['muscle', 'atrophy']):
            return 'Muscle Physiology'
        elif any(word in title_lower for word in ['immune', 'infection', 'microbiome']):
            return 'Immune System'
        elif any(word in title_lower for word in ['plant', 'arabidopsis', 'root']):
            return 'Plant Biology'
        elif any(word in title_lower for word in ['radiation', 'dna', 'genomic']):
            return 'Radiation Effects'
        elif any(word in title_lower for word in ['microgravity', 'gravity']):
            return 'Microgravity Adaptation'
        else:
            return 'Other'
    
    def detect_organism(title):
        title_lower = title.lower()
        if 'mouse' in title_lower or 'mice' in title_lower:
            return 'Mouse'
        elif 'arabidopsis' in title_lower:
            return 'Arabidopsis'
        elif 'drosophila' in title_lower:
            return 'Drosophila'
        elif 'human' in title_lower or 'astronaut' in title_lower:
            return 'Human'
        elif 'rat' in title_lower:
            return 'Rat'
        else:
            return 'Various'
    
    df['topic'] = df['Title'].apply(categorize_topic)
    df['organism'] = df['Title'].apply(detect_organism)
    
    return df

# Function to summarize paper from URL
def summarize_paper(url):
    if not TRANSFORMERS_AVAILABLE:
        return "Summarization feature is not available."
    
    try:
        downloaded = trafilatura.fetch_url(url)
        text = trafilatura.extract(downloaded) if downloaded else None
        
        if text:
            summary = summarize_entire_text(text)
            return summary
        else:
            return None
    except Exception as e:
        st.error(f"Error summarizing paper: {str(e)}")
        return None

# Simple summarizer function
def summarize_from_url(url):
    if not TRANSFORMERS_AVAILABLE:
        return "❌ Summarization feature is not available."
    
    try:
        downloaded = trafilatura.fetch_url(url)
        text = trafilatura.extract(downloaded) if downloaded else None

        if text:
            summary = summarize_entire_text(text)
            return summary
        else:
            return "❌ Failed to extract text from the URL."
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Filter publications
def filter_publications(df, search_term, selected_topics, selected_organisms):
    filtered_df = df.copy()
    
    if selected_topics and len(selected_topics) > 0:
        filtered_df = filtered_df[filtered_df['topic'].isin(selected_topics)]
    
    if selected_organisms and len(selected_organisms) > 0:
        filtered_df = filtered_df[filtered_df['organism'].isin(selected_organisms)]
    
    if search_term and search_term.strip():
        search_terms = search_term.lower().split()
        search_mask = pd.Series(True, index=filtered_df.index)
        for term in search_terms:
            term_mask = (
                filtered_df['Title'].str.lower().str.contains(term, na=False) |
                filtered_df['topic'].str.lower().str.contains(term, na=False) |
                filtered_df['organism'].str.lower().str.contains(term, na=False)
            )
            search_mask &= term_mask
        filtered_df = filtered_df[search_mask]
    
    return filtered_df

# Main app
def main():
    st.title("🚀 NASA Bioscience Explorer")
    st.markdown("Explore 608 NASA life sciences publications")
    
    # Load data
    df = load_data()
    
    if df.empty:
        st.error("Failed to load data. Please check if the data file exists.")
        return
    
    # Initialize session state for summaries
    if 'summaries' not in st.session_state:
        st.session_state.summaries = {}
    
    # Create header section with filters
    st.markdown("### 🔍 Search and Filter Publications")
    
    # Create three columns for filters
    search_col, topic_col, organism_col = st.columns([1, 1, 1])
    
    with search_col:
        search_term = st.text_input(
            "Search publications:", 
            placeholder="Enter keywords..."
        )
    
    with topic_col:
        topic_options = df['topic'].unique().tolist()
        selected_topics = st.multiselect(
            "Research Topics:",
            options=topic_options,
            default=[]
        )
    
    with organism_col:
        organism_options = df['organism'].unique().tolist()
        selected_organisms = st.multiselect(
            "Organisms:",
            options=organism_options,
            default=[]
        )
    
    st.markdown("---")
    
    # Filter data
    filtered_df = filter_publications(df, search_term, selected_topics, selected_organisms)
    
    # Metrics
    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Total Publications", len(df))
    col2.metric("Filtered Publications", len(filtered_df))
    col3.metric("Research Topics", df['topic'].nunique())
    col4.metric("Organisms Studied", df['organism'].nunique())
    
    # Create tabs
    tab1, tab2 = st.tabs(["📊 Research Dashboard", "📄 Paper Summarizer"])
    
    with tab1:
        if not filtered_df.empty:
            col1, col2 = st.columns(2)
            
            with col1:
                topic_counts = filtered_df['topic'].value_counts()
                topic_labels = [f"{topic} ({count})" for topic, count in topic_counts.items()]
                fig1 = px.pie(
                    values=topic_counts.values,
                    names=topic_labels,
                    title="📈 Research Topics Distribution"
                )
                fig1.update_traces(textinfo='percent+label')
                st.plotly_chart(fig1, use_container_width=True)
            
            with col2:
                organism_counts = filtered_df['organism'].value_counts().reset_index()
                organism_counts.columns = ['Organism', 'Count']
                organism_counts['Label'] = organism_counts.apply(lambda x: f"{x['Organism']} ({x['Count']})", axis=1)
                fig2 = px.bar(
                    data_frame=organism_counts,
                    x='Label',
                    y='Count',
                    title="🧬 Publications by Organism"
                )
                fig2.update_xaxes(tickangle=45)
                fig2.update_layout(xaxis_title="")
                st.plotly_chart(fig2, use_container_width=True)
        
        st.markdown("---")
        st.subheader("📚 Publication Browser")
        
        if not filtered_df.empty:
            for idx, row in filtered_df.iterrows():
                with st.expander(f"**{row['Title']}**"):
                    st.write(f"**Topic:** {row['topic']}")
                    st.write(f"**Organism:** {row['organism']}")
                    st.markdown(f"[📄 Read Paper]({row['Link']})")
                    
                    summary_key = f"summary_{idx}"
                    
                    if summary_key not in st.session_state:
                        st.session_state[summary_key] = None
                    
                    if TRANSFORMERS_AVAILABLE and st.button("🔍 Generate Summary", key=f"btn_{idx}"):
                        with st.spinner("Generating summary..."):
                            summary = summarize_paper(row['Link'])
                            if summary:
                                st.session_state[summary_key] = summary
                            else:
                                st.error("❌ Failed to extract text from this paper.")
                    
                    if st.session_state[summary_key]:
                        st.subheader("Summary")
                        st.write(st.session_state[summary_key])
        else:
            st.warning("🔍 No publications match the current filters.")
        
        st.markdown("---")
        st.subheader("💡 Research Insights")
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("### 🎯 Most Studied Areas")
            top_topics = df['topic'].value_counts().head(3)
            for topic, count in top_topics.items():
                st.write(f"- **{topic}**: {count} publications")
        
        with col2:
            st.markdown("### 🔬 Research Gaps")
            gaps = [
                "Limited long-duration human studies",
                "Combined radiation + microgravity effects",
                "Psychological health in space"
            ]
            for gap in gaps:
                st.write(f"- {gap}")
    
    with tab2:
        st.markdown("### 📄 Research Paper Summarizer")
        
        if not TRANSFORMERS_AVAILABLE:
            st.warning("⚠️ Summarization feature is currently unavailable. Please check the logs.")
        else:
            st.markdown("Enter any scientific article URL to get an AI-generated summary")
            
            url_input = st.text_input(
                "Enter Article URL:",
                value="https://pmc.ncbi.nlm.nih.gov/articles/PMC10772081/",
                placeholder="https://pmc.ncbi.nlm.nih.gov/articles/...",
                key="url_input"
            )
            
            col1, col2 = st.columns([1, 4])
            with col1:
                if st.button("🚀 Summarize Paper", type="primary"):
                    if url_input:
                        with st.spinner("🔍 Generating summary..."):
                            summary = summarize_from_url(url_input)
                            if summary:
                                with col2:
                                    st.info(summary)
                    else:
                        with col2:
                            st.warning("⚠️ Please enter a URL to summarize")

if __name__ == "__main__":
    main()