enhanced_research_assistant.py

#!/usr/bin/env python3
"""
Enhanced Research Assistant with PDF Processing and Google Scholar Integration
"""

import asyncio
import os
import re
import requests
from pathlib import Path
from typing import List, Dict, Any, Optional
import json
import tempfile
import aiofiles
import xml.etree.ElementTree as ET
from urllib.parse import quote_plus
from datetime import datetime

# PDF processing
import PyPDF2
import pdfplumber

# Web scraping
from bs4 import BeautifulSoup
import time

# AutoGen imports
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.ui import Console
from autogen_ext.models.openai import OpenAIChatCompletionClient


class PaperProcessor:
    """Handles different types of paper inputs"""
    
    @staticmethod
    async def extract_text_from_pdf(pdf_path: str) -> str:
        """Extract text from PDF file"""
        try:
            text = ""
            with open(pdf_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                for page in pdf_reader.pages:
                    text += page.extract_text() + "\n"
            
            # If PyPDF2 fails, try pdfplumber
            if len(text.strip()) < 100:
                with pdfplumber.open(pdf_path) as pdf:
                    text = ""
                    for page in pdf.pages:
                        page_text = page.extract_text()
                        if page_text:
                            text += page_text + "\n"
            
            return text.strip()
        except Exception as e:
            return f"Error extracting PDF text: {str(e)}"
    
    @staticmethod
    async def download_arxiv_paper(arxiv_url: str) -> tuple[str, str]:
        """Download arXiv paper and extract text"""
        try:
            # Extract arXiv ID
            arxiv_id_match = re.search(r'(\d+\.\d+)', arxiv_url)
            if not arxiv_id_match:
                return "", "Invalid arXiv URL format"
            
            arxiv_id = arxiv_id_match.group(1)
            
            # Get paper metadata from arXiv API
            api_url = f"http://export.arxiv.org/api/query?id_list={arxiv_id}"
            response = requests.get(api_url)
            
            metadata = ""
            if response.status_code == 200:
                soup = BeautifulSoup(response.content, 'xml')
                entry = soup.find('entry')
                if entry:
                    title = entry.find('title').text.strip() if entry.find('title') else "Unknown"
                    authors = [author.find('name').text for author in entry.find_all('author')]
                    summary = entry.find('summary').text.strip() if entry.find('summary') else ""
                    
                    metadata = f"Title: {title}\nAuthors: {', '.join(authors)}\nAbstract: {summary}\n\n"
            
            # Download PDF
            pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
            pdf_response = requests.get(pdf_url)
            pdf_response.raise_for_status()
            
            # Save to temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                tmp_file.write(pdf_response.content)
                pdf_path = tmp_file.name
            
            # Extract text
            pdf_text = await PaperProcessor.extract_text_from_pdf(pdf_path)
            
            # Clean up
            os.unlink(pdf_path)
            
            return metadata + pdf_text, ""
            
        except Exception as e:
            return "", f"Error processing arXiv paper: {str(e)}"
    
    @staticmethod
    async def process_latex_content(latex_content: str) -> str:
        """Process LaTeX content and extract readable text"""
        # Remove common LaTeX commands
        text = re.sub(r'\\[a-zA-Z]+\{[^}]*\}', '', latex_content)
        text = re.sub(r'\\[a-zA-Z]+', '', text)
        text = re.sub(r'\{|\}', '', text)
        text = re.sub(r'%.*', '', text)  # Remove comments
        text = re.sub(r'\n\s*\n', '\n\n', text)  # Clean up whitespace
        
        return text.strip()
    
    @staticmethod
    async def process_paper(paper_input: str) -> dict:
        """Process different types of paper inputs and return structured data"""
        try:
            paper_data = {
                'input': paper_input,
                'type': '',
                'content': '',
                'metadata': {}
            }
            
            # Determine input type and process accordingly
            if paper_input.startswith('http') and 'arxiv.org' in paper_input:
                # arXiv URL
                paper_data['type'] = 'arxiv_url'
                content, error = await PaperProcessor.download_arxiv_paper(paper_input)
                if error:
                    paper_data['error'] = error
                    return paper_data
                paper_data['content'] = content
                
            elif paper_input.endswith('.pdf') and os.path.exists(paper_input):
                # PDF file path
                paper_data['type'] = 'pdf_file'
                content = await PaperProcessor.extract_text_from_pdf(paper_input)
                paper_data['content'] = content
                paper_data['metadata']['file_path'] = paper_input
                
            elif '\\documentclass' in paper_input or '\\begin{document}' in paper_input:
                # LaTeX content
                paper_data['type'] = 'latex_content'
                content = await PaperProcessor.process_latex_content(paper_input)
                paper_data['content'] = content
                
            else:
                # Plain text content
                paper_data['type'] = 'plain_text'
                paper_data['content'] = paper_input
            
            return paper_data
            
        except Exception as e:
            return {
                'input': paper_input,
                'type': 'error',
                'content': '',
                'error': str(e)
            }


class ArxivSearcher:
    """Handles arXiv API searches"""
    
    @staticmethod
    async def search_arxiv(query: str, num_results: int = 10) -> List[Dict[str, str]]:
        """Search arXiv for papers using the official API"""
        try:
            # arXiv API endpoint
            base_url = "http://export.arxiv.org/api/query?"
            search_query = quote_plus(query)
            url = f"{base_url}search_query=all:{search_query}&start=0&max_results={num_results}&sortBy=relevance&sortOrder=descending"
            
            response = requests.get(url)
            response.raise_for_status()
            
            # Parse XML response
            root = ET.fromstring(response.content)
            
            # Define namespace
            ns = {'atom': 'http://www.w3.org/2005/Atom',
                  'arxiv': 'http://arxiv.org/schemas/atom'}
            
            results = []
            for entry in root.findall('atom:entry', ns):
                title = entry.find('atom:title', ns).text.strip().replace('\n', ' ')
                
                # Get arXiv ID and create links
                arxiv_id = entry.find('atom:id', ns).text.split('/')[-1]
                arxiv_url = f"https://arxiv.org/abs/{arxiv_id}"
                pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
                
                # Get authors
                authors = []
                for author in entry.findall('atom:author', ns):
                    name = author.find('atom:name', ns).text
                    authors.append(name)
                authors_str = ", ".join(authors)
                
                # Get abstract
                summary = entry.find('atom:summary', ns).text.strip().replace('\n', ' ')
                
                # Get published date
                published = entry.find('atom:published', ns).text[:10]  # YYYY-MM-DD
                
                # Get categories
                categories = []
                for category in entry.findall('atom:category', ns):
                    categories.append(category.get('term'))
                categories_str = ", ".join(categories)
                
                results.append({
                    'title': title,
                    'arxiv_id': arxiv_id,
                    'arxiv_url': arxiv_url,
                    'pdf_url': pdf_url,
                    'authors': authors_str,
                    'abstract': summary[:300] + "..." if len(summary) > 300 else summary,
                    'published': published,
                    'categories': categories_str
                })
            
            return results
            
        except Exception as e:
            print(f"Error searching arXiv: {str(e)}")
            return []

class ScholarSearcher:
    """Handles Google Scholar searches"""
    
    @staticmethod
    async def search_google_scholar(query: str, num_results: int = 10) -> List[Dict[str, str]]:
        """Search Google Scholar for papers"""
        try:
            # Simple web scraping approach (in production, use scholarly library or API)
            search_url = f"https://scholar.google.com/scholar?q={query.replace(' ', '+')}&hl=en"
            
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
            }
            
            response = requests.get(search_url, headers=headers)
            soup = BeautifulSoup(response.content, 'html.parser')
            
            results = []
            for result in soup.find_all('div', class_='gs_r gs_or gs_scl')[:num_results]:
                title_elem = result.find('h3', class_='gs_rt')
                if title_elem:
                    title = title_elem.get_text()
                    link = title_elem.find('a')['href'] if title_elem.find('a') else ""
                    
                    # Extract authors and venue
                    meta_elem = result.find('div', class_='gs_a')
                    meta_text = meta_elem.get_text() if meta_elem else ""
                    
                    # Extract snippet
                    snippet_elem = result.find('div', class_='gs_rs')
                    snippet = snippet_elem.get_text() if snippet_elem else ""
                    
                    results.append({
                        'title': title,
                        'link': link,
                        'meta': meta_text,
                        'snippet': snippet
                    })
            
            return results
            
        except Exception as e:
            print(f"Error searching Google Scholar: {str(e)}")
            return []


class EnhancedResearchAssistant:
    def __init__(self, openai_api_key: str):
        self.model_client = OpenAIChatCompletionClient(
            model="gpt-4o",
            api_key=openai_api_key
        )
        self.paper_processor = PaperProcessor()
        self.scholar_searcher = ScholarSearcher()
        self.arxiv_searcher = ArxivSearcher()
        self.research_context = {}
        self.context_file = "research_context.json"
        self.setup_agents()
        
        # Try to load existing context
        self.load_context()
    
    def _serialize_context_safely(self, max_length: int = 2000) -> str:
        """Safely serialize research context for JSON, handling non-serializable objects"""
        try:
            # Create a safe copy of the context with only serializable data
            safe_context = {}
            for key, value in self.research_context.items():
                if isinstance(value, (str, int, float, bool, list, dict)):
                    safe_context[key] = value
                else:
                    # Convert complex objects to string representation
                    safe_context[key] = str(value)
            
            context_str = json.dumps(safe_context, indent=2)
            return context_str[:max_length] if len(context_str) > max_length else context_str
        except Exception as e:
            # Fallback to string representation
            return f"Research context summary: {str(self.research_context)[:max_length]}"
    
    def format_arxiv_results(self, results: List[Dict[str, str]]) -> str:
        """Format arXiv search results with clickable links"""
        if not results:
            return "No papers found."
        
        formatted_results = "📚 **Related Papers from arXiv:**\n\n"
        
        for i, paper in enumerate(results, 1):
            formatted_results += f"**{i}. {paper['title']}**\n"
            formatted_results += f"   👥 Authors: {paper['authors']}\n"
            formatted_results += f"   📅 Published: {paper['published']}\n"
            formatted_results += f"   🏷️ Categories: {paper['categories']}\n"
            formatted_results += f"   🔗 **Links:**\n"
            formatted_results += f"      • [View Abstract]({paper['arxiv_url']})\n"
            formatted_results += f"      • [Download PDF]({paper['pdf_url']})\n"
            formatted_results += f"   📝 Abstract: {paper['abstract']}\n\n"
        
        return formatted_results
    
    async def search_related_papers(self, query: str, num_results: int = 8) -> str:
        """Search for related papers and return formatted results with links"""
        try:
            print(f"🔍 Searching arXiv for: {query}")
            arxiv_results = await self.arxiv_searcher.search_arxiv(query, num_results)
            
            if arxiv_results:
                formatted_results = self.format_arxiv_results(arxiv_results)
                
                # Store results in context for future reference
                if 'related_papers' not in self.research_context:
                    self.research_context['related_papers'] = []
                
                self.research_context['related_papers'].extend(arxiv_results)
                
                return formatted_results
            else:
                return f"No papers found for query: {query}"
                
        except Exception as e:
            return f"Error searching for papers: {str(e)}"
    
    def save_context(self, filename: str = None) -> str:
        """Save the current research context to a JSON file"""
        try:
            if filename is None:
                filename = self.context_file
            
            # Create a safe copy of the context for saving
            safe_context = {}
            for key, value in self.research_context.items():
                if isinstance(value, (str, int, float, bool, list, dict)):
                    safe_context[key] = value
                else:
                    # Convert non-serializable objects to string representation
                    safe_context[key] = str(value)
            
            # Add metadata
            safe_context['_metadata'] = {
                'saved_at': datetime.now().isoformat(),
                'version': '1.0'
            }
            
            # Save to file
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(safe_context, f, indent=2, ensure_ascii=False)
            
            return f"✅ Research context saved to {filename}"
            
        except Exception as e:
            return f"❌ Error saving context: {str(e)}"
    
    def load_context(self, filename: str = None) -> str:
        """Load research context from a JSON file"""
        try:
            if filename is None:
                filename = self.context_file
            
            if not os.path.exists(filename):
                return f"No saved context found at {filename}"
            
            with open(filename, 'r', encoding='utf-8') as f:
                loaded_context = json.load(f)
            
            # Remove metadata before loading
            if '_metadata' in loaded_context:
                metadata = loaded_context.pop('_metadata')
                saved_at = metadata.get('saved_at', 'unknown')
                print(f"📂 Loading research context saved at {saved_at}")
            
            # Merge with existing context (in case there's already some data)
            self.research_context.update(loaded_context)
            
            return f"✅ Research context loaded from {filename}"
            
        except Exception as e:
            return f"❌ Error loading context: {str(e)}"
    
    def get_context_summary(self) -> str:
        """Get a summary of the current research context"""
        if not self.research_context:
            return "No research context available."
        
        summary = "📊 **Current Research Context Summary:**\n\n"
        
        # Count different types of data
        if 'processed_papers' in self.research_context:
            papers_count = len(self.research_context['processed_papers'])
            summary += f"📄 Processed Papers: {papers_count}\n"
        
        if 'related_papers' in self.research_context:
            related_count = len(self.research_context['related_papers'])
            summary += f"🔗 Related Papers Found: {related_count}\n"
        
        if 'analysis_results' in self.research_context:
            summary += f"🧠 Analysis Results: Available\n"
        
        if 'research_proposal' in self.research_context:
            summary += f"📝 Research Proposal: Generated\n"
        
        if 'literature_search' in self.research_context:
            summary += f"🔍 Literature Search: Completed\n"
        
        # Show available keys
        summary += f"\n🗂️ Available Data: {', '.join(self.research_context.keys())}\n"
        
        return summary
    
    async def analyze_aspect(self, aspect: str) -> str:
        """Analyze a specific aspect of the research"""
        try:
            context_str = self._serialize_context_safely(3000)
            
            analysis = await self.paper_analyzer.run(
                task=f"Provide detailed analysis of this aspect: {aspect}\n\nResearch Context:\n{context_str}"
            )
            
            result = str(analysis.messages[-1].content) if hasattr(analysis, 'messages') and analysis.messages else str(analysis)
            return result
            
        except Exception as e:
            return f"❌ Error analyzing aspect: {str(e)}"
    
    async def get_research_summary(self) -> str:
        """Get a comprehensive research summary"""
        try:
            context_str = self._serialize_context_safely(4000)
            
            summary = await self.research_synthesizer.run(
                task=f"Provide a comprehensive summary of our research exploration so far.\n\nContext: {context_str}"
            )
            
            result = str(summary.messages[-1].content) if hasattr(summary, 'messages') and summary.messages else str(summary)
            return result
            
        except Exception as e:
            return f"❌ Error generating summary: {str(e)}"
    
    async def investigate_authors(self, author_name: str) -> str:
        """Investigate specific authors"""
        try:
            context_str = self._serialize_context_safely(2000)
            
            investigation = await self.literature_searcher.run(
                task=f"Investigate author: {author_name}. Find their key papers, research areas, and collaborations.\n\nContext: {context_str}"
            )
            
            result = str(investigation.messages[-1].content) if hasattr(investigation, 'messages') and investigation.messages else str(investigation)
            return result
            
        except Exception as e:
            return f"❌ Error investigating authors: {str(e)}"
    
    async def refine_proposal(self) -> str:
        """Refine the research proposal"""
        try:
            current_proposal = str(self.research_context.get('proposal', ''))
            context_str = self._serialize_context_safely(2000)
            
            refined = await self.proposal_writer.run(
                task=f"Refine and improve this research proposal based on our current understanding:\n\nCurrent Proposal:\n{current_proposal}\n\nContext: {context_str}"
            )
            
            result = str(refined.messages[-1].content) if hasattr(refined, 'messages') and refined.messages else str(refined)
            
            # Update the proposal in context
            self.research_context['proposal'] = result
            
            return result
            
        except Exception as e:
            return f"❌ Error refining proposal: {str(e)}"
    
    async def explore_direction(self, direction: str) -> str:
        """Explore a specific research direction"""
        try:
            context_str = self._serialize_context_safely(2000)
            
            exploration = await self.research_synthesizer.run(
                task=f"Explore this research direction: {direction}\n\nProvide insights, potential approaches, and connections to our current research.\n\nContext: {context_str}"
            )
            
            result = str(exploration.messages[-1].content) if hasattr(exploration, 'messages') and exploration.messages else str(exploration)
            return result
            
        except Exception as e:
            return f"❌ Error exploring direction: {str(e)}"
    
    async def general_query(self, query: str) -> str:
        """Handle general research queries"""
        try:
            context_str = self._serialize_context_safely(2000)
            
            response = await self.research_synthesizer.run(
                task=f"Help with this research question: {query}\n\nContext: {context_str}"
            )
            
            result = str(response.messages[-1].content) if hasattr(response, 'messages') and response.messages else str(response)
            return result
            
        except Exception as e:
            return f"❌ Error processing query: {str(e)}"
    
    async def add_new_paper(self, paper_input: str) -> str:
        """Add and analyze a new paper during interactive session"""
        try:
            print(f"📄 Processing new paper: {paper_input[:100]}...")
            
            # Process the new paper
            processed_paper = await self.paper_processor.process_paper(paper_input)
            
            if not processed_paper or not processed_paper.get('content'):
                return "❌ Failed to process the paper. Please check the input format."
            
            # Initialize processed_papers list if it doesn't exist
            if 'processed_papers' not in self.research_context:
                self.research_context['processed_papers'] = []
            
            # Add to processed papers
            paper_id = len(self.research_context['processed_papers']) + 1
            processed_paper['id'] = paper_id
            self.research_context['processed_papers'].append(processed_paper)
            
            print(f"🧠 Analyzing new paper {paper_id}...")
            
            # Analyze the new paper
            analysis_result = await self.paper_analyzer.run(
                task=f"Analyze this research paper in detail:\n\n{processed_paper['content'][:3000]}..."
            )
            new_analysis = str(analysis_result.messages[-1].content) if hasattr(analysis_result, 'messages') and analysis_result.messages else str(analysis_result)
            
            # Update analysis in context
            if 'analysis' not in self.research_context:
                self.research_context['analysis'] = {}
            
            self.research_context['analysis'][f'paper_{paper_id}'] = new_analysis
            
            print(f"🔍 Searching for related literature...")
            
            # Search for related papers
            related_search = await self.literature_searcher.run(
                task=f"Based on this new paper analysis, suggest search strategies and find related work:\n\n{str(new_analysis)[:2000]}..."
            )
            related_papers = str(related_search.messages[-1].content) if hasattr(related_search, 'messages') and related_search.messages else str(related_search)
            
            # Update literature search in context
            if 'literature' not in self.research_context:
                self.research_context['literature'] = {}
            
            self.research_context['literature'][f'paper_{paper_id}'] = related_papers
            
            print(f"🧩 Updating research synthesis...")
            
            # Update synthesis with new paper
            current_synthesis = str(self.research_context.get('synthesis', ''))
            synthesis_update = await self.research_synthesizer.run(
                task=f"Update the research landscape synthesis with this new paper analysis:\n\nNew Paper Analysis:\n{str(new_analysis)[:2000]}...\n\nCurrent Synthesis:\n{current_synthesis[:2000]}..."
            )
            updated_synthesis = str(synthesis_update.messages[-1].content) if hasattr(synthesis_update, 'messages') and synthesis_update.messages else str(synthesis_update)
            
            self.research_context['synthesis'] = updated_synthesis
            
            print(f"📝 Updating research proposal...")
            
            # Update proposal with new insights
            current_proposal = str(self.research_context.get('proposal', ''))
            proposal_update = await self.proposal_writer.run(
                task=f"Update the research proposal with insights from this new paper:\n\nNew Paper Analysis:\n{str(new_analysis)[:2000]}...\n\nCurrent Proposal:\n{current_proposal[:2000]}..."
            )
            updated_proposal = str(proposal_update.messages[-1].content) if hasattr(proposal_update, 'messages') and proposal_update.messages else str(proposal_update)
            
            self.research_context['proposal'] = updated_proposal
            
            return f"✅ Successfully added and analyzed paper {paper_id}!\n📊 Research context updated with new insights\n🔄 Synthesis and proposal have been refreshed"
            
        except Exception as e:
            return f"❌ Error adding new paper: {str(e)}"
    
    def setup_agents(self):
        """Setup specialized research agents"""
        
        self.paper_analyzer = AssistantAgent(
            name="PaperAnalyzer",
            model_client=self.model_client,
            system_message="""You are an expert research paper analyzer. Analyze papers thoroughly and extract:

            1. **Basic Information**: Title, authors, venue, year
            2. **Research Context**: Problem being addressed, motivation, related work
            3. **Methodology**: Approach, techniques, experimental setup
            4. **Contributions**: Main contributions, novelty, significance
            5. **Results**: Key findings, performance metrics, comparisons
            6. **Limitations**: Acknowledged limitations, potential weaknesses
            7. **Future Work**: Suggested directions, open problems
            8. **Keywords**: Technical terms, research areas, methods

            Provide structured, detailed analysis that captures the essence of the research.
            Focus on understanding the research question, approach, and significance.
            """,
            description="Analyzes research papers comprehensively"
        )
        
        self.literature_searcher = AssistantAgent(
            name="LiteratureSearcher",
            model_client=self.model_client,
            system_message="""You are an expert at literature search and discovery. Your expertise includes:

            1. **Search Strategy**: Generate effective search queries for different databases
            2. **Related Work**: Identify papers that cite or are cited by the analyzed papers
            3. **Key Authors**: Find leading researchers and research groups in the field
            4. **Venues**: Identify top conferences and journals in the research area
            5. **Trends**: Spot emerging trends and hot topics
            6. **Gaps**: Identify under-explored areas and research opportunities

            Provide comprehensive search strategies and explain the reasoning behind them.
            Suggest both broad surveys and specific technical papers.
            """,
            description="Expert at finding and organizing related literature"
        )
        
        self.research_synthesizer = AssistantAgent(
            name="ResearchSynthesizer",
            model_client=self.model_client,
            system_message="""You are an expert research synthesizer who understands the big picture. Your role:

            1. **Field Analysis**: Map out the current state of the research field
            2. **Problem Importance**: Assess why research problems matter and their impact
            3. **Approach Comparison**: Compare different methodological approaches
            4. **Gap Identification**: Find unexplored areas and research opportunities
            5. **Trend Analysis**: Identify emerging directions and future possibilities
            6. **Critical Assessment**: Evaluate strengths and weaknesses of current work

            Think like a senior researcher who can see connections across papers and identify
            the most promising research directions. Focus on impact and significance.
            """,
            description="Synthesizes research landscape and identifies opportunities"
        )
        
        self.proposal_writer = AssistantAgent(
            name="ProposalWriter",
            model_client=self.model_client,
            system_message="""You are an expert research proposal writer. Create compelling, well-structured proposals:

            1. **Problem Statement**: Clear, motivated problem definition
            2. **Literature Review**: Concise summary of relevant work and gaps
            3. **Research Questions**: Specific, answerable research questions
            4. **Methodology**: Detailed approach with justification
            5. **Expected Outcomes**: Clear contributions and impact
            6. **Timeline**: Realistic milestones and deliverables
            7. **References**: Comprehensive, well-organized bibliography

            Write proposals that are convincing, feasible, and impactful.
            Make sure the research questions are novel and the approach is sound.
            """,
            description="Writes structured research proposals"
        )

    async def process_papers(self, paper_inputs: List[str]) -> List[Dict[str, Any]]:
        """Process different types of paper inputs"""
        processed_papers = []
        
        for i, paper_input in enumerate(paper_inputs):
            print(f"📄 Processing paper {i+1}/{len(paper_inputs)}...")
            
            paper_data = {"input": paper_input, "content": "", "error": ""}
            
            try:
                if "arxiv.org" in paper_input:
                    content, error = await self.paper_processor.download_arxiv_paper(paper_input)
                    paper_data["content"] = content
                    paper_data["error"] = error
                    paper_data["type"] = "arxiv"
                    
                elif paper_input.endswith('.pdf'):
                    if os.path.exists(paper_input):
                        content = await self.paper_processor.extract_text_from_pdf(paper_input)
                        paper_data["content"] = content
                        paper_data["type"] = "pdf"
                    else:
                        paper_data["error"] = f"PDF file not found: {paper_input}"
                        
                elif '\\documentclass' in paper_input or '\\begin{document}' in paper_input:
                    content = await self.paper_processor.process_latex_content(paper_input)
                    paper_data["content"] = content
                    paper_data["type"] = "latex"
                    
                else:
                    paper_data["content"] = paper_input
                    paper_data["type"] = "text"
                    
            except Exception as e:
                paper_data["error"] = str(e)
            
            processed_papers.append(paper_data)
        
        return processed_papers

    async def analyze_papers(self, processed_papers: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Analyze processed papers"""
        print("🔍 Analyzing papers...")
        
        analyses = []
        for i, paper in enumerate(processed_papers):
            if paper["error"]:
                print(f"⚠️  Skipping paper {i+1} due to error: {paper['error']}")
                continue
            
            print(f"🧠 Analyzing paper {i+1}...")
            
            analysis_prompt = f"""
            Please provide a comprehensive analysis of this research paper:
            
            {paper['content'][:8000]}  # Limit content to avoid token limits
            
            Provide detailed analysis covering all aspects mentioned in your system message.
            Focus on understanding the research contribution and its significance.
            """
            
            try:
                analysis_result = await self.paper_analyzer.run(task=analysis_prompt)
                # Extract the actual text content from the TaskResult
                analysis_text = str(analysis_result.messages[-1].content) if hasattr(analysis_result, 'messages') and analysis_result.messages else str(analysis_result)
                analyses.append({
                    "paper_index": i,
                    "analysis": analysis_text,
                    "paper_type": paper["type"]
                })
            except Exception as e:
                print(f"❌ Error analyzing paper {i+1}: {str(e)}")
        
        return {"analyses": analyses, "total_papers": len(processed_papers)}

    async def search_related_work(self, analysis_results: Dict[str, Any]) -> Dict[str, Any]:
        """Search for related work based on analysis"""
        print("🔎 Searching for related literature...")
        
        # Generate search queries based on analyses
        search_prompt = f"""
        Based on these paper analyses, generate comprehensive search strategies:
        
        {json.dumps(analysis_results, indent=2)[:4000]}
        
        Provide:
        1. Specific search queries for Google Scholar
        2. Key authors and research groups to investigate
        3. Important venues and conferences
        4. Related research areas and keywords
        5. Potential survey papers and seminal works
        """
        
        search_result = await self.literature_searcher.run(task=search_prompt)
        search_strategy = str(search_result.messages[-1].content) if hasattr(search_result, 'messages') and search_result.messages else str(search_result)
        
        # Extract search queries and perform actual searches
        # For demo purposes, we'll simulate this
        related_papers = []
        
        # In a real implementation, you would:
        # 1. Extract search queries from the strategy
        # 2. Perform actual Google Scholar searches
        # 3. Process and rank the results
        
        return {
            "search_strategy": search_strategy,
            "related_papers": related_papers
        }

    async def synthesize_research_landscape(self, analysis_results: Dict[str, Any], 
                                          literature_search: Dict[str, Any]) -> Dict[str, Any]:
        """Synthesize the overall research landscape"""
        print("🧩 Synthesizing research landscape...")
        
        synthesis_prompt = f"""
        Synthesize the research landscape based on the paper analyses and literature search:
        
        Paper Analyses: {json.dumps(analysis_results, indent=2)[:3000]}
        Literature Search: {json.dumps(literature_search, indent=2)[:3000]}
        
        Provide comprehensive synthesis covering:
        1. Current state of the field - what's being studied and why
        2. Key research questions and their importance
        3. Main approaches and their trade-offs
        4. Significant findings and breakthroughs
        5. Open problems and research gaps
        6. Emerging trends and future directions
        7. Potential high-impact research opportunities
        
        Think like a senior researcher mapping out the field.
        """
        
        synthesis_result = await self.research_synthesizer.run(task=synthesis_prompt)
        synthesis = str(synthesis_result.messages[-1].content) if hasattr(synthesis_result, 'messages') and synthesis_result.messages else str(synthesis_result)
        return {"synthesis": synthesis}

    async def generate_research_proposal(self, synthesis_results: Dict[str, Any]) -> Dict[str, Any]:
        """Generate research proposal outline"""
        print("📝 Generating research proposal...")
        
        proposal_prompt = f"""
        Create a comprehensive research proposal based on this synthesis:
        
        {json.dumps(synthesis_results, indent=2)[:4000]}
        
        Generate a detailed proposal with:
        
        1. **Executive Summary**
        2. **Problem Statement and Motivation**
        3. **Literature Review and Gap Analysis**
        4. **Research Objectives and Questions**
        5. **Proposed Methodology**
        6. **Expected Contributions and Impact**
        7. **Timeline and Milestones**
        8. **Resource Requirements**
        9. **Risk Assessment**
        10. **Comprehensive Reference List**
        
        Make it compelling, feasible, and impactful.
        """
        
        proposal_result = await self.proposal_writer.run(task=proposal_prompt)
        proposal = str(proposal_result.messages[-1].content) if hasattr(proposal_result, 'messages') and proposal_result.messages else str(proposal_result)
        return {"proposal": proposal}

    async def interactive_mode(self):
        """Interactive exploration mode"""
        print("\n" + "="*60)
        print("🤖 INTERACTIVE RESEARCH EXPLORATION MODE")
        print("="*60)
        
        # Show context status
        if self.research_context:
            print(f"📂 Research context loaded with {len(self.research_context)} data items")
        else:
            print("📂 Starting with empty research context")
        print("="*60)
        print("Available commands:")
        print("• 'search [topic]' - Find papers with clickable links from arXiv")
        print("• 'add [paper]' - Add and analyze a new paper (arXiv URL, PDF path, or text)")
        print("• 'analyze [aspect]' - Deep dive into specific aspects")
        print("• 'refine proposal' - Refine the research proposal")
        print("• 'explore [direction]' - Explore research directions")
        print("• 'authors [name]' - Investigate specific authors")
        print("• 'summary' - Get summary of current research context")
        print("• 'save' - Save current research context")
        print("• 'load [filename]' - Load saved research context")
        print("• 'context' - Show detailed context summary")
        print("• 'export' - Export research findings")
        print("• 'exit' - End session")
        print("="*60)
        
        while True:
            try:
                user_input = input("\n💭 Your request: ").strip()
                
                if user_input.lower() in ['exit', 'quit', 'done']:
                    # Auto-save context before exiting
                    if self.research_context:
                        save_result = self.save_context()
                        print(f"💾 {save_result}")
                    print("👋 Research session completed!")
                    break
                
                if not user_input:
                    continue
                
                # Route requests to appropriate agents
                if user_input.lower().startswith('search'):
                    # Extract search query from user input
                    search_query = user_input[6:].strip()  # Remove 'search' prefix
                    if not search_query:
                        response = "Please provide a search query. Example: 'search machine learning transformers'"
                    else:
                        # Get actual paper links from arXiv
                        paper_results = await self.search_related_papers(search_query)
                        
                        # Also get AI analysis of the search
                        response_result = await self.literature_searcher.run(
                            task=f"Provide search strategy and analysis for: {search_query}\n\nContext: {self._serialize_context_safely(2000)}"
                        )
                        ai_analysis = str(response_result.messages[-1].content) if hasattr(response_result, 'messages') and response_result.messages else str(response_result)
                        
                        # Combine AI analysis with actual paper links
                        response = f"{ai_analysis}\n\n{paper_results}"
                elif user_input.lower().startswith('add'):
                    # Extract paper input
                    paper_input = user_input[3:].strip()  # Remove 'add' prefix
                    if paper_input:
                        response = await self.add_new_paper(paper_input)
                    else:
                        response = "Please provide a paper to add. Example: add https://arxiv.org/abs/2301.00001"
                elif user_input.lower().startswith('refine') or 'proposal' in user_input.lower():
                    response_result = await self.proposal_writer.run(
                        task=f"Help refine the research proposal: {user_input}\n\nContext: {self._serialize_context_safely(2000)}"
                    )
                    response = str(response_result.messages[-1].content) if hasattr(response_result, 'messages') and response_result.messages else str(response_result)
                elif user_input.lower().startswith('analyze'):
                    response_result = await self.paper_analyzer.run(
                        task=f"Provide detailed analysis: {user_input}\n\nContext: {self._serialize_context_safely(2000)}"
                    )
                    response = str(response_result.messages[-1].content) if hasattr(response_result, 'messages') and response_result.messages else str(response_result)
                elif user_input.lower() == 'summary':
                    response_result = await self.research_synthesizer.run(
                        task=f"Provide a comprehensive summary of our research exploration so far.\n\nContext: {self._serialize_context_safely(3000)}"
                    )
                    response = str(response_result.messages[-1].content) if hasattr(response_result, 'messages') and response_result.messages else str(response_result)
                elif user_input.lower() == 'save':
                    response = self.save_context()
                elif user_input.lower().startswith('load'):
                    # Extract filename if provided
                    parts = user_input.split(maxsplit=1)
                    filename = parts[1] if len(parts) > 1 else None
                    response = self.load_context(filename)
                elif user_input.lower() == 'context':
                    response = self.get_context_summary()
                else:
                    # Default to synthesizer for general questions
                    response_result = await self.research_synthesizer.run(
                        task=f"Help with this research question: {user_input}\n\nContext: {self._serialize_context_safely(2000)}"
                    )
                    response = str(response_result.messages[-1].content) if hasattr(response_result, 'messages') and response_result.messages else str(response_result)
                
                print(f"\n🤖 **Response:**\n{response}")
                
            except KeyboardInterrupt:
                print("\n👋 Session interrupted. Goodbye!")
                break
            except Exception as e:
                print(f"❌ Error: {str(e)}")

    async def run_complete_workflow(self, paper_inputs: List[str]):
        """Run the complete research assistant workflow"""
        print("🚀 ADVANCED RESEARCH ASSISTANT")
        print("="*60)
        
        try:
            # Step 1: Process paper inputs
            processed_papers = await self.process_papers(paper_inputs)
            self.research_context['processed_papers'] = processed_papers
            
            # Step 2: Analyze papers
            analysis_results = await self.analyze_papers(processed_papers)
            self.research_context['analysis'] = analysis_results
            
            # Step 3: Search related literature
            literature_search = await self.search_related_work(analysis_results)
            self.research_context['literature'] = literature_search
            
            # Step 4: Synthesize research landscape
            synthesis = await self.synthesize_research_landscape(analysis_results, literature_search)
            self.research_context['synthesis'] = synthesis
            
            # Step 5: Generate research proposal
            proposal = await self.generate_research_proposal(synthesis)
            self.research_context['proposal'] = proposal
            
            print("\n" + "="*60)
            print("✅ INITIAL RESEARCH WORKFLOW COMPLETED!")
            print("📊 Research context built successfully")
            print("📝 Research proposal generated")
            print("="*60)
            
            # Step 6: Interactive exploration
            await self.interactive_mode()
            
        except Exception as e:
            print(f"❌ Error in workflow: {str(e)}")
        finally:
            await self.model_client.close()


async def main():
    """Main function"""
    
    # Check API key
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("❌ OPENAI_API_KEY not found!")
        print("Set it with: export OPENAI_API_KEY='your-key-here'")
        return
    
    print("🎓 ENHANCED RESEARCH ASSISTANT")
    print("="*50)
    print("Capabilities:")
    print("✅ Process arXiv papers, PDFs, LaTeX, and text")
    print("✅ Comprehensive paper analysis")
    print("✅ Literature search and discovery")
    print("✅ Research landscape synthesis")
    print("✅ Research proposal generation")
    print("✅ Interactive exploration and refinement")
    print("="*50)
    
    # Initialize the research assistant first to check for existing context
    assistant = EnhancedResearchAssistant(api_key)
    
    # Check if we have existing context
    has_existing_context = bool(assistant.research_context)
    
    # Collect paper inputs
    papers = []
    print("\n📚 Please provide your research papers:")
    print("(arXiv URLs, PDF file paths, LaTeX code, or paper text)")
    if has_existing_context:
        print("(Type 'done' to skip and use existing context, or 'skip' to go directly to interactive mode)")
    else:
        print("(Type 'done' when finished)")
    
    while True:
        paper_input = input(f"\nPaper {len(papers) + 1} (or 'done'): ").strip()
        
        if paper_input.lower() in ['done', 'skip']:
            break
        
        if paper_input:
            papers.append(paper_input)
            print(f"✅ Added paper {len(papers)}")
    
    # Handle different scenarios
    if papers:
        print(f"\n🎯 Starting analysis of {len(papers)} paper(s)...")
        await assistant.run_complete_workflow(papers)
    elif has_existing_context:
        print("\n🔄 Using existing research context...")
        print("📂 Loaded research context with existing data")
        await assistant.interactive_mode()
    else:
        print("❌ No papers provided and no existing context found. Exiting.")
        return


if __name__ == "__main__":
    asyncio.run(main())