#!/usr/bin/env python3
"""200 parallel AI paper research agents on 50 Daytona sandboxes (4 agents each).

Each agent searches a unique AI research niche, finds recent papers,
and identifies under-explored ideas with high probability of success.

Usage:
    python scripts/paper_hunt.py
"""

import asyncio
import json
import logging
import os
import time
import warnings
from pathlib import Path

warnings.filterwarnings('ignore')
logging.getLogger('asyncio').setLevel(logging.CRITICAL)
logging.getLogger('daytona').setLevel(logging.CRITICAL)

env_path = Path(__file__).resolve().parent.parent / '.env'
if env_path.exists():
    for line in env_path.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith('#') and '=' in line:
            key, _, value = line.partition('=')
            os.environ.setdefault(key.strip(), value.strip())

from daytona_sdk import (
    AsyncDaytona,
    CreateSandboxFromImageParams,
    DaytonaConfig,
    Image,
)

API_KEY = os.environ['DAYTONA_API_KEY']
ANTHROPIC_KEY = os.environ['ANTHROPIC_API_KEY']

# 200 unique AI research niches — no duplicates
NICHES = [
    # RL & Decision Making (20)
    "GRPO algorithm improvements for LLM training 2025 2026",
    "offline reinforcement learning real world deployment",
    "multi-agent reinforcement learning emergent communication",
    "reward hacking detection and prevention in RLHF",
    "constitutional AI and RLAIF new methods",
    "inverse reinforcement learning from human demonstrations",
    "hierarchical reinforcement learning for long-horizon tasks",
    "model-based RL with world models for planning",
    "RL for chip design and hardware optimization",
    "safe reinforcement learning with formal guarantees",
    "reinforcement learning from AI feedback RLAIF scaling",
    "process reward models vs outcome reward models",
    "RLHF alternatives direct preference optimization DPO",
    "multi-objective reinforcement learning Pareto optimization",
    "RL for autonomous driving decision making 2026",
    "curiosity-driven exploration in sparse reward environments",
    "meta-reinforcement learning few-shot adaptation",
    "reinforcement learning for combinatorial optimization",
    "language model self-play debate for alignment",
    "reward model ensembles for robust RLHF",
    # LLM Architecture (20)
    "mixture of experts scaling beyond trillion parameters",
    "state space models Mamba replacing transformers",
    "linear attention mechanisms subquadratic transformers",
    "retrieval augmented generation RAG improvements 2026",
    "long context transformers million token windows",
    "sparse transformers efficient inference at scale",
    "multi-modal foundation models unified architecture",
    "small language models on-device deployment",
    "speculative decoding faster LLM inference",
    "KV cache optimization for long context",
    "quantization-aware training for LLM compression",
    "neural architecture search for transformer variants",
    "diffusion transformers for language generation",
    "byte-level tokenization replacing BPE",
    "mixture of depths adaptive compute allocation",
    "continuous batching and PagedAttention inference",
    "early exit strategies for efficient LLM inference",
    "modular neural networks compositional generalization",
    "test-time compute scaling and inference-time training",
    "distillation from large to small language models 2026",
    # AI Agents (20)
    "autonomous coding agents benchmark evaluation",
    "tool use and function calling in LLMs",
    "multi-agent systems task decomposition orchestration",
    "agent memory architectures long-term episodic",
    "web browsing agents autonomous navigation",
    "computer use agents GUI interaction",
    "sandboxed agent execution environments",
    "agent-to-agent communication protocols",
    "self-improving AI agents recursive enhancement",
    "agent planning with tree search and MCTS",
    "LLM agents for scientific research automation",
    "code generation agents real world evaluation",
    "agent safety guardrails and monitoring",
    "multi-modal agents vision language action",
    "agent frameworks comparison LangChain CrewAI AutoGen",
    "persistent agent memory and knowledge graphs",
    "agent skill libraries reusable capabilities",
    "human-agent collaboration hybrid workflows",
    "agent evaluation benchmarks SWE-bench WebArena",
    "autonomous agents for data analysis and science",
    # AI Safety & Alignment (20)
    "mechanistic interpretability latest breakthroughs 2026",
    "scalable oversight and recursive reward modeling",
    "AI alignment via debate and amplification",
    "red teaming LLMs automated adversarial attacks",
    "watermarking AI generated text detection",
    "jailbreaking LLMs defense mechanisms 2026",
    "AI governance frameworks and regulation",
    "deceptive alignment and sleeper agents detection",
    "corrigibility and shutdown problem in AI",
    "value learning from human preferences at scale",
    "AI safety benchmarks and evaluation frameworks",
    "representation engineering for model behavior control",
    "circuit-level understanding of transformer features",
    "backdoor attacks on fine-tuned language models",
    "differential privacy in LLM training",
    "unlearning specific knowledge from trained models",
    "constitutional AI methods for harmlessness",
    "AI risk assessment frameworks for frontier models",
    "cooperative AI multi-agent safety",
    "formal verification of neural network properties",
    # Computer Vision & Multimodal (20)
    "video generation models Sora competitors 2026",
    "3D scene generation from text descriptions",
    "image segmentation foundation models SAM2",
    "visual language models for robotics",
    "neural radiance fields NeRF real-time rendering",
    "gaussian splatting 3D reconstruction advances",
    "medical imaging AI diagnosis new methods",
    "autonomous driving perception models 2026",
    "document understanding and OCR with LLMs",
    "visual reasoning and spatial understanding",
    "image generation consistency and controllability",
    "video understanding and temporal reasoning",
    "point cloud processing for 3D understanding",
    "visual grounding and referring expression",
    "few-shot visual classification with foundation models",
    "synthetic data generation for vision training",
    "visual instruction tuning methods",
    "depth estimation monocular and stereo",
    "action recognition in videos transformer models",
    "visual anomaly detection industrial applications",
    # AI for Science (20)
    "protein structure prediction beyond AlphaFold",
    "AI for drug discovery molecular generation 2026",
    "weather forecasting AI models GraphCast GenCast",
    "AI for materials science discovery",
    "scientific literature mining and synthesis",
    "AI for theorem proving Lean4 formalization",
    "computational chemistry with machine learning",
    "genomics and gene editing with AI",
    "AI for climate modeling and prediction",
    "AI accelerated physics simulations",
    "retrosynthesis planning with AI",
    "AI for battery and energy storage research",
    "astronomical object detection and classification",
    "AI for epidemiology and disease modeling",
    "quantum chemistry with neural networks",
    "AI for experimental design and optimization",
    "foundation models for scientific data",
    "AI for protein engineering and design",
    "symbolic regression and equation discovery",
    "AI for systems biology and metabolomics",
    # Efficient AI & Edge (20)
    "on-device LLM inference optimization mobile",
    "model pruning structured unstructured methods",
    "knowledge distillation for edge deployment",
    "neural network quantization INT4 INT2",
    "efficient fine-tuning LoRA QLoRA adapters",
    "hardware-aware neural architecture search",
    "TinyML microcontroller machine learning",
    "federated learning privacy-preserving training",
    "split computing edge-cloud inference",
    "energy-efficient AI training green computing",
    "model compression for real-time inference",
    "binary neural networks extreme quantization",
    "dynamic inference early exit networks",
    "efficient attention mechanisms for edge",
    "neural network compilation and optimization",
    "AI chip architecture innovation 2026",
    "sparsity-aware training and inference",
    "continual learning without catastrophic forgetting",
    "parameter-efficient transfer learning methods",
    "efficient data loading and preprocessing pipelines",
    # NLP & Language (20)
    "low-resource language NLP underserved languages",
    "multilingual LLMs cross-lingual transfer",
    "structured output generation JSON schemas",
    "factual consistency and hallucination reduction",
    "summarization faithfulness metrics",
    "dialogue systems and conversational AI 2026",
    "information extraction from unstructured text",
    "text-to-SQL natural language databases",
    "sentiment analysis and opinion mining advances",
    "named entity recognition few-shot methods",
    "question answering over long documents",
    "machine translation low-resource language pairs",
    "text style transfer and controllable generation",
    "coreference resolution in long documents",
    "argument mining and reasoning in text",
    "biomedical NLP clinical text processing",
    "legal NLP contract analysis automation",
    "financial NLP earnings call analysis",
    "code-switching and multilingual text processing",
    "pragmatics and implicit meaning understanding",
    # Robotics & Embodied AI (20)
    "foundation models for robot manipulation",
    "sim-to-real transfer in robotics 2026",
    "humanoid robots learning locomotion",
    "robot learning from human video demonstrations",
    "tactile sensing and dexterous manipulation",
    "autonomous mobile robots navigation SLAM",
    "drone swarm coordination and planning",
    "surgical robotics AI-assisted procedures",
    "robot task planning with LLMs",
    "deformable object manipulation",
    "multi-robot collaboration and communication",
    "robot safety and human-robot interaction",
    "legged locomotion rough terrain",
    "grasping unknown objects zero-shot",
    "imitation learning from few demonstrations",
    "robot foundation models RT-2 and beyond",
    "embodied AI in simulation environments",
    "soft robotics and compliant manipulation",
    "agricultural robotics harvesting automation",
    "underwater robotics autonomous exploration",
    # Emerging & Wild (20)
    "AI-generated music composition and production",
    "AI for game design procedural content generation",
    "neuro-symbolic AI combining neural and logical",
    "causal inference and causal discovery with AI",
    "AI for education personalized tutoring systems",
    "digital twins AI-powered simulation",
    "AI for cybersecurity threat detection 2026",
    "AI for supply chain optimization",
    "generative AI for architectural design",
    "AI companionship and emotional intelligence",
    "swarm intelligence simulation prediction MiroFish",
    "AI for accessibility assistive technology",
    "AI-powered code review and bug detection",
    "synthetic biology and AI protein design",
    "AI for fraud detection and financial crime",
    "neural interface brain-computer interaction AI",
    "AI for urban planning and smart cities",
    "AI for agriculture precision farming",
    "AI for legal reasoning and case prediction",
    "consciousness and sentience in AI systems debate",
]

# Script that runs 4 searches inside one sandbox
AGENT_SCRIPT = '''
import json
import os
import sys
import time
import urllib.request
import urllib.error
import concurrent.futures

ANTHROPIC_KEY = os.environ["ANTHROPIC_KEY"]
NICHES = json.loads(os.environ["NICHES_JSON"])
BATCH_ID = os.environ["BATCH_ID"]

def search_niche(niche_id, niche):
    try:
        payload = json.dumps({
            "model": "claude-sonnet-4-6",
            "max_tokens": 600,
            "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 2}],
            "messages": [{"role": "user", "content":
                f"Search for recent AI research papers (2025-2026) on: {niche}\\n\\n"
                "Find the most promising UNDER-EXPLORED niche idea in this area. "
                "Return ONLY:\\n"
                "1) NICHE: one-line description of the under-explored opportunity\\n"
                "2) WHY: why this is likely to succeed (1 sentence)\\n"
                "3) PAPER: most relevant recent paper title + authors\\n"
                "4) GAP: what specific gap exists that nobody is filling\\n"
                "5) FEASIBILITY: low/medium/high for a small team\\n"
                "Be specific and contrarian. Find what others are missing."
            }]
        }).encode("utf-8")

        headers = {
            "Content-Type": "application/json",
            "x-api-key": ANTHROPIC_KEY,
            "anthropic-version": "2023-06-01",
        }

        req = urllib.request.Request(
            "https://api.anthropic.com/v1/messages",
            data=payload, headers=headers, method="POST"
        )

        with urllib.request.urlopen(req, timeout=60) as resp:
            body = json.loads(resp.read().decode("utf-8"))

        text = " ".join(b["text"] for b in body.get("content", []) if b.get("type") == "text").strip()
        tokens = body.get("usage", {}).get("input_tokens", 0) + body.get("usage", {}).get("output_tokens", 0)

        result = {"id": niche_id, "niche": niche, "result": text, "tokens": tokens}
        print(f"RESULT:{json.dumps(result)}", flush=True)
        return result

    except Exception as e:
        result = {"id": niche_id, "niche": niche, "error": str(e)[:200]}
        print(f"RESULT:{json.dumps(result)}", flush=True)
        return result

# Run 4 niches in parallel threads inside this sandbox
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
    futures = []
    for i, niche in enumerate(NICHES):
        niche_id = f"batch{BATCH_ID}_agent{i}"
        futures.append(pool.submit(search_niche, niche_id, niche))
    for f in concurrent.futures.as_completed(futures):
        f.result()

print("BATCH_DONE", flush=True)
'''

RESET = '\033[0m'
BOLD = '\033[1m'
DIM = '\033[2m'
GREEN = '\033[92m'
CYAN = '\033[96m'
RED = '\033[91m'
YELLOW = '\033[93m'


async def run_batch(daytona, batch_id, niches, sem, results, counters, total, log_file):
    async with sem:
        sandbox = None
        try:
            sandbox = await daytona.create(
                CreateSandboxFromImageParams(image=Image.debian_slim('3.12'), ephemeral=True)
            )

            niches_json = json.dumps(niches).replace("'", "'\\''")
            await sandbox.fs.upload_file(AGENT_SCRIPT.encode('utf-8'), '/home/user/search.py')

            env_cmd = f'ANTHROPIC_KEY={ANTHROPIC_KEY} BATCH_ID={batch_id} NICHES_JSON=\'{niches_json}\''
            response = await sandbox.process.exec(f'{env_cmd} python -u /home/user/search.py', timeout=120)
            stdout = response.result if hasattr(response, 'result') else ''

            for line in stdout.strip().splitlines():
                if line.startswith('RESULT:'):
                    data = json.loads(line[7:])
                    results.append(data)
                    counters['done'] += 1
                    n = counters['done']

                    if 'error' in data:
                        msg = f"  [{RED}ERR{RESET}] {n:3d}/{total}  {data['niche'][:50]}  {data['error'][:40]}"
                    else:
                        short = data['result'][:80].replace('\n', ' ')
                        msg = f"  [{GREEN}OK{RESET} ] {n:3d}/{total}  {short}"

                    print(msg, flush=True)
                    with open(log_file, 'a') as f:
                        f.write(msg + '\n')

        except Exception as exc:
            counters['done'] += len(niches)
            msg = f"  [{RED}FAIL{RESET}] batch {batch_id}: {type(exc).__name__}: {str(exc)[:60]}"
            print(msg, flush=True)
            with open(log_file, 'a') as f:
                f.write(msg + '\n')

        finally:
            if sandbox:
                try:
                    await daytona.delete(sandbox)
                except Exception:
                    pass


async def main():
    niches = NICHES[:200]
    total = len(niches)

    # 4 niches per sandbox = 50 sandboxes
    batch_size = 4
    batches = [niches[i:i+batch_size] for i in range(0, total, batch_size)]

    log_file = '/tmp/paper_hunt.log'
    with open(log_file, 'w') as f:
        f.write(f"Paper Hunt — {total} agents across {len(batches)} sandboxes\n")
        f.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")

    print(f'\n{BOLD}{"=" * 70}')
    print(f'  AI PAPER HUNT — {total} research agents on {len(batches)} Daytona sandboxes')
    print(f'  4 agents per sandbox, each searching a unique AI research niche')
    print(f'  Log: tail -f {log_file}')
    print(f'{"=" * 70}{RESET}\n')

    config = DaytonaConfig(api_key=API_KEY)
    daytona = AsyncDaytona(config)
    sem = asyncio.Semaphore(25)
    results = []
    counters = {'done': 0}

    start = time.perf_counter()

    tasks = [
        run_batch(daytona, i, batch, sem, results, counters, total, log_file)
        for i, batch in enumerate(batches)
    ]
    await asyncio.gather(*tasks)

    elapsed = time.perf_counter() - start

    # Write final report
    succeeded = [r for r in results if 'error' not in r]
    failed = [r for r in results if 'error' in r]

    report = f'\n{"=" * 70}\n'
    report += f'  PAPER HUNT RESULTS — {len(succeeded)}/{total} niches analyzed\n'
    report += f'  Time: {elapsed:.1f}s | Failed: {len(failed)}\n'
    report += f'{"=" * 70}\n\n'

    # Sort by feasibility mentions
    for r in sorted(succeeded, key=lambda x: x.get('niche', '')):
        report += f'--- {r["niche"][:60]} ---\n'
        report += f'{r["result"][:300]}\n\n'

    report += f'{"=" * 70}\n'

    print(report)
    with open(log_file, 'a') as f:
        f.write(report)

    # Also save as JSON
    with open('/tmp/paper_hunt_results.json', 'w') as f:
        json.dump({'results': results, 'elapsed': elapsed, 'total': total}, f, indent=2)

    print(f'  Full results: /tmp/paper_hunt_results.json')
    print(f'  Log: {log_file}\n')


if __name__ == '__main__':
    asyncio.run(main())
