"""Template strings for ``orze --init`` scaffolding.

Calling spec:
    from orze.cli_demo import BASELINE_TRAIN_PY, RESEARCH_RULES_TEMPLATE

Both are plain strings written to disk by the --init handler in cli.py.
No runtime logic lives here — only the template content.
"""

# ---------------------------------------------------------------------------
# Baseline train.py template
# ---------------------------------------------------------------------------

BASELINE_TRAIN_PY = '''\
#!/usr/bin/env python3
"""Baseline training script generated by orze --init.

Demonstrates the full orze contract with a task-agnostic demo
(linear regression on synthetic data). Only depends on pyyaml.

Replace the training logic with your own — keep the argument
parsing and metrics output as-is.

Orze calls this script as:
    python train.py --idea-id <id> --results-dir <dir> --ideas-md <file> --config <yaml>

CONTRACT:
  1. Read your idea config from results-dir/idea-id/idea_config.yaml
     (orze writes this file before launching)
  2. Do your training
  3. Write metrics to results-dir/idea-id/metrics.json with at minimum:
     {"status": "COMPLETED", "<primary_metric>": <value>, "training_time": <seconds>}
  4. Exit 0 on success, non-zero on failure
"""
import argparse
import json
import math
import random
import re
import time
from pathlib import Path

import yaml


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def deep_merge(base: dict, override: dict) -> dict:
    """Recursively merge *override* into *base* (returns new dict)."""
    merged = dict(base)
    for k, v in override.items():
        if k in merged and isinstance(merged[k], dict) and isinstance(v, dict):
            merged[k] = deep_merge(merged[k], v)
        else:
            merged[k] = v
    return merged


def load_idea_config(ideas_md: str, idea_id: str) -> dict:
    """Parse ideas.md and return the YAML config block for *idea_id*."""
    text = Path(ideas_md).read_text(encoding="utf-8")
    pattern = rf"^##\\s+{re.escape(idea_id)}\\b.*?```ya?ml\\s*\\n(.*?)```"
    m = re.search(pattern, text, re.MULTILINE | re.DOTALL)
    if not m:
        raise ValueError(f"Idea {idea_id} not found in {ideas_md}")
    return yaml.safe_load(m.group(1)) or {}


# ---------------------------------------------------------------------------
# Demo training loop — linear regression on synthetic data
# ---------------------------------------------------------------------------

def _generate_data(n=200, noise=0.1, seed=42):
    """y = 3*x + 2 + noise."""
    rng = random.Random(seed)
    xs = [rng.gauss(0, 1) for _ in range(n)]
    ys = [3.0 * x + 2.0 + rng.gauss(0, noise) for x in xs]
    return xs, ys


def _train_linear(xs, ys, lr=0.01, epochs=50):
    """Gradient descent on y = w*x + b.  Returns (w, b, history)."""
    w, b = 0.0, 0.0
    n = len(xs)
    history = []
    for epoch in range(1, epochs + 1):
        # forward
        preds = [w * x + b for x in xs]
        loss = sum((p - y) ** 2 for p, y in zip(preds, ys)) / n
        # gradients
        dw = sum(2 * (p - y) * x for p, y, x in zip(preds, ys, xs)) / n
        db = sum(2 * (p - y) for p, y in zip(preds, ys)) / n
        w -= lr * dw
        b -= lr * db
        history.append(loss)
    return w, b, history


def _r2_score(ys, preds):
    mean_y = sum(ys) / len(ys)
    ss_res = sum((y - p) ** 2 for y, p in zip(ys, preds))
    ss_tot = sum((y - mean_y) ** 2 for y in ys)
    return 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--idea-id", required=True)
    parser.add_argument("--results-dir", required=True)
    parser.add_argument("--ideas-md", required=True)
    parser.add_argument("--config", required=True, help="Path to base config YAML")
    args = parser.parse_args()

    idea_dir = Path(args.results_dir) / args.idea_id
    idea_dir.mkdir(parents=True, exist_ok=True)

    # 1. Load base config
    base_cfg = {}
    if Path(args.config).exists():
        base_cfg = yaml.safe_load(Path(args.config).read_text()) or {}

    # 2. Load idea-specific config (written by orze before launch)
    idea_cfg_path = idea_dir / "idea_config.yaml"
    idea_cfg = {}
    if idea_cfg_path.exists():
        idea_cfg = yaml.safe_load(idea_cfg_path.read_text()) or {}

    # 3. Merge: idea overrides base
    cfg = deep_merge(base_cfg, idea_cfg)
    print(f"Training {args.idea_id} with config: {cfg}")

    # 4. Extract hyperparameters (with defaults)
    lr = float(cfg.get("learning_rate", 0.01))
    epochs = int(cfg.get("epochs", 50))
    noise = float(cfg.get("noise", 0.1))
    seed = int(cfg.get("seed", 42))

    # 5. Train
    t0 = time.time()
    xs_train, ys_train = _generate_data(n=160, noise=noise, seed=seed)
    xs_val, ys_val = _generate_data(n=40, noise=noise, seed=seed + 1)

    w, b, history = _train_linear(xs_train, ys_train, lr=lr, epochs=epochs)
    training_time = time.time() - t0

    # 6. Evaluate
    val_preds = [w * x + b for x in xs_val]
    val_loss = sum((p - y) ** 2 for p, y in zip(val_preds, ys_val)) / len(ys_val)
    train_loss = history[-1] if history else float("inf")
    r2 = _r2_score(ys_val, val_preds)

    # 7. Write metrics (orze reads this to determine success/failure)
    metrics = {
        "status": "COMPLETED",
        "val_loss": round(val_loss, 6),
        "train_loss": round(train_loss, 6),
        "r2": round(r2, 6),
        "epochs": epochs,
        "training_time": round(training_time, 2),
    }
    (idea_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
    print(f"Done. Metrics: {metrics}")


if __name__ == "__main__":
    main()
'''

# ---------------------------------------------------------------------------
# RESEARCH_RULES.md template
# ---------------------------------------------------------------------------

RESEARCH_RULES_TEMPLATE = '''\
# Research Rules
<!-- Edit this file to steer the research agent. -->

## Goal
Improve performance on the task defined in this project.
Read `{results_dir}/report.md` for current results.

## Current state
- Cycle: {cycle}
- Completed experiments: {completed}
- Queued experiments: {queued}

## Strategy
1. Start with simple baselines (different learning rates, epochs).
2. Try architectural changes one at a time.
3. Combine the best-performing ideas.
4. Avoid changing more than two hyperparameters at once.

## Idea format
Append new ideas to `{ideas_file}` using this format:

```markdown
## idea-XXXX: Short descriptive title
- **Priority**: high|medium|low

\\```yaml
learning_rate: 0.001
epochs: 20
# ... any hyperparameters your train script accepts
\\```
```

## Rules
- **Append-only** — never edit or delete existing ideas.
- **Unique IDs** — increment from the highest existing idea number.
- **Complete configs** — every idea must have all required hyperparameters.
- **One change at a time** — isolate variables for clear attribution.
'''
