"""Full research pipeline: idea -> literature -> methods -> results.

This demonstrates the complete project workflow on a simple harmonic
oscillator study. Each step builds on the previous one:

  1. Idea generation: propose a research direction
  2. Literature search: find relevant papers via Semantic Scholar
  3. Methods development: design the analysis approach
  4. Results: execute code locally to produce data and plots

The results step uses a WebSocket session — code generated by the AI
runs locally on your machine in an isolated virtual environment.

Output files are pulled to output_pipeline_oscillator/ locally.

Examples of available models (pass via default_model):
  - gemini-3.1-flash-lite-preview  (fast, cheapest)
  - gemini-3.1-pro-preview         (stronger reasoning)
  - gpt-5-nano                     (fast, good quality)
  - gpt-5.2                        (strongest OpenAI model)
  - claude-sonnet-4-6              (strong all-round)
"""

import os
import air

client = air.AIR(
    api_key=os.environ["AIR_API_KEY"],
    base_url=os.environ.get("AIR_BASE_URL", "http://localhost:8000"),
    # local_dir: where project files are pulled to locally after each step.
    local_dir=os.path.join(os.path.dirname(os.path.abspath(__file__)), "output_pipeline_oscillator"),
)

project_name = "harmonic-oscillator"

data_description = """
We want to study the simple and damped harmonic oscillator from both
analytical and numerical perspectives.

Context and resources:
- Python with numpy, scipy, and matplotlib
- Limited hardware: 2 CPUs, 8 GB RAM — keep simulations lightweight
"""

# # Clean start
# try:
#     client.delete_project(project_name)
# except Exception:
#     pass

print(f"Data description:\n{data_description.strip()}\n")
print(f"Creating project: {project_name}\n")

project = client.create_project(
    name=project_name,
    data_description=data_description,
)

# --- Step 1: Idea ---
print("=" * 60)
print("[1/4] Generating research idea ...")
print("=" * 60)
idea = project.idea()
print(f"\n{idea}\n")

# --- Step 2: Literature ---
print("=" * 60)
print("[2/4] Running literature search ...")
print("=" * 60)
lit = project.literature(timeout=900)
print(f"\n{lit}\n")

# --- Step 3: Methods ---
print("=" * 60)
print("[3/4] Developing methods ...")
print("=" * 60)
meth = project.methods()
print(f"\n{meth}\n")

# --- Step 4: Results (local code execution) ---
print("=" * 60)
print("[4/4] Generating results (code runs locally) ...")
print("=" * 60)
# Examples of available models:
#   gemini-3.1-flash-lite-preview, gemini-3.1-pro-preview,
#   gpt-5-nano, gpt-5.2, claude-sonnet-4-6
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output_pipeline_oscillator")
res = project.results(
    max_steps=2,                                    # number of plan steps
    max_attempts=5,                                 # retries per step on failure
    n_plan_reviews=1,                               # plan review rounds
    engineer_model="gemini-3.1-flash-lite-preview", # LLM for engineer
    researcher_model="gemini-3.1-flash-lite-preview",# LLM for researcher
    # planner_model="gemini-3.1-flash-lite-preview",  # LLM for planner
    # plan_reviewer_model="gemini-3.1-pro-preview",   # LLM for plan reviewer
    # default_model="gemini-3.1-flash-lite-preview",  # default for all agents
    # formatter_model="gemini-3.1-flash-lite-preview", # LLM for formatters
    work_dir=output_dir,                            # local execution directory
)
print(f"\nResults: {len(res)} chars")

# Pull all files
project.pull_files()

print(f"\nOutput files pulled to: {client._local_dir}")

client.close()
