import gradio as gr
import torch
import requests
import pandas as pd
from transformers import pipeline, AutoTokenizer, AutoModel
from sentence_transformers import SentenceTransformer
import json
from typing import List, Dict
import logging
import time
from datetime import datetime

class ResearchExplorer:
    def __init__(self):
        # Initialize models based on available hardware
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        
        # Initialize models
        self.summarizer = pipeline(
            "summarization",
            model="facebook/bart-large-cnn",
            device=0 if self.device == "cuda" else -1
        )
        
        # Initialize sentence transformer for semantic search
        self.sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
        
        # API endpoints
        self.semantic_scholar_api = "https://api.semanticscholar.org/graph/v1/paper/search"
        self.arxiv_api = "http://export.arxiv.org/api/query"
        
        # Configure logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def fetch_papers_semantic_scholar(self, query: str, limit: int = 5) -> List[Dict]:
        """Fetch papers from Semantic Scholar API"""
        try:
            params = {
                "query": query,
                "limit": limit,
                "fields": "title,abstract,url,year,authors,venue"
            }
            response = requests.get(self.semantic_scholar_api, params=params)
            response.raise_for_status()
            return response.json().get("data", [])
        except Exception as e:
            self.logger.error(f"Error fetching from Semantic Scholar: {e}")
            return []

    def summarize_text(self, text: str, max_length: int = 130) -> str:
        """Generate summary using BART model"""
        try:
            if not text or len(text) < 50:
                return text
            summary = self.summarizer(text, max_length=max_length, min_length=30, do_sample=False)
            return summary[0]['summary_text']
        except Exception as e:
            self.logger.error(f"Summarization error: {e}")
            return text[:max_length] + "..."

    def process_papers(self, papers: List[Dict]) -> List[Dict]:
        """Process and format paper information"""
        processed_papers = []
        for paper in papers:
            try:
                # Extract and format paper information
                title = paper.get("title", "No Title")
                abstract = paper.get("abstract", "No abstract available")
                year = paper.get("year", "N/A")
                venue = paper.get("venue", "N/A")
                url = paper.get("url", "#")
                
                # Generate summary
                summary = self.summarize_text(abstract)
                
                # Format authors
                authors = paper.get("authors", [])
                author_names = ", ".join([author.get("name", "") for author in authors[:3]])
                if len(authors) > 3:
                    author_names += " et al."
                
                processed_papers.append({
                    "Title": title,
                    "Authors": author_names,
                    "Year": year,
                    "Venue": venue,
                    "Summary": summary,
                    "URL": url
                })
            except Exception as e:
                self.logger.error(f"Error processing paper: {e}")
                continue
                
        return processed_papers

    def search_papers(self, query: str, limit: int = 5) -> List[Dict]:
        """Main function to search and process papers"""
        # Fetch papers
        papers = self.fetch_papers_semantic_scholar(query, limit)
        
        # Process and return results
        return self.process_papers(papers)

# Initialize the explorer
explorer = ResearchExplorer()

# Define Gradio interface
def gradio_interface(query: str, num_papers: int) -> pd.DataFrame:
    """Gradio interface function"""
    if not query:
        return pd.DataFrame()
    
    results = explorer.search_papers(query, limit=num_papers)
    return pd.DataFrame(results)

# Create and launch the interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Enter Research Topic", placeholder="e.g., machine learning, quantum computing"),
        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Number of Papers")
    ],
    outputs=gr.DataFrame(label="Research Papers"),
    title="Research Paper Explorer",
    description="Search and explore academic papers with AI-powered summaries",
    examples=[
        ["machine learning in healthcare", 5],
        ["quantum computing applications", 3],
        ["climate change mitigation", 4]
    ],
    theme="default"
)

# Launch the interface
iface.launch(share=True)