app.py

import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta
from dotenv import load_dotenv
import os
from typing import List, Dict

from src.clients.pubmed_client import PubMedClient
from src.clients.biorxiv_client import BioRxivClient
from src.embedder import PaperEmbedder
from src.storage import PaperStorage

load_dotenv()

st.set_page_config(
    page_title="Literature Intelligence Tool",
    page_icon="\U0001f52c",
    layout="wide"
)


@st.cache_resource
def initialize_components():
    """Initialize clients, embedder, and storage."""
    email = os.getenv("PUBMED_EMAIL", "user@example.com")
    api_key = os.getenv("PUBMED_API_KEY")
    persist_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
    model_name = os.getenv("EMBEDDING_MODEL", "pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb")
    verify_ssl = os.getenv("VERIFY_SSL", "false").lower() == "true"

    pubmed_client = PubMedClient(email=email, api_key=api_key, verify_ssl=verify_ssl)
    biorxiv_client = BioRxivClient()
    embedder = PaperEmbedder(model_name=model_name)
    storage = PaperStorage(persist_directory=persist_dir)

    return pubmed_client, biorxiv_client, embedder, storage


def fetch_and_index_papers(
    query: str,
    sources: List[str],
    max_results: int,
    days_back: int,
    pubmed_client: PubMedClient,
    biorxiv_client: BioRxivClient,
    embedder: PaperEmbedder,
    storage: PaperStorage
) -> Dict[str, int]:
    """Fetch papers from selected sources and index them."""
    all_papers = []

    if "PubMed" in sources:
        with st.spinner("Fetching from PubMed..."):
            pubmed_papers = pubmed_client.search(
                query=query,
                max_results=max_results,
                days_back=days_back
            )
            all_papers.extend(pubmed_papers)

    if "bioRxiv" in sources:
        with st.spinner("Fetching from bioRxiv..."):
            biorxiv_papers = biorxiv_client.search(
                query=query,
                max_results=max_results,
                days_back=days_back
            )
            all_papers.extend(biorxiv_papers)

    if all_papers:
        with st.spinner(f"Generating embeddings for {len(all_papers)} papers..."):
            abstracts = [f"{p.get('title', '')}. {p.get('abstract', '')}" for p in all_papers]
            embeddings = embedder.embed_batch(abstracts)

        with st.spinner("Indexing papers..."):
            added_count = storage.add_papers(all_papers, embeddings)

        return {
            'total_fetched': len(all_papers),
            'newly_added': added_count,
            'already_existed': len(all_papers) - added_count
        }

    return {'total_fetched': 0, 'newly_added': 0, 'already_existed': 0}


def search_papers(
    query: str,
    n_results: int,
    source_filter: str,
    embedder: PaperEmbedder,
    storage: PaperStorage
) -> List[Dict]:
    """Semantic search for papers."""
    query_embedding = embedder.embed_text(query)

    source = None if source_filter == "All" else source_filter.lower()

    results = storage.search_papers(
        query_embedding=query_embedding,
        n_results=n_results,
        source_filter=source
    )

    return results


def main():
    st.title("\U0001f52c Literature Intelligence Tool")
    st.markdown("Track, index, and analyze scientific papers from PubMed and bioRxiv")

    pubmed_client, biorxiv_client, embedder, storage = initialize_components()

    stats = storage.get_database_stats()

    st.sidebar.header("Database Statistics")
    st.sidebar.metric("Total Papers", stats['total_papers'])

    if stats['by_source']:
        st.sidebar.markdown("**By Source:**")
        for source, count in stats['by_source'].items():
            st.sidebar.text(f"{source}: {count}")

    tab1, tab2, tab3 = st.tabs(["Fetch & Index", "Semantic Search", "Trends Analysis"])

    with tab1:
        st.header("Fetch and Index Papers")

        col1, col2 = st.columns(2)

        with col1:
            query = st.text_input(
                "Search Query",
                placeholder="e.g., EGFR inhibitor, CAR-T cell therapy, Alzheimer's disease",
                help="Enter gene names, drug names, diseases, or any search terms"
            )

            sources = st.multiselect(
                "Data Sources",
                options=["PubMed", "bioRxiv"],
                default=["PubMed"]
            )

        with col2:
            max_results = st.slider(
                "Maximum Results per Source",
                min_value=10,
                max_value=500,
                value=100,
                step=10
            )

            days_back = st.slider(
                "Days to Look Back",
                min_value=7,
                max_value=365,
                value=90,
                step=7
            )

        if st.button("Fetch and Index Papers", type="primary"):
            if not query:
                st.error("Please enter a search query")
            elif not sources:
                st.error("Please select at least one data source")
            else:
                result = fetch_and_index_papers(
                    query=query,
                    sources=sources,
                    max_results=max_results,
                    days_back=days_back,
                    pubmed_client=pubmed_client,
                    biorxiv_client=biorxiv_client,
                    embedder=embedder,
                    storage=storage
                )

                st.success(f"Fetching complete!")
                st.info(f"Total fetched: {result['total_fetched']} | Newly added: {result['newly_added']} | Already existed: {result['already_existed']}")

                st.rerun()

    with tab2:
        st.header("Semantic Search")

        search_query = st.text_input(
            "Search Query",
            placeholder="e.g., novel therapeutic approaches for cancer",
            help="Describe what you're looking for in natural language"
        )

        col1, col2 = st.columns(2)
        with col1:
            n_results = st.slider("Number of Results", min_value=5, max_value=50, value=10)
        with col2:
            source_filter = st.selectbox("Filter by Source", options=["All", "PubMed", "bioRxiv"])

        if st.button("Search", type="primary"):
            if not search_query:
                st.error("Please enter a search query")
            else:
                with st.spinner("Searching..."):
                    results = search_papers(
                        query=search_query,
                        n_results=n_results,
                        source_filter=source_filter,
                        embedder=embedder,
                        storage=storage
                    )

                if results:
                    st.success(f"Found {len(results)} relevant papers")

                    for i, paper in enumerate(results):
                        with st.expander(
                            f"{i+1}. [{paper['source'].upper()}] {paper['title']} (Similarity: {paper['similarity_score']:.3f})",
                            expanded=(i == 0)
                        ):
                            st.markdown(f"**Paper ID:** {paper['paper_id']}")
                            st.markdown(f"**Publication Date:** {paper['publication_date']}")
                            st.markdown(f"**Authors:** {', '.join(paper['authors'][:5])}" + (
                                f" + {len(paper['authors']) - 5} more" if len(paper['authors']) > 5 else ""
                            ))
                            st.markdown(f"**Abstract:**\n{paper['abstract']}")
                            if paper['keywords']:
                                st.markdown(f"**Keywords:** {', '.join(paper['keywords'][:10])}")
                else:
                    st.warning("No papers found. Try fetching and indexing papers first.")

    with tab3:
        st.header("Trends Analysis")

        if stats['total_papers'] > 0:
            start_date = st.date_input(
                "Start Date",
                value=datetime.now() - timedelta(days=90)
            )
            end_date = st.date_input(
                "End Date",
                value=datetime.now()
            )

            if st.button("Analyze Trends", type="primary"):
                with st.spinner("Analyzing trends..."):
                    papers = storage.get_papers_by_date_range(
                        start_date=start_date.strftime("%Y-%m-%d"),
                        end_date=end_date.strftime("%Y-%m-%d")
                    )

                    if papers:
                        df = pd.DataFrame(papers)
                        df['publication_date'] = pd.to_datetime(df['publication_date'])

                        st.subheader("Publications Over Time")
                        daily_counts = df.groupby(df['publication_date'].dt.date).size().reset_index()
                        daily_counts.columns = ['Date', 'Count']

                        fig = px.line(
                            daily_counts,
                            x='Date',
                            y='Count',
                            title='Publications per Day'
                        )
                        st.plotly_chart(fig, use_container_width=True)

                        col1, col2 = st.columns(2)

                        with col1:
                            st.subheader("By Source")
                            source_counts = df['source'].value_counts()
                            fig = px.pie(
                                values=source_counts.values,
                                names=source_counts.index,
                                title='Distribution by Source'
                            )
                            st.plotly_chart(fig, use_container_width=True)

                        with col2:
                            st.subheader("By Query Tag")
                            tag_counts = df['query_tag'].value_counts().head(10)
                            fig = px.bar(
                                x=tag_counts.index,
                                y=tag_counts.values,
                                title='Top 10 Query Tags',
                                labels={'x': 'Query Tag', 'y': 'Count'}
                            )
                            st.plotly_chart(fig, use_container_width=True)

                        st.subheader("Top Keywords")
                        all_keywords = []
                        for keywords in df['keywords']:
                            all_keywords.extend(keywords)

                        if all_keywords:
                            keyword_counts = pd.Series(all_keywords).value_counts().head(20)
                            fig = px.bar(
                                x=keyword_counts.values,
                                y=keyword_counts.index,
                                orientation='h',
                                title='Top 20 Keywords',
                                labels={'x': 'Count', 'y': 'Keyword'}
                            )
                            st.plotly_chart(fig, use_container_width=True)

                    else:
                        st.warning("No papers found in the selected date range")
        else:
            st.info("No papers in database yet. Use the 'Fetch & Index' tab to add papers.")


if __name__ == "__main__":
    main()