Research Lookup

#!/usr/bin/env python3
"""
Example usage of the Research Lookup skill with automatic model selection.

This script demonstrates:
1. Automatic model selection based on query complexity
2. Manual model override options
3. Batch query processing
4. Integration with scientific writing workflows
"""

import os
from research_lookup import ResearchLookup


def example_automatic_selection():
    """Demonstrate automatic model selection."""
    print("=" * 80)
    print("EXAMPLE 1: Automatic Model Selection")
    print("=" * 80)
    print()
    
    research = ResearchLookup()
    
    # Simple lookup - will use Sonar Pro Search
    query1 = "Recent advances in CRISPR gene editing 2024"
    print(f"Query: {query1}")
    print(f"Expected model: Sonar Pro Search (fast lookup)")
    result1 = research.lookup(query1)
    print(f"Actual model: {result1.get('model')}")
    print()
    
    # Complex analysis - will use Sonar Reasoning Pro
    query2 = "Compare and contrast the efficacy of mRNA vaccines versus traditional vaccines"
    print(f"Query: {query2}")
    print(f"Expected model: Sonar Reasoning Pro (analytical)")
    result2 = research.lookup(query2)
    print(f"Actual model: {result2.get('model')}")
    print()


def example_manual_override():
    """Demonstrate manual model override."""
    print("=" * 80)
    print("EXAMPLE 2: Manual Model Override")
    print("=" * 80)
    print()
    
    # Force Sonar Pro Search for budget-constrained rapid lookup
    research_pro = ResearchLookup(force_model='pro')
    query = "Explain the mechanism of CRISPR-Cas9"
    print(f"Query: {query}")
    print(f"Forced model: Sonar Pro Search")
    result = research_pro.lookup(query)
    print(f"Model used: {result.get('model')}")
    print()
    
    # Force Sonar Reasoning Pro for critical analysis
    research_reasoning = ResearchLookup(force_model='reasoning')
    print(f"Query: {query}")
    print(f"Forced model: Sonar Reasoning Pro")
    result = research_reasoning.lookup(query)
    print(f"Model used: {result.get('model')}")
    print()


def example_batch_queries():
    """Demonstrate batch query processing."""
    print("=" * 80)
    print("EXAMPLE 3: Batch Query Processing")
    print("=" * 80)
    print()
    
    research = ResearchLookup()
    
    # Mix of simple and complex queries
    queries = [
        "Recent clinical trials for Alzheimer's disease",  # Sonar Pro Search
        "Compare deep learning vs traditional ML in drug discovery",  # Sonar Reasoning Pro
        "Statistical power analysis methods",  # Sonar Pro Search
    ]
    
    print("Processing batch queries...")
    print("Each query will automatically select the appropriate model")
    print()
    
    results = research.batch_lookup(queries, delay=1.0)
    
    for i, result in enumerate(results):
        print(f"Query {i+1}: {result['query'][:50]}...")
        print(f"  Model: {result.get('model')}")
        print(f"  Type: {result.get('model_type')}")
        print()


def example_scientific_writing_workflow():
    """Demonstrate integration with scientific writing workflow."""
    print("=" * 80)
    print("EXAMPLE 4: Scientific Writing Workflow")
    print("=" * 80)
    print()
    
    research = ResearchLookup()
    
    # Literature review phase - use Pro for breadth
    print("PHASE 1: Literature Review (Breadth)")
    lit_queries = [
        "Recent papers on machine learning in genomics 2024",
        "Clinical applications of AI in radiology",
        "RNA sequencing analysis methods"
    ]
    
    for query in lit_queries:
        print(f"  - {query}")
        # These will automatically use Sonar Pro Search
    print()
    
    # Discussion phase - use Reasoning Pro for synthesis
    print("PHASE 2: Discussion (Synthesis & Analysis)")
    discussion_queries = [
        "Compare the advantages and limitations of different ML approaches in genomics",
        "Explain the relationship between model interpretability and clinical adoption",
        "Analyze the ethical implications of AI in medical diagnosis"
    ]
    
    for query in discussion_queries:
        print(f"  - {query}")
        # These will automatically use Sonar Reasoning Pro
    print()


def main():
    """Run all examples (requires OPENROUTER_API_KEY to be set)."""
    
    if not os.getenv("OPENROUTER_API_KEY"):
        print("Note: Set OPENROUTER_API_KEY environment variable to run live queries")
        print("These examples show the structure without making actual API calls")
        print()
    
    # Uncomment to run examples (requires API key)
    # example_automatic_selection()
    # example_manual_override()
    # example_batch_queries()
    # example_scientific_writing_workflow()
    
    # Show complexity assessment without API calls
    print("=" * 80)
    print("COMPLEXITY ASSESSMENT EXAMPLES (No API calls required)")
    print("=" * 80)
    print()
    
    os.environ.setdefault("OPENROUTER_API_KEY", "test")
    research = ResearchLookup()
    
    test_queries = [
        ("Recent CRISPR studies", "pro"),
        ("Compare CRISPR vs TALENs", "reasoning"),
        ("Explain how CRISPR works", "reasoning"),
        ("Western blot protocol", "pro"),
        ("Pros and cons of different sequencing methods", "reasoning"),
    ]
    
    for query, expected in test_queries:
        complexity = research._assess_query_complexity(query)
        model_name = "Sonar Reasoning Pro" if complexity == "reasoning" else "Sonar Pro Search"
        status = "✓" if complexity == expected else "✗"
        print(f"{status} '{query}'")
        print(f"  → {model_name}")
        print()


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Research Lookup Tool for Claude Code
Performs research queries using Perplexity Sonar Pro Search via OpenRouter.
"""

import os
import sys
import json
from typing import Dict, List, Optional

# Import the main research lookup class
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'scripts'))
from research_lookup import ResearchLookup


def format_response(result: Dict) -> str:
    """Format the research result for display."""
    if not result["success"]:
        return f"❌ Research lookup failed: {result['error']}"

    response = result["response"]
    citations = result["citations"]
    sources = result.get("sources", [])

    # Format the output for Claude Code
    output = f"""🔍 **Research Results**

**Query:** {result['query']}
**Model:** {result['model']}
**Timestamp:** {result['timestamp']}
**Note:** Results prioritized by citation count, venue prestige, and author reputation

---

{response}

"""

    # Display API-provided sources with venue/citation info
    if sources:
        output += f"\n📚 **Sources ({len(sources)}):**\n"
        output += "_Prioritized by venue quality and citation impact_\n\n"
        for i, source in enumerate(sources, 1):
            title = source.get("title", "Untitled")
            url = source.get("url", "")
            date = source.get("date", "")
            snippet = source.get("snippet", "")
            
            # Format source entry with available metadata
            date_str = f" ({date})" if date else ""
            output += f"{i}. **{title}**{date_str}\n"
            
            # Add venue indicator if detectable from URL
            venue_indicator = _detect_venue_tier(url)
            if venue_indicator:
                output += f"   📊 Venue: {venue_indicator}\n"
            
            if url:
                output += f"   🔗 {url}\n"
            if snippet:
                output += f"   _{snippet[:150]}{'...' if len(snippet) > 150 else ''}_\n"
            output += "\n"

    # Display extracted citations (DOIs, etc.)
    if citations:
        doi_citations = [c for c in citations if c.get("type") == "doi"]
        url_citations = [c for c in citations if c.get("type") == "url"]
        
        if doi_citations:
            output += f"\n🔗 **DOI References ({len(doi_citations)}):**\n"
            for i, citation in enumerate(doi_citations, 1):
                output += f"{i}. DOI: {citation.get('doi', '')} → {citation.get('url', '')}\n"
        
        if url_citations:
            output += f"\n🌐 **Additional URLs ({len(url_citations)}):**\n"
            for i, citation in enumerate(url_citations, 1):
                url = citation.get('url', '')
                venue = _detect_venue_tier(url)
                venue_str = f" [{venue}]" if venue else ""
                output += f"{i}. {url}{venue_str}\n"

    if result.get("usage"):
        usage = result["usage"]
        output += f"\n**Usage:** {usage.get('total_tokens', 'N/A')} tokens"

    return output


def _detect_venue_tier(url: str) -> Optional[str]:
    """Detect venue tier from URL to indicate source quality."""
    if not url:
        return None
    
    url_lower = url.lower()
    
    # Tier 1 - Premier venues
    tier1_indicators = {
        "nature.com": "Nature (Tier 1)",
        "science.org": "Science (Tier 1)",
        "cell.com": "Cell Press (Tier 1)",
        "nejm.org": "NEJM (Tier 1)",
        "thelancet.com": "Lancet (Tier 1)",
        "jamanetwork.com": "JAMA (Tier 1)",
        "pnas.org": "PNAS (Tier 1)",
    }
    
    # Tier 2 - High-impact specialized
    tier2_indicators = {
        "neurips.cc": "NeurIPS (Tier 2 - Top ML)",
        "icml.cc": "ICML (Tier 2 - Top ML)",
        "openreview.net": "Top ML Conference (Tier 2)",
        "aacrjournals.org": "AACR Journals (Tier 2)",
        "ahajournals.org": "AHA Journals (Tier 2)",
        "bloodjournal.org": "Blood (Tier 2)",
        "jci.org": "JCI (Tier 2)",
    }
    
    # Tier 3 - Respected academic sources
    tier3_indicators = {
        "springer.com": "Springer",
        "wiley.com": "Wiley",
        "elsevier.com": "Elsevier",
        "oup.com": "Oxford University Press",
        "arxiv.org": "arXiv (Preprint)",
        "biorxiv.org": "bioRxiv (Preprint)",
        "medrxiv.org": "medRxiv (Preprint)",
        "pubmed": "PubMed",
        "ncbi.nlm.nih.gov": "NCBI/PubMed",
        "ieee.org": "IEEE",
        "acm.org": "ACM",
    }
    
    for domain, label in tier1_indicators.items():
        if domain in url_lower:
            return label
    
    for domain, label in tier2_indicators.items():
        if domain in url_lower:
            return label
    
    for domain, label in tier3_indicators.items():
        if domain in url_lower:
            return label
    
    return None


def main():
    """Main entry point for Claude Code tool."""
    # Check for API key
    if not os.getenv("OPENROUTER_API_KEY"):
        print("❌ Error: OPENROUTER_API_KEY environment variable not set")
        print("Please set it in your .env file or export it:")
        print("  export OPENROUTER_API_KEY='your_openrouter_api_key'")
        return 1

    # Get query from command line arguments
    if len(sys.argv) < 2:
        print("❌ Error: No query provided")
        print("Usage: python lookup.py 'your research query here'")
        return 1

    query = " ".join(sys.argv[1:])

    try:
        # Initialize research tool
        research = ResearchLookup()

        # Perform lookup
        print(f"🔍 Researching: {query}")
        result = research.lookup(query)

        # Format and output result
        formatted_output = format_response(result)
        print(formatted_output)

        # Return success code
        return 0 if result["success"] else 1

    except Exception as e:
        print(f"❌ Error: {str(e)}")
        return 1


if __name__ == "__main__":
    exit(main())

#!/usr/bin/env python3
"""
Research Information Lookup Tool

Routes research queries to the best backend:
  - Parallel Chat API (core model): Default for all general research queries
  - Perplexity sonar-pro-search (via OpenRouter): Academic-specific paper searches

Environment variables:
  PARALLEL_API_KEY    - Required for Parallel Chat API (primary backend)
  OPENROUTER_API_KEY  - Required for Perplexity academic searches (fallback)
"""

import os
import sys
import json
import re
import time
import requests
from datetime import datetime
from typing import Any, Dict, List, Optional


class ResearchLookup:
    """Research information lookup with intelligent backend routing.

    Routes queries to the Parallel Chat API (default) or Perplexity
    sonar-pro-search (academic paper searches only).
    """

    ACADEMIC_KEYWORDS = [
        "find papers", "find paper", "find articles", "find article",
        "cite ", "citation", "citations for",
        "doi ", "doi:", "pubmed", "pmid",
        "journal article", "peer-reviewed",
        "systematic review", "meta-analysis",
        "literature search", "literature on",
        "academic papers", "academic paper",
        "research papers on", "research paper on",
        "published studies", "published study",
        "scholarly", "scholar",
        "arxiv", "preprint",
        "foundational papers", "seminal papers", "landmark papers",
        "highly cited", "most cited",
    ]

    PARALLEL_SYSTEM_PROMPT = (
        "You are a deep research analyst. Provide a comprehensive, well-cited "
        "research report on the user's topic. Include:\n"
        "- Key findings with specific data, statistics, and quantitative evidence\n"
        "- Detailed analysis organized by themes\n"
        "- Multiple authoritative sources cited inline\n"
        "- Methodologies and implications where relevant\n"
        "- Future outlook and research gaps\n"
        "Use markdown formatting with clear section headers. "
        "Prioritize authoritative and recent sources."
    )

    CHAT_BASE_URL = "https://api.parallel.ai"

    def __init__(self, force_backend: Optional[str] = None):
        """Initialize the research lookup tool.

        Args:
            force_backend: Force a specific backend ('parallel' or 'perplexity').
                          If None, backend is auto-selected based on query content.
        """
        self.force_backend = force_backend
        self.parallel_available = bool(os.getenv("PARALLEL_API_KEY"))
        self.perplexity_available = bool(os.getenv("OPENROUTER_API_KEY"))

        if not self.parallel_available and not self.perplexity_available:
            raise ValueError(
                "No API keys found. Set at least one of:\n"
                "  PARALLEL_API_KEY (for Parallel Chat API - primary)\n"
                "  OPENROUTER_API_KEY (for Perplexity academic search - fallback)"
            )

    def _select_backend(self, query: str) -> str:
        """Select the best backend for a query."""
        if self.force_backend:
            if self.force_backend == "perplexity" and self.perplexity_available:
                return "perplexity"
            if self.force_backend == "parallel" and self.parallel_available:
                return "parallel"

        query_lower = query.lower()
        is_academic = any(kw in query_lower for kw in self.ACADEMIC_KEYWORDS)

        if is_academic and self.perplexity_available:
            return "perplexity"

        if self.parallel_available:
            return "parallel"

        if self.perplexity_available:
            return "perplexity"

        raise ValueError("No backend available. Check API keys.")

    # ------------------------------------------------------------------
    # Parallel Chat API backend
    # ------------------------------------------------------------------

    def _get_chat_client(self):
        """Lazy-load and cache the OpenAI client for Parallel Chat API."""
        if not hasattr(self, "_chat_client"):
            try:
                from openai import OpenAI
            except ImportError:
                raise ImportError(
                    "The 'openai' package is required for Parallel Chat API.\n"
                    "Install it with: pip install openai"
                )
            self._chat_client = OpenAI(
                api_key=os.getenv("PARALLEL_API_KEY"),
                base_url=self.CHAT_BASE_URL,
            )
        return self._chat_client

    def _parallel_lookup(self, query: str) -> Dict[str, Any]:
        """Run research via the Parallel Chat API (core model)."""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        model = "core"

        try:
            client = self._get_chat_client()

            print(f"[Research] Parallel Chat API (model={model})...", file=sys.stderr)

            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": self.PARALLEL_SYSTEM_PROMPT},
                    {"role": "user", "content": query},
                ],
                stream=False,
            )

            content = ""
            if response.choices and len(response.choices) > 0:
                content = response.choices[0].message.content or ""

            api_citations = self._extract_basis_citations(response)
            text_citations = self._extract_citations_from_text(content)

            return {
                "success": True,
                "query": query,
                "response": content,
                "citations": api_citations + text_citations,
                "sources": api_citations,
                "timestamp": timestamp,
                "backend": "parallel",
                "model": f"parallel-chat/{model}",
            }

        except Exception as e:
            return {
                "success": False,
                "query": query,
                "error": str(e),
                "timestamp": timestamp,
                "backend": "parallel",
                "model": f"parallel-chat/{model}",
            }

    def _extract_basis_citations(self, response) -> List[Dict[str, str]]:
        """Extract citation sources from the Chat API research basis."""
        citations = []
        basis = getattr(response, "basis", None)
        if not basis:
            return citations

        seen_urls = set()
        if isinstance(basis, list):
            for item in basis:
                cits = (
                    item.get("citations", []) if isinstance(item, dict)
                    else getattr(item, "citations", None) or []
                )
                for cit in cits:
                    url = cit.get("url", "") if isinstance(cit, dict) else getattr(cit, "url", "")
                    if url and url not in seen_urls:
                        seen_urls.add(url)
                        title = cit.get("title", "") if isinstance(cit, dict) else getattr(cit, "title", "")
                        excerpts = cit.get("excerpts", []) if isinstance(cit, dict) else getattr(cit, "excerpts", [])
                        citations.append({
                            "type": "source",
                            "url": url,
                            "title": title,
                            "excerpts": excerpts,
                        })

        return citations

    # ------------------------------------------------------------------
    # Perplexity academic search backend
    # ------------------------------------------------------------------

    def _perplexity_lookup(self, query: str) -> Dict[str, Any]:
        """Run academic search via Perplexity sonar-pro-search through OpenRouter."""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        api_key = os.getenv("OPENROUTER_API_KEY")
        model = "perplexity/sonar-pro-search"

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://scientific-writer.local",
            "X-Title": "Scientific Writer Research Tool",
        }

        research_prompt = self._format_academic_prompt(query)

        messages = [
            {
                "role": "system",
                "content": (
                    "You are an academic research assistant specializing in finding "
                    "HIGH-IMPACT, INFLUENTIAL research.\n\n"
                    "QUALITY PRIORITIZATION (CRITICAL):\n"
                    "- ALWAYS prefer highly-cited papers over obscure publications\n"
                    "- ALWAYS prioritize Tier-1 venues: Nature, Science, Cell, NEJM, Lancet, JAMA, PNAS\n"
                    "- ALWAYS prefer papers from established researchers\n"
                    "- Include citation counts when known (e.g., 'cited 500+ times')\n"
                    "- Quality matters more than quantity\n\n"
                    "VENUE HIERARCHY:\n"
                    "1. Nature/Science/Cell family, NEJM, Lancet, JAMA (highest)\n"
                    "2. High-impact specialized journals (IF>10), top conferences (NeurIPS, ICML, ICLR)\n"
                    "3. Respected field-specific journals (IF 5-10)\n"
                    "4. Other peer-reviewed sources (only if no better option)\n\n"
                    "Focus exclusively on scholarly sources. Prioritize recent literature (2020-2026) "
                    "and provide complete citations with DOIs."
                ),
            },
            {"role": "user", "content": research_prompt},
        ]

        data = {
            "model": model,
            "messages": messages,
            "max_tokens": 8000,
            "temperature": 0.1,
            "search_mode": "academic",
            "search_context_size": "high",
        }

        try:
            response = requests.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers=headers,
                json=data,
                timeout=90,
            )
            response.raise_for_status()
            resp_json = response.json()

            if "choices" in resp_json and len(resp_json["choices"]) > 0:
                choice = resp_json["choices"][0]
                if "message" in choice and "content" in choice["message"]:
                    content = choice["message"]["content"]

                    api_citations = self._extract_api_citations(resp_json, choice)
                    text_citations = self._extract_citations_from_text(content)
                    citations = api_citations + text_citations

                    return {
                        "success": True,
                        "query": query,
                        "response": content,
                        "citations": citations,
                        "sources": api_citations,
                        "timestamp": timestamp,
                        "backend": "perplexity",
                        "model": model,
                        "usage": resp_json.get("usage", {}),
                    }
                else:
                    raise Exception("Invalid response format from API")
            else:
                raise Exception("No response choices received from API")

        except Exception as e:
            return {
                "success": False,
                "query": query,
                "error": str(e),
                "timestamp": timestamp,
                "backend": "perplexity",
                "model": model,
            }

    # ------------------------------------------------------------------
    # Shared utilities
    # ------------------------------------------------------------------

    def _format_academic_prompt(self, query: str) -> str:
        """Format a query for academic research results via Perplexity."""
        return f"""You are an expert research assistant. Please provide comprehensive, accurate research information for the following query: "{query}"

IMPORTANT INSTRUCTIONS:
1. Focus on ACADEMIC and SCIENTIFIC sources (peer-reviewed papers, reputable journals, institutional research)
2. Include RECENT information (prioritize 2020-2026 publications)
3. Provide COMPLETE citations with authors, title, journal/conference, year, and DOI when available
4. Structure your response with clear sections and proper attribution
5. Be comprehensive but concise - aim for 800-1200 words
6. Include key findings, methodologies, and implications when relevant
7. Note any controversies, limitations, or conflicting evidence

PAPER QUALITY PRIORITIZATION (CRITICAL):
8. ALWAYS prioritize HIGHLY-CITED papers over obscure publications
9. ALWAYS prioritize papers from TOP-TIER VENUES (Nature, Science, Cell, NEJM, Lancet, JAMA, PNAS)
10. PREFER papers from ESTABLISHED, REPUTABLE AUTHORS
11. For EACH citation include when available: citation count, venue tier, author credentials
12. PRIORITIZE papers that DIRECTLY address the research question

RESPONSE FORMAT:
- Start with a brief summary (2-3 sentences)
- Present key findings and studies in organized sections
- Rank papers by impact: most influential/cited first
- End with future directions or research gaps if applicable
- Include 5-8 high-quality citations

Remember: Quality over quantity. Prioritize influential, highly-cited papers from prestigious venues."""

    def _extract_api_citations(self, response: Dict[str, Any], choice: Dict[str, Any]) -> List[Dict[str, str]]:
        """Extract citations from Perplexity API response fields."""
        citations = []

        search_results = (
            response.get("search_results")
            or choice.get("search_results")
            or choice.get("message", {}).get("search_results")
            or []
        )

        for result in search_results:
            citation = {
                "type": "source",
                "title": result.get("title", ""),
                "url": result.get("url", ""),
                "date": result.get("date", ""),
            }
            if result.get("snippet"):
                citation["snippet"] = result["snippet"]
            citations.append(citation)

        legacy_citations = (
            response.get("citations")
            or choice.get("citations")
            or choice.get("message", {}).get("citations")
            or []
        )

        for url in legacy_citations:
            if isinstance(url, str):
                citations.append({"type": "source", "url": url, "title": "", "date": ""})
            elif isinstance(url, dict):
                citations.append({
                    "type": "source",
                    "url": url.get("url", ""),
                    "title": url.get("title", ""),
                    "date": url.get("date", ""),
                })

        return citations

    def _extract_citations_from_text(self, text: str) -> List[Dict[str, str]]:
        """Extract DOIs and academic URLs from response text as fallback."""
        citations = []

        doi_pattern = r'(?:doi[:\s]*|https?://(?:dx\.)?doi\.org/)(10\.[0-9]{4,}/[^\s\)\]\,\[\<\>]+)'
        doi_matches = re.findall(doi_pattern, text, re.IGNORECASE)
        seen_dois = set()

        for doi in doi_matches:
            doi_clean = doi.strip().rstrip(".,;:)]")
            if doi_clean and doi_clean not in seen_dois:
                seen_dois.add(doi_clean)
                citations.append({
                    "type": "doi",
                    "doi": doi_clean,
                    "url": f"https://doi.org/{doi_clean}",
                })

        url_pattern = (
            r'https?://[^\s\)\]\,\<\>\"\']+(?:arxiv\.org|pubmed|ncbi\.nlm\.nih\.gov|'
            r'nature\.com|science\.org|wiley\.com|springer\.com|ieee\.org|acm\.org)'
            r'[^\s\)\]\,\<\>\"\']*'
        )
        url_matches = re.findall(url_pattern, text, re.IGNORECASE)
        seen_urls = set()

        for url in url_matches:
            url_clean = url.rstrip(".")
            if url_clean not in seen_urls:
                seen_urls.add(url_clean)
                citations.append({"type": "url", "url": url_clean})

        return citations

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def lookup(self, query: str) -> Dict[str, Any]:
        """Perform a research lookup, routing to the best backend.

        Parallel Chat API is used by default. Perplexity sonar-pro-search
        is used only for academic-specific queries (paper searches, DOI lookups).
        """
        backend = self._select_backend(query)
        print(f"[Research] Backend: {backend} | Query: {query[:80]}...", file=sys.stderr)

        if backend == "parallel":
            return self._parallel_lookup(query)
        else:
            return self._perplexity_lookup(query)

    def batch_lookup(self, queries: List[str], delay: float = 1.0) -> List[Dict[str, Any]]:
        """Perform multiple research lookups with delay between requests."""
        results = []
        for i, query in enumerate(queries):
            if i > 0 and delay > 0:
                time.sleep(delay)
            result = self.lookup(query)
            results.append(result)
            print(f"[Research] Completed query {i+1}/{len(queries)}: {query[:50]}...", file=sys.stderr)
        return results


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    """Command-line interface for the research lookup tool."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Research Information Lookup Tool (Parallel Chat API + Perplexity)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # General research (uses Parallel Chat API, core model)
  python research_lookup.py "latest advances in quantum computing 2025"

  # Academic paper search (auto-routes to Perplexity)
  python research_lookup.py "find papers on CRISPR gene editing clinical trials"

  # Force a specific backend
  python research_lookup.py "topic" --force-backend parallel
  python research_lookup.py "topic" --force-backend perplexity

  # Save output to file
  python research_lookup.py "topic" -o results.txt

  # JSON output
  python research_lookup.py "topic" --json -o results.json
        """,
    )
    parser.add_argument("query", nargs="?", help="Research query to look up")
    parser.add_argument("--batch", nargs="+", help="Run multiple queries")
    parser.add_argument(
        "--force-backend",
        choices=["parallel", "perplexity"],
        help="Force a specific backend (default: auto-select)",
    )
    parser.add_argument("-o", "--output", help="Write output to file")
    parser.add_argument("--json", action="store_true", help="Output as JSON")

    args = parser.parse_args()

    output_file = None
    if args.output:
        output_file = open(args.output, "w", encoding="utf-8")

    def write_output(text):
        if output_file:
            output_file.write(text + "\n")
        else:
            print(text)

    has_parallel = bool(os.getenv("PARALLEL_API_KEY"))
    has_perplexity = bool(os.getenv("OPENROUTER_API_KEY"))
    if not has_parallel and not has_perplexity:
        print("Error: No API keys found. Set at least one:", file=sys.stderr)
        print("  export PARALLEL_API_KEY='...'    (primary - Parallel Chat API)", file=sys.stderr)
        print("  export OPENROUTER_API_KEY='...'   (fallback - Perplexity academic)", file=sys.stderr)
        if output_file:
            output_file.close()
        return 1

    if not args.query and not args.batch:
        parser.print_help()
        if output_file:
            output_file.close()
        return 1

    try:
        research = ResearchLookup(force_backend=args.force_backend)

        if args.batch:
            print(f"Running batch research for {len(args.batch)} queries...", file=sys.stderr)
            results = research.batch_lookup(args.batch)
        else:
            print(f"Researching: {args.query}", file=sys.stderr)
            results = [research.lookup(args.query)]

        if args.json:
            write_output(json.dumps(results, indent=2, ensure_ascii=False, default=str))
            if output_file:
                output_file.close()
            return 0

        for i, result in enumerate(results):
            if result["success"]:
                write_output(f"\n{'='*80}")
                write_output(f"Query {i+1}: {result['query']}")
                write_output(f"Timestamp: {result['timestamp']}")
                write_output(f"Backend: {result.get('backend', 'unknown')} | Model: {result.get('model', 'unknown')}")
                write_output(f"{'='*80}")
                write_output(result["response"])

                sources = result.get("sources", [])
                if sources:
                    write_output(f"\nSources ({len(sources)}):")
                    for j, source in enumerate(sources):
                        title = source.get("title", "Untitled")
                        url = source.get("url", "")
                        date = source.get("date", "")
                        date_str = f" ({date})" if date else ""
                        write_output(f"  [{j+1}] {title}{date_str}")
                        if url:
                            write_output(f"      {url}")

                citations = result.get("citations", [])
                text_citations = [c for c in citations if c.get("type") in ("doi", "url")]
                if text_citations:
                    write_output(f"\nAdditional References ({len(text_citations)}):")
                    for j, citation in enumerate(text_citations):
                        if citation.get("type") == "doi":
                            write_output(f"  [{j+1}] DOI: {citation.get('doi', '')} - {citation.get('url', '')}")
                        elif citation.get("type") == "url":
                            write_output(f"  [{j+1}] {citation.get('url', '')}")

                if result.get("usage"):
                    write_output(f"\nUsage: {result['usage']}")
            else:
                write_output(f"\nError in query {i+1}: {result['error']}")

        if output_file:
            output_file.close()
        return 0

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        if output_file:
            output_file.close()
        return 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
"""
AI-powered scientific schematic generation using Nano Banana 2.

This script uses a smart iterative refinement approach:
1. Generate initial image with Nano Banana 2
2. AI quality review using Gemini 3.1 Pro Preview for scientific critique
3. Only regenerate if quality is below threshold for document type
4. Repeat until quality meets standards (max iterations)

Requirements:
    - OPENROUTER_API_KEY environment variable
    - requests library

Usage:
    python generate_schematic_ai.py "Create a flowchart showing CONSORT participant flow" -o flowchart.png
    python generate_schematic_ai.py "Neural network architecture diagram" -o architecture.png --iterations 2
    python generate_schematic_ai.py "Simple block diagram" -o diagram.png --doc-type poster
"""

import argparse
import base64
import json
import os
import sys
import time
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple

try:
    import requests
except ImportError:
    print("Error: requests library not found. Install with: pip install requests")
    sys.exit(1)

# Try to load .env file from multiple potential locations
def _load_env_file():
    """Load .env file from current directory or script directory only."""
    try:
        from dotenv import load_dotenv
    except ImportError:
        return False

    for candidate in [Path.cwd() / ".env", Path(__file__).resolve().parent / ".env"]:
        if candidate.exists():
            load_dotenv(dotenv_path=candidate, override=False)
            return True

    return False


class ScientificSchematicGenerator:
    """Generate scientific schematics using AI with smart iterative refinement.
    
    Uses Gemini 3.1 Pro Preview for quality review to determine if regeneration is needed.
    Multiple passes only occur if the generated schematic doesn't meet the
    quality threshold for the target document type.
    """
    
    # Quality thresholds by document type (score out of 10)
    # Higher thresholds for more formal publications
    QUALITY_THRESHOLDS = {
        "journal": 8.5,      # Nature, Science, etc. - highest standards
        "conference": 8.0,   # Conference papers - high standards
        "poster": 7.0,       # Academic posters - good quality
        "presentation": 6.5, # Slides/talks - clear but less formal
        "report": 7.5,       # Technical reports - professional
        "grant": 8.0,        # Grant proposals - must be compelling
        "thesis": 8.0,       # Dissertations - formal academic
        "preprint": 7.5,     # arXiv, etc. - good quality
        "default": 7.5,      # Default threshold
    }
    
    # Scientific diagram best practices prompt template
    SCIENTIFIC_DIAGRAM_GUIDELINES = """
Create a high-quality scientific diagram with these requirements:

VISUAL QUALITY:
- Clean white or light background (no textures or gradients)
- High contrast for readability and printing
- Professional, publication-ready appearance
- Sharp, clear lines and text
- Adequate spacing between elements to prevent crowding

TYPOGRAPHY:
- Clear, readable sans-serif fonts (Arial, Helvetica style)
- Minimum 10pt font size for all labels
- Consistent font sizes throughout
- All text horizontal or clearly readable
- No overlapping text

SCIENTIFIC STANDARDS:
- Accurate representation of concepts
- Clear labels for all components
- Include scale bars, legends, or axes where appropriate
- Use standard scientific notation and symbols
- Include units where applicable

ACCESSIBILITY:
- Colorblind-friendly color palette (use Okabe-Ito colors if using color)
- High contrast between elements
- Redundant encoding (shapes + colors, not just colors)
- Works well in grayscale

LAYOUT:
- Logical flow (left-to-right or top-to-bottom)
- Clear visual hierarchy
- Balanced composition
- Appropriate use of whitespace
- No clutter or unnecessary decorative elements

IMPORTANT - NO FIGURE NUMBERS:
- Do NOT include "Figure 1:", "Fig. 1", or any figure numbering in the image
- Do NOT add captions or titles like "Figure: ..." at the top or bottom
- Figure numbers and captions are added separately in the document/LaTeX
- The diagram should contain only the visual content itself
"""
    
    def __init__(self, api_key: Optional[str] = None, verbose: bool = False):
        """
        Initialize the generator.
        
        Args:
            api_key: OpenRouter API key (or use OPENROUTER_API_KEY env var)
            verbose: Print detailed progress information
        """
        # Priority: 1) explicit api_key param, 2) environment variable, 3) .env file
        self.api_key = api_key or os.getenv("OPENROUTER_API_KEY")
        
        # If not found in environment, try loading from .env file
        if not self.api_key:
            _load_env_file()
            self.api_key = os.getenv("OPENROUTER_API_KEY")
        
        if not self.api_key:
            raise ValueError(
                "OPENROUTER_API_KEY not found. Please either:\n"
                "  1. Set the OPENROUTER_API_KEY environment variable\n"
                "  2. Add OPENROUTER_API_KEY to your .env file\n"
                "  3. Pass api_key parameter to the constructor\n"
                "Get your API key from: https://openrouter.ai/keys"
            )
        
        self.verbose = verbose
        self._last_error = None  # Track last error for better reporting
        self.base_url = "https://openrouter.ai/api/v1"
        # Nano Banana 2 - Google's advanced image generation model
        # https://openrouter.ai/google/gemini-3-pro-image-preview
        self.image_model = "google/gemini-3.1-flash-image-preview"
        # Gemini 3.1 Pro Preview for quality review - excellent vision and reasoning
        self.review_model = "google/gemini-3.1-pro-preview"
        
    def _log(self, message: str):
        """Log message if verbose mode is enabled."""
        if self.verbose:
            print(f"[{time.strftime('%H:%M:%S')}] {message}")
    
    def _make_request(self, model: str, messages: List[Dict[str, Any]], 
                     modalities: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        Make a request to OpenRouter API.
        
        Args:
            model: Model identifier
            messages: List of message dictionaries
            modalities: Optional list of modalities (e.g., ["image", "text"])
            
        Returns:
            API response as dictionary
        """
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com/scientific-writer",
            "X-Title": "Scientific Schematic Generator"
        }
        
        payload = {
            "model": model,
            "messages": messages
        }
        
        if modalities:
            payload["modalities"] = modalities
        
        self._log(f"Making request to {model}...")
        
        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=120
            )
            
            # Try to get response body even on error
            try:
                response_json = response.json()
            except json.JSONDecodeError:
                response_json = {"raw_text": response.text[:500]}
            
            # Check for HTTP errors but include response body in error message
            if response.status_code != 200:
                error_detail = response_json.get("error", response_json)
                self._log(f"HTTP {response.status_code}: {error_detail}")
                raise RuntimeError(f"API request failed (HTTP {response.status_code}): {error_detail}")
            
            return response_json
        except requests.exceptions.Timeout:
            raise RuntimeError("API request timed out after 120 seconds")
        except requests.exceptions.RequestException as e:
            raise RuntimeError(f"API request failed: {str(e)}")
    
    def _extract_image_from_response(self, response: Dict[str, Any]) -> Optional[bytes]:
        """
        Extract base64-encoded image from API response.
        
        For Nano Banana 2, images are returned in the 'images' field of the message,
        not in the 'content' field.
        
        Args:
            response: API response dictionary
            
        Returns:
            Image bytes or None if not found
        """
        try:
            choices = response.get("choices", [])
            if not choices:
                self._log("No choices in response")
                return None
            
            message = choices[0].get("message", {})
            
            # IMPORTANT: Nano Banana 2 returns images in the 'images' field
            images = message.get("images", [])
            if images and len(images) > 0:
                self._log(f"Found {len(images)} image(s) in 'images' field")
                
                # Get first image
                first_image = images[0]
                if isinstance(first_image, dict):
                    # Extract image_url
                    if first_image.get("type") == "image_url":
                        url = first_image.get("image_url", {})
                        if isinstance(url, dict):
                            url = url.get("url", "")
                        
                        if url and url.startswith("data:image"):
                            # Extract base64 data after comma
                            if "," in url:
                                base64_str = url.split(",", 1)[1]
                                # Clean whitespace
                                base64_str = base64_str.replace('\n', '').replace('\r', '').replace(' ', '')
                                self._log(f"Extracted base64 data (length: {len(base64_str)})")
                                return base64.b64decode(base64_str)
            
            # Fallback: check content field (for other models or future changes)
            content = message.get("content", "")
            
            if self.verbose:
                self._log(f"Content type: {type(content)}, length: {len(str(content))}")
            
            # Handle string content
            if isinstance(content, str) and "data:image" in content:
                import re
                match = re.search(r'data:image/[^;]+;base64,([A-Za-z0-9+/=\n\r]+)', content, re.DOTALL)
                if match:
                    base64_str = match.group(1).replace('\n', '').replace('\r', '').replace(' ', '')
                    self._log(f"Found image in content field (length: {len(base64_str)})")
                    return base64.b64decode(base64_str)
            
            # Handle list content
            if isinstance(content, list):
                for i, block in enumerate(content):
                    if isinstance(block, dict) and block.get("type") == "image_url":
                        url = block.get("image_url", {})
                        if isinstance(url, dict):
                            url = url.get("url", "")
                        if url and url.startswith("data:image") and "," in url:
                            base64_str = url.split(",", 1)[1].replace('\n', '').replace('\r', '').replace(' ', '')
                            self._log(f"Found image in content block {i}")
                            return base64.b64decode(base64_str)
            
            self._log("No image data found in response")
            return None
            
        except Exception as e:
            self._log(f"Error extracting image: {str(e)}")
            import traceback
            if self.verbose:
                traceback.print_exc()
            return None
    
    def _image_to_base64(self, image_path: str) -> str:
        """
        Convert image file to base64 data URL.
        
        Args:
            image_path: Path to image file
            
        Returns:
            Base64 data URL string
        """
        with open(image_path, "rb") as f:
            image_data = f.read()
        
        # Determine image type from extension
        ext = Path(image_path).suffix.lower()
        mime_type = {
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".gif": "image/gif",
            ".webp": "image/webp"
        }.get(ext, "image/png")
        
        base64_data = base64.b64encode(image_data).decode("utf-8")
        return f"data:{mime_type};base64,{base64_data}"
    
    def generate_image(self, prompt: str) -> Optional[bytes]:
        """
        Generate an image using Nano Banana 2.
        
        Args:
            prompt: Description of the diagram to generate
            
        Returns:
            Image bytes or None if generation failed
        """
        self._last_error = None  # Reset error
        
        messages = [
            {
                "role": "user",
                "content": prompt
            }
        ]
        
        try:
            response = self._make_request(
                model=self.image_model,
                messages=messages,
                modalities=["image", "text"]
            )
            
            # Debug: print response structure if verbose
            if self.verbose:
                self._log(f"Response keys: {response.keys()}")
                if "error" in response:
                    self._log(f"API Error: {response['error']}")
                if "choices" in response and response["choices"]:
                    msg = response["choices"][0].get("message", {})
                    self._log(f"Message keys: {msg.keys()}")
                    # Show content preview without printing huge base64 data
                    content = msg.get("content", "")
                    if isinstance(content, str):
                        preview = content[:200] + "..." if len(content) > 200 else content
                        self._log(f"Content preview: {preview}")
                    elif isinstance(content, list):
                        self._log(f"Content is list with {len(content)} items")
                        for i, item in enumerate(content[:3]):
                            if isinstance(item, dict):
                                self._log(f"  Item {i}: type={item.get('type')}")
            
            # Check for API errors in response
            if "error" in response:
                error_msg = response["error"]
                if isinstance(error_msg, dict):
                    error_msg = error_msg.get("message", str(error_msg))
                self._last_error = f"API Error: {error_msg}"
                print(f"✗ {self._last_error}")
                return None
            
            image_data = self._extract_image_from_response(response)
            if image_data:
                self._log(f"✓ Generated image ({len(image_data)} bytes)")
            else:
                self._last_error = "No image data in API response - model may not support image generation"
                self._log(f"✗ {self._last_error}")
                # Additional debug info when image extraction fails
                if self.verbose and "choices" in response:
                    msg = response["choices"][0].get("message", {})
                    self._log(f"Full message structure: {json.dumps({k: type(v).__name__ for k, v in msg.items()})}")
            
            return image_data
        except RuntimeError as e:
            self._last_error = str(e)
            self._log(f"✗ Generation failed: {self._last_error}")
            return None
        except Exception as e:
            self._last_error = f"Unexpected error: {str(e)}"
            self._log(f"✗ Generation failed: {self._last_error}")
            import traceback
            if self.verbose:
                traceback.print_exc()
            return None
    
    def review_image(self, image_path: str, original_prompt: str, 
                    iteration: int, doc_type: str = "default",
                    max_iterations: int = 2) -> Tuple[str, float, bool]:
        """
        Review generated image using Gemini 3.1 Pro Preview for quality analysis.
        
        Uses Gemini 3.1 Pro Preview's superior vision and reasoning capabilities to
        evaluate the schematic quality and determine if regeneration is needed.
        
        Args:
            image_path: Path to the generated image
            original_prompt: Original user prompt
            iteration: Current iteration number
            doc_type: Document type (journal, poster, presentation, etc.)
            max_iterations: Maximum iterations allowed
            
        Returns:
            Tuple of (critique text, quality score 0-10, needs_improvement bool)
        """
        # Use Gemini 3.1 Pro Preview for review - excellent vision and analysis
        image_data_url = self._image_to_base64(image_path)
        
        # Get quality threshold for this document type
        threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(), 
                                                 self.QUALITY_THRESHOLDS["default"])
        
        review_prompt = f"""You are an expert reviewer evaluating a scientific diagram for publication quality.

ORIGINAL REQUEST: {original_prompt}

DOCUMENT TYPE: {doc_type} (quality threshold: {threshold}/10)
ITERATION: {iteration}/{max_iterations}

Carefully evaluate this diagram on these criteria:

1. **Scientific Accuracy** (0-2 points)
   - Correct representation of concepts
   - Proper notation and symbols
   - Accurate relationships shown

2. **Clarity and Readability** (0-2 points)
   - Easy to understand at a glance
   - Clear visual hierarchy
   - No ambiguous elements

3. **Label Quality** (0-2 points)
   - All important elements labeled
   - Labels are readable (appropriate font size)
   - Consistent labeling style

4. **Layout and Composition** (0-2 points)
   - Logical flow (top-to-bottom or left-to-right)
   - Balanced use of space
   - No overlapping elements

5. **Professional Appearance** (0-2 points)
   - Publication-ready quality
   - Clean, crisp lines and shapes
   - Appropriate colors/contrast

RESPOND IN THIS EXACT FORMAT:
SCORE: [total score 0-10]

STRENGTHS:
- [strength 1]
- [strength 2]

ISSUES:
- [issue 1 if any]
- [issue 2 if any]

VERDICT: [ACCEPTABLE or NEEDS_IMPROVEMENT]

If score >= {threshold}, the diagram is ACCEPTABLE for {doc_type} publication.
If score < {threshold}, mark as NEEDS_IMPROVEMENT with specific suggestions."""

        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": review_prompt
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_data_url
                        }
                    }
                ]
            }
        ]
        
        try:
            # Use Gemini 3.1 Pro Preview for high-quality review
            response = self._make_request(
                model=self.review_model,
                messages=messages
            )
            
            # Extract text response
            choices = response.get("choices", [])
            if not choices:
                return "Image generated successfully", 8.0
            
            message = choices[0].get("message", {})
            content = message.get("content", "")
            
            # Check reasoning field (Nano Banana 2 puts analysis here)
            reasoning = message.get("reasoning", "")
            if reasoning and not content:
                content = reasoning
            
            if isinstance(content, list):
                # Extract text from content blocks
                text_parts = []
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "text":
                        text_parts.append(block.get("text", ""))
                content = "\n".join(text_parts)
            
            # Try to extract score
            score = 7.5  # Default score if extraction fails
            import re
            
            # Look for SCORE: X or SCORE: X/10 format
            score_match = re.search(r'SCORE:\s*(\d+(?:\.\d+)?)', content, re.IGNORECASE)
            if score_match:
                score = float(score_match.group(1))
            else:
                # Fallback: look for any score pattern
                score_match = re.search(r'(?:score|rating|quality)[:\s]+(\d+(?:\.\d+)?)\s*(?:/\s*10)?', content, re.IGNORECASE)
                if score_match:
                    score = float(score_match.group(1))
            
            # Determine if improvement is needed based on verdict or score
            needs_improvement = False
            if "NEEDS_IMPROVEMENT" in content.upper():
                needs_improvement = True
            elif score < threshold:
                needs_improvement = True
            
            self._log(f"✓ Review complete (Score: {score}/10, Threshold: {threshold}/10)")
            self._log(f"  Verdict: {'Needs improvement' if needs_improvement else 'Acceptable'}")
            
            return (content if content else "Image generated successfully", 
                    score, 
                    needs_improvement)
        except Exception as e:
            self._log(f"Review skipped: {str(e)}")
            # Don't fail the whole process if review fails - assume acceptable
            return "Image generated successfully (review skipped)", 7.5, False
    
    def improve_prompt(self, original_prompt: str, critique: str, 
                      iteration: int) -> str:
        """
        Improve the generation prompt based on critique.
        
        Args:
            original_prompt: Original user prompt
            critique: Review critique from previous iteration
            iteration: Current iteration number
            
        Returns:
            Improved prompt for next generation
        """
        improved_prompt = f"""{self.SCIENTIFIC_DIAGRAM_GUIDELINES}

USER REQUEST: {original_prompt}

ITERATION {iteration}: Based on previous feedback, address these specific improvements:
{critique}

Generate an improved version that addresses all the critique points while maintaining scientific accuracy and professional quality."""
        
        return improved_prompt
    
    def generate_iterative(self, user_prompt: str, output_path: str,
                          iterations: int = 2, 
                          doc_type: str = "default") -> Dict[str, Any]:
        """
        Generate scientific schematic with smart iterative refinement.
        
        Only regenerates if the quality score is below the threshold for the
        specified document type. This saves API calls and time when the first
        generation is already good enough.
        
        Args:
            user_prompt: User's description of desired diagram
            output_path: Path to save final image
            iterations: Maximum refinement iterations (default: 2, max: 2)
            doc_type: Document type for quality threshold (journal, poster, etc.)
            
        Returns:
            Dictionary with generation results and metadata
        """
        output_path = Path(output_path)
        output_dir = output_path.parent
        output_dir.mkdir(parents=True, exist_ok=True)
        
        base_name = output_path.stem
        extension = output_path.suffix or ".png"
        
        # Get quality threshold for this document type
        threshold = self.QUALITY_THRESHOLDS.get(doc_type.lower(), 
                                                 self.QUALITY_THRESHOLDS["default"])
        
        results = {
            "user_prompt": user_prompt,
            "doc_type": doc_type,
            "quality_threshold": threshold,
            "iterations": [],
            "final_image": None,
            "final_score": 0.0,
            "success": False,
            "early_stop": False,
            "early_stop_reason": None
        }
        
        current_prompt = f"""{self.SCIENTIFIC_DIAGRAM_GUIDELINES}

USER REQUEST: {user_prompt}

Generate a publication-quality scientific diagram that meets all the guidelines above."""
        
        print(f"\n{'='*60}")
        print(f"Generating Scientific Schematic")
        print(f"{'='*60}")
        print(f"Description: {user_prompt}")
        print(f"Document Type: {doc_type}")
        print(f"Quality Threshold: {threshold}/10")
        print(f"Max Iterations: {iterations}")
        print(f"Output: {output_path}")
        print(f"{'='*60}\n")
        
        for i in range(1, iterations + 1):
            print(f"\n[Iteration {i}/{iterations}]")
            print("-" * 40)
            
            # Generate image
            print(f"Generating image...")
            image_data = self.generate_image(current_prompt)
            
            if not image_data:
                error_msg = getattr(self, '_last_error', 'Image generation failed - no image data returned')
                print(f"✗ Generation failed: {error_msg}")
                results["iterations"].append({
                    "iteration": i,
                    "success": False,
                    "error": error_msg
                })
                continue
            
            # Save iteration image
            iter_path = output_dir / f"{base_name}_v{i}{extension}"
            with open(iter_path, "wb") as f:
                f.write(image_data)
            print(f"✓ Saved: {iter_path}")
            
            # Review image using Gemini 3.1 Pro Preview
            print(f"Reviewing image with Gemini 3.1 Pro Preview...")
            critique, score, needs_improvement = self.review_image(
                str(iter_path), user_prompt, i, doc_type, iterations
            )
            print(f"✓ Score: {score}/10 (threshold: {threshold}/10)")
            
            # Save iteration results
            iteration_result = {
                "iteration": i,
                "image_path": str(iter_path),
                "prompt": current_prompt,
                "critique": critique,
                "score": score,
                "needs_improvement": needs_improvement,
                "success": True
            }
            results["iterations"].append(iteration_result)
            
            # Check if quality is acceptable - STOP EARLY if so
            if not needs_improvement:
                print(f"\n✓ Quality meets {doc_type} threshold ({score} >= {threshold})")
                print(f"  No further iterations needed!")
                results["final_image"] = str(iter_path)
                results["final_score"] = score
                results["success"] = True
                results["early_stop"] = True
                results["early_stop_reason"] = f"Quality score {score} meets threshold {threshold} for {doc_type}"
                break
            
            # If this is the last iteration, we're done regardless
            if i == iterations:
                print(f"\n⚠ Maximum iterations reached")
                results["final_image"] = str(iter_path)
                results["final_score"] = score
                results["success"] = True
                break
            
            # Quality below threshold - improve prompt for next iteration
            print(f"\n⚠ Quality below threshold ({score} < {threshold})")
            print(f"Improving prompt based on feedback...")
            current_prompt = self.improve_prompt(user_prompt, critique, i + 1)
        
        # Copy final version to output path
        if results["success"] and results["final_image"]:
            final_iter_path = Path(results["final_image"])
            if final_iter_path != output_path:
                import shutil
                shutil.copy(final_iter_path, output_path)
                print(f"\n✓ Final image: {output_path}")
        
        # Save review log
        log_path = output_dir / f"{base_name}_review_log.json"
        with open(log_path, "w") as f:
            json.dump(results, f, indent=2)
        print(f"✓ Review log: {log_path}")
        
        print(f"\n{'='*60}")
        print(f"Generation Complete!")
        print(f"Final Score: {results['final_score']}/10")
        if results["early_stop"]:
            print(f"Iterations Used: {len([r for r in results['iterations'] if r.get('success')])}/{iterations} (early stop)")
        print(f"{'='*60}\n")
        
        return results


def main():
    """Command-line interface."""
    parser = argparse.ArgumentParser(
        description="Generate scientific schematics using AI with smart iterative refinement",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Generate a flowchart for a journal paper
  python generate_schematic_ai.py "CONSORT participant flow diagram" -o flowchart.png --doc-type journal
  
  # Generate neural network architecture for presentation (lower threshold)
  python generate_schematic_ai.py "Transformer encoder-decoder architecture" -o transformer.png --doc-type presentation
  
  # Generate with custom max iterations for poster
  python generate_schematic_ai.py "Biological signaling pathway" -o pathway.png --iterations 2 --doc-type poster
  
  # Verbose output
  python generate_schematic_ai.py "Circuit diagram" -o circuit.png -v

Document Types (quality thresholds):
  journal      8.5/10  - Nature, Science, peer-reviewed journals
  conference   8.0/10  - Conference papers
  thesis       8.0/10  - Dissertations, theses
  grant        8.0/10  - Grant proposals
  preprint     7.5/10  - arXiv, bioRxiv, etc.
  report       7.5/10  - Technical reports
  poster       7.0/10  - Academic posters
  presentation 6.5/10  - Slides, talks
  default      7.5/10  - General purpose

Note: Multiple iterations only occur if quality is BELOW the threshold.
      If the first generation meets the threshold, no extra API calls are made.

Environment:
  OPENROUTER_API_KEY    OpenRouter API key (required)
        """
    )
    
    parser.add_argument("prompt", help="Description of the diagram to generate")
    parser.add_argument("-o", "--output", required=True, 
                       help="Output image path (e.g., diagram.png)")
    parser.add_argument("--iterations", type=int, default=2,
                       help="Maximum refinement iterations (default: 2, max: 2)")
    parser.add_argument("--doc-type", default="default",
                       choices=["journal", "conference", "poster", "presentation", 
                               "report", "grant", "thesis", "preprint", "default"],
                       help="Document type for quality threshold (default: default)")
    parser.add_argument("--api-key", help="OpenRouter API key (or set OPENROUTER_API_KEY)")
    parser.add_argument("-v", "--verbose", action="store_true",
                       help="Verbose output")
    
    args = parser.parse_args()
    
    # Check for API key
    api_key = args.api_key or os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        print("Error: OPENROUTER_API_KEY environment variable not set")
        print("\nSet it with:")
        print("  export OPENROUTER_API_KEY='your_api_key'")
        print("\nOr provide via --api-key flag")
        sys.exit(1)
    
    # Validate iterations - enforce max of 2
    if args.iterations < 1 or args.iterations > 2:
        print("Error: Iterations must be between 1 and 2")
        sys.exit(1)
    
    try:
        generator = ScientificSchematicGenerator(api_key=api_key, verbose=args.verbose)
        results = generator.generate_iterative(
            user_prompt=args.prompt,
            output_path=args.output,
            iterations=args.iterations,
            doc_type=args.doc_type
        )
        
        if results["success"]:
            print(f"\n✓ Success! Image saved to: {args.output}")
            if results.get("early_stop"):
                print(f"  (Completed in {len([r for r in results['iterations'] if r.get('success')])} iteration(s) - quality threshold met)")
            sys.exit(0)
        else:
            print(f"\n✗ Generation failed. Check review log for details.")
            sys.exit(1)
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()