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.literature_review import LiteratureReview
from scientist_ai.proof_engine import ProofEngine
from scientist_ai.scientific_writer import ScientificWriter

# --- CONFIG ---
OPENAI_API_KEY = "YOUR_OPENAI_API_KEY"  # Replace with your key

# --- 1. Literature Review ---
lr = LiteratureReview()
papers = lr.fetch_arxiv("class number L-functions", max_results=2)
lr.add_paper(
    title="On the Growth of Class Numbers",
    authors=["A. Mathematician"],
    summary="This paper studies the growth of class numbers in number fields.",
    link="https://arxiv.org/abs/1234.5678"
)

# --- 2. Knowledge Brain ---
kb = KnowledgeBrain()
for paper in lr.papers:
    if 'title' in paper:
        kb.add_concept(paper['title'], {'type': 'paper', 'summary': paper['summary']})
kb.add_concept('L-functions', {'type': 'concept', 'description': 'Complex functions encoding number-theoretic information'})
kb.add_relation('On the Growth of Class Numbers', 'L-functions', 'related_to')

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

# --- 4. Proof Engine ---
pe = ProofEngine()
proof_results = []
for hyp in hyps:
    # For demo, try to prove the statement as a symbolic string (if possible)
    proof_results.append(pe.prove(hyp['statement']))

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

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

# --- 7. Scientific Writer ---
sw = ScientificWriter(api_key=OPENAI_API_KEY)
article = sw.write_article(
    title="Automated Discovery in Class Number Theory",
    abstract="We use an AI system to generate and test new hypotheses about class numbers and L-functions.",
    hypotheses=hyps,
    results=proof_results,
    bibliography=lr.get_bibliography(),
    style="arxiv"
)

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