Chat With Arxiv

"""
ArXiv Chatbot Module

Conversational interface for exploring ArXiv papers.
"""

from datetime import datetime
from typing import Dict, List


class ArXivChatbot:
    """Chatbot for interacting with ArXiv papers."""

    def __init__(self):
        """Initialize chatbot."""
        self.answerer = None  # Will be initialized with PaperQuestionAnswerer
        self.synthesizer = None  # Will be initialized with ResearchSynthesizer
        self.conversation_history = []
        self.context_manager = None  # Will be initialized with ConversationContextManager

    def chat(self, user_message: str) -> str:
        """Handle user message in conversation."""
        # Store in history
        self.conversation_history.append({
            "role": "user",
            "content": user_message,
            "timestamp": datetime.now()
        })

        # Determine query type
        query_type = self.classify_query(user_message)

        # Generate response
        if query_type == "single_paper_qa":
            response = self.handle_single_paper_query(user_message)
        elif query_type == "multi_paper_synthesis":
            response = self.handle_synthesis_query(user_message)
        elif query_type == "research_trend":
            response = self.handle_trend_query(user_message)
        else:
            response = self.handle_general_query(user_message)

        # Store response
        self.conversation_history.append({
            "role": "assistant",
            "content": response,
            "timestamp": datetime.now()
        })

        return response

    def classify_query(self, message: str) -> str:
        """Classify type of user query."""
        message_lower = message.lower()

        if any(word in message_lower for word in ["trend", "recent", "latest"]):
            return "research_trend"
        elif any(word in message_lower for word in ["compare", "difference", "similar", "across"]):
            return "multi_paper_synthesis"
        elif any(word in message_lower for word in ["summarize", "explain", "about"]):
            return "single_paper_qa"
        else:
            return "general"

    def handle_single_paper_query(self, message: str) -> str:
        """Handle question about specific papers."""
        if not self.answerer:
            return "Paper question answerer not configured."

        result = self.answerer.answer_with_citations(message)

        response = f"{result['answer']}\n\n**Sources:**\n"
        response += "\n".join([f"- {citation}" for citation in result['citations']])

        return response

    def handle_synthesis_query(self, message: str) -> str:
        """Handle synthesis query across papers."""
        if not self.synthesizer:
            return "Synthesis capabilities not configured."

        topic = self.extract_topic_from_query(message)
        synthesis = self.synthesizer.synthesize_research_area(topic, num_papers=15)

        response = f"## Research Synthesis: {topic}\n\n"
        response += f"Analyzed {synthesis['papers_analyzed']} papers\n\n"
        response += f"**Key Themes:** {', '.join(synthesis.get('key_themes', [])[:5])}\n\n"
        response += f"**Summary:** {synthesis.get('synthesis', 'Summary not available')}\n"

        return response

    def handle_trend_query(self, message: str) -> str:
        """Handle research trend query."""
        if not self.answerer or not self.answerer.retriever:
            return "Trend analysis not configured."

        try:
            result = self.answerer.retriever.get_trending_papers("cs.AI", days=7)

            response = "## Recent Research Trends\n\n"
            for i, paper in enumerate(result[:5], 1):
                response += f"{i}. **{paper['title']}**\n"
                response += f"   Authors: {', '.join(paper['authors'][:2])}\n"
                response += f"   {paper['summary'][:200]}...\n\n"

            return response
        except Exception as e:
            return f"Error retrieving trends: {str(e)}"

    def handle_general_query(self, message: str) -> str:
        """Handle general knowledge query."""
        if not self.answerer:
            return "Question answerer not configured."

        return self.answerer.answer_question(message)

    def extract_topic_from_query(self, message: str) -> str:
        """Extract main topic from query."""
        return message.replace("?", "").replace(".", "")

    def get_conversation_context(self) -> List[Dict]:
        """Get conversation history for context."""
        return self.conversation_history

    def reset_conversation(self):
        """Reset conversation history."""
        self.conversation_history = []

"""
ArXiv Paper Retriever Module

Handles paper discovery and retrieval from ArXiv.
"""

import arxiv
from datetime import datetime, timedelta
from typing import List, Dict


class ArXivPaperRetriever:
    """Retrieves papers from ArXiv."""

    def __init__(self, max_results: int = 50):
        """
        Initialize retriever.

        Args:
            max_results: Maximum results per search
        """
        self.client = arxiv.Client()
        self.max_results = max_results

    def search_papers(self, query: str, sort_by=arxiv.SortCriterion.Relevance) -> List[Dict]:
        """Search ArXiv for papers."""
        search = arxiv.Search(
            query=query,
            sort_by=sort_by,
            max_results=self.max_results
        )

        papers = []
        for result in self.client.results(search):
            papers.append({
                "title": result.title,
                "authors": [author.name for author in result.authors],
                "summary": result.summary,
                "published": result.published,
                "arxiv_id": result.arxiv_id,
                "pdf_url": result.pdf_url,
                "categories": result.categories,
                "doi": result.doi,
            })

        return papers

    def search_by_category(self, category: str, from_date: datetime, to_date: datetime) -> List[Dict]:
        """Search papers in specific category within date range."""
        date_range_query = f'submittedDate:[{from_date.strftime("%Y%m%d%H%M%S")}Z TO {to_date.strftime("%Y%m%d%H%M%S")}Z]'
        query = f'cat:{category} AND {date_range_query}'
        return self.search_papers(query)

    def search_by_author(self, author_name: str) -> List[Dict]:
        """Search papers by specific author."""
        query = f'au:"{author_name}"'
        return self.search_papers(query)

    def search_by_title(self, title_keywords: List[str]) -> List[Dict]:
        """Search papers by title keywords."""
        title_query = ' AND '.join([f'ti:"{keyword}"' for keyword in title_keywords])
        return self.search_papers(title_query)

    def get_trending_papers(self, category: str, days: int = 7) -> List[Dict]:
        """Get trending papers from recent submissions."""
        to_date = datetime.now()
        from_date = to_date - timedelta(days=days)
        return self.search_by_category(category, from_date, to_date)

    def search_similar_papers(self, arxiv_id: str) -> List[Dict]:
        """Find papers similar to given paper."""
        paper = self.get_paper_by_id(arxiv_id)
        if not paper:
            return []

        key_terms = self.extract_key_terms(paper["summary"])
        query = ' OR '.join(key_terms[:5])
        return self.search_papers(query)

    def get_paper_by_id(self, arxiv_id: str) -> Dict:
        """Get specific paper by ArXiv ID."""
        search = arxiv.Search(id_list=[arxiv_id])

        for result in self.client.results(search):
            return {
                "title": result.title,
                "authors": [author.name for author in result.authors],
                "summary": result.summary,
                "published": result.published,
                "arxiv_id": result.arxiv_id,
                "pdf_url": result.pdf_url,
                "categories": result.categories,
            }
        return None

    def extract_key_terms(self, text: str) -> List[str]:
        """Extract key terms from text."""
        try:
            from nltk.tokenize import sent_tokenize
            from nltk.corpus import stopwords

            words = text.lower().split()
            stop_words = set(stopwords.words('english'))

            key_terms = [
                word.strip('.,;:!?') for word in words
                if len(word) > 5 and word.lower() not in stop_words
            ]

            return list(set(key_terms))[:10]
        except:
            # Fallback if NLTK not available
            return text.split()[:10]

"""
Paper Content Processing Module

Handles downloading and extracting content from papers.
"""

import PyPDF2
import requests
from io import BytesIO
from typing import Dict, List


class PaperContentProcessor:
    """Processes paper content from PDFs."""

    def __init__(self):
        """Initialize processor."""
        self.cache = {}

    def download_and_process_paper(self, pdf_url: str, arxiv_id: str) -> Dict:
        """Download and extract content from paper."""
        try:
            pdf_content = self.download_pdf(pdf_url)
            text = self.extract_text_from_pdf(pdf_content)

            sections = self.parse_paper_structure(text)

            self.cache[arxiv_id] = {
                "text": text,
                "sections": sections,
                "citations": self.extract_citations(text)
            }

            return self.cache[arxiv_id]

        except Exception as e:
            print(f"Error processing paper: {e}")
            return None

    def download_pdf(self, pdf_url: str) -> BytesIO:
        """Download PDF from URL."""
        response = requests.get(pdf_url, timeout=10)
        response.raise_for_status()
        return BytesIO(response.content)

    def extract_text_from_pdf(self, pdf_content: BytesIO) -> str:
        """Extract text from PDF."""
        reader = PyPDF2.PdfReader(pdf_content)
        text = ""

        for page in reader.pages:
            text += page.extract_text()

        return text

    def parse_paper_structure(self, text: str) -> Dict:
        """Parse paper into sections."""
        sections = {
            "abstract": self.extract_section(text, "abstract"),
            "introduction": self.extract_section(text, "introduction"),
            "methodology": self.extract_section(text, "methodology|method|approach"),
            "results": self.extract_section(text, "results|findings"),
            "conclusion": self.extract_section(text, "conclusion|discussion"),
            "references": self.extract_section(text, "references|bibliography")
        }

        return sections

    def extract_section(self, text: str, section_keywords: str) -> str:
        """Extract specific section from paper."""
        import re

        pattern = f"(?i)({section_keywords})\\s*\\n"
        matches = list(re.finditer(pattern, text))

        if not matches:
            return ""

        start_pos = matches[0].end()

        section_pattern = r"(?i)(abstract|introduction|related work|methodology|method|results|conclusion|references|bibliography)\s*\n"
        next_matches = list(re.finditer(section_pattern, text[start_pos:]))

        if next_matches:
            end_pos = start_pos + next_matches[0].start()
        else:
            end_pos = len(text)

        return text[start_pos:end_pos].strip()

    def extract_citations(self, text: str) -> List[Dict]:
        """Extract citations from paper."""
        import re

        citations = []
        patterns = [
            r'\[(\d+)\]\s*(.+?)(?=\[|\Z)',
            r'(\w+\s+et\s+al\.?.*?\(\d{4}\))',
        ]

        for pattern in patterns:
            matches = re.finditer(pattern, text, re.IGNORECASE | re.DOTALL)
            for match in matches:
                citations.append({
                    "text": match.group(0)[:200],
                    "position": match.start()
                })

        return citations

    def get_cached_paper(self, arxiv_id: str) -> Dict:
        """Get cached paper content."""
        return self.cache.get(arxiv_id)

    def chunk_paper_for_rag(self, paper_content: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]:
        """Split paper into chunks for RAG."""
        chunks = []
        start = 0

        while start < len(paper_content):
            end = min(start + chunk_size, len(paper_content))
            chunks.append(paper_content[start:end])
            start = end - overlap

        return chunks

"""
Paper Question Answerer Module

Handles Q&A about research papers using RAG.
"""

from typing import Dict, List
import numpy as np


class PaperQuestionAnswerer:
    """Answers questions about research papers."""

    def __init__(self):
        """Initialize answerer."""
        self.retriever = None  # Will be initialized with ArXivPaperRetriever
        self.processor = None  # Will be initialized with PaperContentProcessor
        self.embeddings = None  # Will be initialized with embedding model
        self.llm = None  # Will be initialized with LLM

    def answer_question(self, question: str, top_k_papers: int = 5) -> Dict:
        """Answer question using papers from ArXiv."""
        # Step 1: Search for relevant papers
        papers = self.retriever.search_papers(question, max_results=top_k_papers)

        # Step 2: Process papers and chunk content
        paper_chunks = []
        paper_sources = []

        for paper in papers:
            try:
                content = self.processor.download_and_process_paper(
                    paper["pdf_url"],
                    paper["arxiv_id"]
                )

                if content:
                    chunks = self.processor.chunk_paper_for_rag(content["text"])
                    paper_chunks.extend(chunks)
                    paper_sources.append(paper)

            except Exception as e:
                print(f"Error processing {paper['arxiv_id']}: {e}")

        # Step 3: Retrieve most relevant chunks
        relevant_chunks = self.retrieve_relevant_chunks(question, paper_chunks)

        # Step 4: Generate answer
        context = "\n\n".join(relevant_chunks)
        answer = self.generate_answer(question, context)

        return {
            "question": question,
            "answer": answer,
            "sources": paper_sources,
            "relevant_chunks": relevant_chunks
        }

    def retrieve_relevant_chunks(self, query: str, chunks: List[str], top_k: int = 5) -> List[str]:
        """Retrieve most relevant chunks for query."""
        if not chunks or not self.embeddings:
            return chunks[:top_k]

        query_embedding = self.embeddings.encode(query)

        chunk_embeddings = [self.embeddings.encode(chunk) for chunk in chunks]

        similarities = [
            np.dot(query_embedding, chunk_emb) /
            (np.linalg.norm(query_embedding) * np.linalg.norm(chunk_emb))
            for chunk_emb in chunk_embeddings
        ]

        top_indices = np.argsort(similarities)[-top_k:][::-1]

        return [chunks[i] for i in top_indices if similarities[i] > 0.3]

    def generate_answer(self, question: str, context: str) -> str:
        """Generate answer using LLM."""
        prompt = f"""Based on the following research paper excerpts, answer the question.

Question: {question}

Context from research papers:
{context}

Answer: """

        if self.llm:
            answer = self.llm.generate(prompt, max_tokens=500)
            return answer.strip()
        else:
            return "LLM not configured"

    def answer_with_citations(self, question: str) -> Dict:
        """Answer question with proper citations."""
        result = self.answer_question(question)

        citations = self.format_citations(result["sources"])

        return {
            "answer": result["answer"],
            "citations": citations,
            "source_papers": result["sources"]
        }

    def format_citations(self, papers: List[Dict]) -> List[str]:
        """Format paper citations."""
        citations = []

        for paper in papers:
            authors = paper['authors'][:3] if paper['authors'] else ["Unknown"]
            citation = f"{', '.join(authors)} et al. ({paper['published'].year}). "
            citation += f"{paper['title']}. ArXiv:{paper['arxiv_id']}"
            citations.append(citation)

        return citations