import gradio as gr
import pandas as pd
import numpy as np
import os
from datetime import datetime

# Load the papers data
def load_papers():
    try:
        papers_df = pd.read_csv('all_papers_0328.csv')
        # Clean up columns if needed and handle missing values
        papers_df = papers_df.fillna('')
        
        # Filter out papers with empty titles
        papers_df = papers_df[papers_df['Title'].str.strip() != '']
        
        # Ensure Year is integer
        papers_df['Year'] = pd.to_numeric(papers_df['Year'], errors='coerce').fillna(0).astype(int)
        
        return papers_df
    except Exception as e:
        print(f"Error loading papers: {e}")
        # Return empty dataframe with expected columns
        return pd.DataFrame(columns=['Title', 'TLDR-EN', 'Section', 'url', 'Year', 'Publish Venue'])

# Search function
def search_papers(search_term, section_filter, year_filter, sort_by):
    papers_df = load_papers()
    
    if search_term:
        # Case-insensitive search across multiple columns
        search_mask = (
            papers_df['Title'].str.contains(search_term, case=False, na=False, regex=True) |
            papers_df['TLDR-EN'].str.contains(search_term, case=False, na=False, regex=True) |
            papers_df['Section'].str.contains(search_term, case=False, na=False, regex=True) |
            papers_df['Publish Venue'].str.contains(search_term, case=False, na=False, regex=True)
        )
        papers_df = papers_df[search_mask]
    
    # Apply section filter if selected
    if section_filter != "All Sections":
        papers_df = papers_df[papers_df['Section'] == section_filter]
    
    # Apply year filter if selected
    if year_filter != "All Years":
        papers_df = papers_df[papers_df['Year'] == int(year_filter)]
    
    # Sort based on selection
    if sort_by == "Year (newest first)":
        papers_df = papers_df.sort_values(by=['Year', 'Title'], ascending=[False, True])
    elif sort_by == "Year (oldest first)":
        papers_df = papers_df.sort_values(by=['Year', 'Title'], ascending=[True, True])
    elif sort_by == "Title (A-Z)":
        papers_df = papers_df.sort_values(by='Title')
    elif sort_by == "Section":
        papers_df = papers_df.sort_values(by=['Section', 'Year', 'Title'], ascending=[True, False, True])
    
    # Format for display
    html_output = "<div class='papers-container'>"
    
    if len(papers_df) == 0:
        html_output += "<p>No papers found matching your criteria.</p>"
    else:
        for i, row in papers_df.iterrows():
            html_output += f"""
            <div class='paper-card'>
                <div class='paper-title'>
                    <a href='{row['url']}' target='_blank'>{row['Title']}</a>
                </div>
                <div class='paper-tldr'>{row['TLDR-EN']}</div>
                <div class='paper-meta'>
                    <span class='meta-item section'>{row['Section']}</span>
                    <span class='meta-item year'>{row['Year']}</span>
                    <span class='meta-item venue'>{row['Publish Venue']}</span>
                </div>
            </div>
            """
    
    html_output += "</div>"
    
    # Add paper count
    paper_count = len(papers_df)
    count_text = f"<p><strong>{paper_count} papers</strong> found</p>"
    
    return count_text + html_output

# Get unique sections and years for filtering
def get_filter_options():
    papers_df = load_papers()
    sections = ["All Sections"] + sorted(papers_df['Section'].unique().tolist())
    years = ["All Years"] + [str(year) for year in sorted(papers_df['Year'].unique().tolist(), reverse=True) if year > 0]
    return sections, years

# Custom CSS
custom_css = """
/* Main container */
body {
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, 'Open Sans', 'Helvetica Neue', sans-serif;
}

.papers-container {
    display: flex;
    flex-direction: column;
    gap: 18px;
    margin-top: 20px;
}

/* Paper card styling */
.paper-card {
    border: 1px solid #e0e0e0;
    border-radius: 12px;
    padding: 20px;
    background-color: #ffffff;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
    transition: all 0.2s ease;
    display: flex;
    flex-direction: column;
    gap: 10px;
}

.paper-card:hover {
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
    transform: translateY(-2px);
    border-color: #d0d0d0;
}

.paper-title {
    font-size: 18px;
    font-weight: 600;
    line-height: 1.4;
    margin-bottom: 4px;
}

.paper-title a {
    color: #2563EB;
    text-decoration: none;
}

.paper-title a:hover {
    text-decoration: underline;
}

.paper-tldr {
    font-size: 14px;
    color: #4B5563;
    line-height: 1.5;
    margin: 8px 0;
}

.paper-meta {
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
    margin-top: 4px;
}

.meta-item {
    background-color: #F3F4F6;
    border-radius: 16px;
    padding: 4px 12px;
    font-size: 12px;
    color: #4B5563;
    font-weight: 500;
}

/* Section colors */
.meta-item.section {
    background-color: #DBEAFE;
    color: #1E40AF;
}

.meta-item.year {
    background-color: #FEE2E2;
    color: #991B1B;
}

.meta-item.venue {
    background-color: #E0E7FF;
    color: #3730A3;
}

/* Responsive design */
@media (max-width: 768px) {
    .paper-card {
        padding: 16px;
    }
    
    .paper-title {
        font-size: 16px;
    }
    
    .paper-tldr {
        font-size: 13px;
    }
    
    .meta-item {
        font-size: 11px;
        padding: 3px 10px;
    }
}

/* Results count styling */
p strong {
    color: #2563EB;
}
"""

# Create the Gradio interface
def create_interface():
    sections, years = get_filter_options()
    
    # Get paper statistics
    papers_df = load_papers()
    total_papers = len(papers_df)
    paper_counts_by_section = papers_df['Section'].value_counts().to_dict()
    paper_counts_by_year = papers_df['Year'].value_counts().to_dict()
    
    # Filter out year 0 if it exists
    min_year = min([year for year in paper_counts_by_year.keys() if year > 0]) if paper_counts_by_year else 'N/A'
    max_year = max(paper_counts_by_year.keys()) if paper_counts_by_year else 'N/A'
    
    # Project description with linked paper
    project_description = f"""
    # Large Language Model Agent: A Survey on Methodology, Applications and Challenges
    
    This application showcases papers from our comprehensive survey on Large Language Model (LLM) agents. We organize papers across key categories including agent construction, collaboration mechanisms, evolution, tools, security, benchmarks, and applications.
    
    ## About the Survey
    
    The era of intelligent agents is upon us, driven by revolutionary advancements in large language models. Large Language Model (LLM) agents, with goal-driven behaviors and dynamic adaptation capabilities, potentially represent a critical pathway toward artificial general intelligence.
    
    This survey systematically deconstructs LLM agent systems through a methodology-centered taxonomy, linking architectural foundations, collaboration mechanisms, and evolutionary pathways. We unify fragmented research threads by revealing fundamental connections between agent design principles and their emergent behaviors in complex environments.
    
    [View the full paper on arXiv](https://arxiv.org/abs/2503.21460)  
    [Explore our GitHub repository](https://github.com/luo-junyu/Awesome-Agent-Papers)
    
    ## Submit Your Paper
    
    We welcome contributions to expand our collection. To submit your paper:
    - Email us at luo.junyu@outlook.com with your paper details
    - Create a pull request on our [GitHub repository](https://github.com/luo-junyu/Awesome-Agent-Papers)
    
    ## Collection Overview
    
    - **Total Papers**: {total_papers}
    - **Categories**: {len(paper_counts_by_section)}
    - **Year Range**: {min_year} - {max_year}
    """
    
    with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
        gr.Markdown(project_description)
        
        with gr.Row():
            with gr.Column(scale=3):
                search_input = gr.Textbox(
                    label="Search Papers", 
                    placeholder="Enter keywords to search titles, summaries, sections, or venues",
                    show_label=True
                )
            
            with gr.Column(scale=1):
                section_dropdown = gr.Dropdown(
                    choices=sections, 
                    value="All Sections", 
                    label="Filter by Section"
                )
        
        with gr.Row():
            with gr.Column(scale=1):
                year_dropdown = gr.Dropdown(
                    choices=years, 
                    value="All Years", 
                    label="Filter by Year"
                )
            
            with gr.Column(scale=1):
                sort_dropdown = gr.Dropdown(
                    choices=[
                        "Year (newest first)", 
                        "Year (oldest first)", 
                        "Title (A-Z)", 
                        "Section"
                    ],
                    value="Year (newest first)",
                    label="Sort by"
                )
        
        search_button = gr.Button("Search", variant="primary")
        
        # Results display
        results_html = gr.HTML(label="Search Results")
        
        # Section distribution chart
        section_data = [[section, count] for section, count in paper_counts_by_section.items()]
        section_data.sort(key=lambda x: x[1], reverse=True)
        
        with gr.Accordion("Paper Distribution by Section", open=False):
            gr.Dataframe(
                headers=["Section", "Count"],
                datatype=["str", "number"],
                value=section_data
            )
        
        # Year distribution chart
        year_data = [[str(year), count] for year, count in paper_counts_by_year.items() if year > 0]
        year_data.sort(key=lambda x: int(x[0]), reverse=True)
        
        with gr.Accordion("Paper Distribution by Year", open=False):
            gr.Dataframe(
                headers=["Year", "Count"],
                datatype=["str", "number"],
                value=year_data
            )
        
        # # Add example searches
        # gr.Examples(
        #     examples=[
        #         ["agent collaboration", "All Sections", "All Years", "Year (newest first)"],
        #         ["security", "Security", "All Years", "Year (newest first)"],
        #         ["benchmark", "Datasets & Benchmarks", "2024", "Year (newest first)"],
        #         ["tools", "Tools", "All Years", "Year (newest first)"],
        #     ],
        #     inputs=[search_input, section_dropdown, year_dropdown, sort_dropdown],
        #     outputs=results_html,
        #     fn=search_papers,
        #     cache_examples=True,
        # )
        
        # Set up search on button click and input changes
        search_button.click(
            fn=search_papers,
            inputs=[search_input, section_dropdown, year_dropdown, sort_dropdown],
            outputs=results_html
        )
        
        # Also search when dropdown values change
        section_dropdown.change(
            fn=search_papers,
            inputs=[search_input, section_dropdown, year_dropdown, sort_dropdown],
            outputs=results_html
        )
        
        year_dropdown.change(
            fn=search_papers,
            inputs=[search_input, section_dropdown, year_dropdown, sort_dropdown],
            outputs=results_html
        )
        
        sort_dropdown.change(
            fn=search_papers,
            inputs=[search_input, section_dropdown, year_dropdown, sort_dropdown],
            outputs=results_html
        )
        
        # Load initial results on page load
        demo.load(
            fn=lambda: search_papers("", "All Sections", "All Years", "Year (newest first)"),
            inputs=None,
            outputs=results_html
        )
    
    return demo

# Create and launch the interface
demo = create_interface()

if __name__ == "__main__":
    demo.launch()