import os
from scientist_ai.literature_review import LiteratureReview
from scientist_ai.knowledge_brain import KnowledgeBrain
from scientist_ai.hypothesis_brain import HypothesisBrain
from scientist_ai.critic_brain import CriticBrain
from scientist_ai.evolution_brain import EvolutionBrain
from scientist_ai.proof_engine import ProofEngine
from scientist_ai.scientific_writer import ScientificWriter

# --- CONFIG ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "YOUR_OPENAI_API_KEY")  # Set as env var or replace
SUBJECT = "class number L-functions"
AUTHOR_STYLE = "Terence Tao"

# --- 1. Literature Review ---
lr = LiteratureReview()
papers = lr.fetch_arxiv(SUBJECT, max_results=3)
# Optionally, add more papers manually or from other APIs

# --- 2. Summarize Papers ---
sw = ScientificWriter(api_key=OPENAI_API_KEY)
summaries = []
for paper in lr.papers:
    if 'raw' in paper:
        summaries.append(sw.summarize(paper['raw'], style=AUTHOR_STYLE))
    elif 'summary' in paper:
        summaries.append(sw.summarize(paper['summary'], style=AUTHOR_STYLE))

# --- 3. Knowledge Brain ---
kb = KnowledgeBrain()
for i, paper in enumerate(lr.papers):
    kb.add_concept(f"Paper_{i+1}", {'type': 'paper', 'summary': summaries[i] if i < len(summaries) else ''})

# --- 4. Hypothesis Generation ---
hb = HypothesisBrain()
hyps = hb.generate(kb)

# --- 5. Proof Engine ---
pe = ProofEngine()
proof_results = []
for hyp in hyps:
    proof_results.append(pe.prove(hyp['statement']))

# --- 6. Critic Brain ---
cb = CriticBrain()
critiques = [cb.attack(h, kb) for h in hyps]

# --- 7. Evolution Brain ---
eb = EvolutionBrain()
survivors = eb.evolve(hyps)

# --- 8. Write Article in Author Style ---
article = sw.write_article(
    title=f"Automated Research on {SUBJECT}",
    abstract=f"We use an AI system to generate and test new hypotheses about {SUBJECT}.",
    hypotheses=hyps,
    results=proof_results,
    bibliography=lr.get_bibliography(),
    style=AUTHOR_STYLE
)

print("\n--- ARTICLE DRAFT ---\n")
print(article)
