#!/usr/bin/env python3
"""
Run the Research & Generate agent: research → generate any protocol → run → learn from feedback.

Uses env.run_assay_with_protocol() so the agent can try arbitrary parameter values
(not limited to presets). Feedback from each run is passed into the next trial.
Works with any spec that has evaluate_custom_protocol (PCR, ELISA, etc.).

Requires: pip install openai, OPENAI_API_KEY set.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from lab_env.env import LabEnv
from lab_env.spec import get_spec_for_workflow
from agents.research_generate_agent import ResearchGenerateAgent


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run Research & Generate agent (research → generate protocol → run → learn)"
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=5,
        help="Number of episodes to run",
    )
    parser.add_argument(
        "--workflow",
        type=str,
        default="pcr-amplification",
        choices=["pcr-amplification", "elisa-readout"],
        help="Experiment type (uses spec with custom protocol support)",
    )
    parser.add_argument(
        "--max-trials",
        type=int,
        default=6,
        help="Max protocol attempts per episode",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print each trial's protocol and result",
    )
    args = parser.parse_args()

    spec = get_spec_for_workflow(args.workflow)
    env = LabEnv(spec=spec)
    agent = ResearchGenerateAgent(max_trials=args.max_trials)

    print(f"Research & Generate agent — workflow={args.workflow}, episodes={args.episodes}")
    print("(Research → generate protocol → run in lab → learn from feedback)\n")

    results = []
    for ep in range(args.episodes):
        seed = args.seed + ep * 1000
        if args.verbose:
            print(f"--- Episode {ep + 1} (seed={seed}) ---")
        out = agent.run_episode(env, seed=seed, verbose=args.verbose)
        results.append(out)
        if not args.verbose:
            print(
                f"  Episode {ep + 1}: reward={out['reward']:.1f}, "
                f"success={out['success']}, protocols_tried={out['num_protocols_tried']}"
            )

    env.close()

    n = len(results)
    print("\n--- Summary ---")
    print(f"  Success rate:     {sum(r['success'] for r in results) / n:.1%}")
    print(f"  Partial rate:     {sum(r['partial'] for r in results) / n:.1%}")
    print(f"  Avg reward:       {sum(r['reward'] for r in results) / n:.1f}")
    print(f"  Avg protocols:    {sum(r['num_protocols_tried'] for r in results) / n:.1f} per episode")


if __name__ == "__main__":
    main()