"""
Toy REINFORCE training for CompilerOptimizationEnv (CPU, no PyTorch).

This is meant for the Hugging Face Space: fast smoke training that uses the
same `runtime_core` mock engine and passes as the Gradio demo.

For full GRPO + Unsloth + LLM, use your Colab / GPU notebooks
(`compiler_optimization_grpo.ipynb`, `role2_deliverable3_*`).

CLI:  python train.py --episodes 50 --max-steps 8 --seed 0
Or import:  from train import run_toy_training; print(run_toy_training(20, 8, 0))
"""

from __future__ import annotations

import argparse
import json
import math
import random
import textwrap
from typing import List, Sequence, Tuple

from runtime_core import (
    CompilerOptimizationEnv,
    MOCK_PASSES,
    MockEngine,
    SAMPLE_PROGRAM,
)

# Additional tiny programs so the policy is not overfit to a single IR.
TRAINING_PROGRAMS: List[List[dict]] = [
    SAMPLE_PROGRAM,
    [
        {"op": "const", "dest": "a", "args": ["1"], "type": "int"},
        {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"},
        {"op": "ret", "args": ["b"]},
    ],
    [
        {"op": "const", "dest": "x", "args": ["2"], "type": "int"},
        {"op": "const", "dest": "y", "args": ["4"], "type": "int"},
        {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"},
        {"op": "const", "dest": "k", "args": ["1"], "type": "int"},
        {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"},
        {"op": "ret", "args": ["w"]},
    ],
]

ACTIONS: Tuple[str, ...] = tuple(sorted(MOCK_PASSES.keys()))


def _softmax(logits: Sequence[float]) -> List[float]:
    m = max(logits) if logits else 0.0
    ex = [math.exp(x - m) for x in logits]
    s = sum(ex) or 1.0
    return [e / s for e in ex]


def _sample_action(rng: random.Random, logits: List[float]) -> Tuple[int, List[float], float]:
    """Return (action index, prob vector, log prob of chosen action)."""
    p = _softmax(logits)
    u = rng.random()
    acc = 0.0
    idx = len(p) - 1
    for i, pi in enumerate(p):
        acc += pi
        if u <= acc:
            idx = i
            break
    log_p = math.log(p[idx] + 1e-12)
    return idx, p, log_p


def _rollout(
    engine: MockEngine,
    program: List[dict],
    max_steps: int,
    logits: List[float],
    rng: random.Random,
) -> Tuple[float, List[Tuple[int, List[float], float]]]:
    """
    One episode. Returns total reward and per-step (action index, prob vector, step reward)
    for REINFORCE with returns G_t = sum of rewards from t onward.
    """
    env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=max_steps)
    env.reset(program)
    total = 0.0
    trace: List[Tuple[int, List[float], float]] = []
    for _ in range(max_steps):
        a_idx, pvec, _ = _sample_action(rng, logits)
        action = ACTIONS[a_idx]
        step = env.step(action)
        r = float(step.reward)
        total += r
        trace.append((a_idx, pvec, r))
        if step.done:
            break
    return total, trace


def _reinforce_update(
    logits: List[float], trace: List[Tuple[int, List[float], float]], lr: float
) -> None:
    """REINFORCE with Monte Carlo return G_t from each step."""
    G = 0.0
    for t in range(len(trace) - 1, -1, -1):
        a_idx, pvec, r = trace[t]
        G += r
        for i in range(len(logits)):
            delta = 1.0 if i == a_idx else 0.0
            logits[i] += lr * G * (delta - pvec[i])


def run_toy_training(
    episodes: int,
    max_steps: int,
    seed: int = 0,
    lr: float = 0.15,
) -> str:
    """
    Train a stateless categorical policy over the three mock passes; print-friendly report.
    """
    if episodes < 1:
        return "episodes must be >= 1"
    if max_steps < 1:
        return "max_steps must be >= 1"
    episodes = int(episodes)
    max_steps = int(max_steps)
    seed = int(seed)
    lr = float(lr)

    rng = random.Random(seed)
    engine = MockEngine()
    logits = [0.0 for _ in ACTIONS]

    history: List[Tuple[int, float]] = []
    for ep in range(1, episodes + 1):
        program = TRAINING_PROGRAMS[(ep - 1) % len(TRAINING_PROGRAMS)]
        G, trace = _rollout(engine, program, max_steps, logits, rng)
        if trace:
            _reinforce_update(logits, trace, lr)
        history.append((ep, G))

    final_p = _softmax(logits)
    lines = [
        "Toy REINFORCE (stateless policy over pass names, CPU, stdlib only)",
        f"  episodes={episodes}  max_steps={max_steps}  seed={seed}  lr={lr}",
        f"  actions order: {list(ACTIONS)}",
        "",
        f"  final logits: {[round(x, 4) for x in logits]}",
        f"  final policy:  {', '.join(f'{a}={p:.3f}' for a, p in zip(ACTIONS, final_p))}",
        "",
        "  return per episode (last 10): " + ", ".join(f"{G:+.1f}" for _, G in history[-10:]),
        "",
    ]
    if history:
        mean_r = sum(G for _, G in history) / len(history)
        lines.append(f"  mean return over all episodes: {mean_r:+.3f}")
    lines.append("")
    lines.append(
        textwrap.dedent(
            """
            This does not train an LLM. For GRPO + Qwen + Unsloth, run the project notebooks
            on a GPU machine (e.g. Colab), not the CPU Space.
            """
        ).strip()
    )
    return "\n".join(lines)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Toy REINFORCE on CompilerOptimizationEnv (CPU)")
    p.add_argument("--episodes", type=int, default=50, help="Number of training episodes")
    p.add_argument("--max-steps", type=int, default=8, help="max_steps per env episode")
    p.add_argument("--seed", type=int, default=0, help="RNG seed")
    p.add_argument("--lr", type=float, default=0.15, help="REINFORCE learning rate")
    p.add_argument(
        "--out-json",
        type=str,
        default="",
        help="If set, write a small run summary to this path (e.g. training_log.json).",
    )
    return p.parse_args()


def main() -> None:
    args = _parse_args()
    report = run_toy_training(
        episodes=args.episodes,
        max_steps=args.max_steps,
        seed=args.seed,
        lr=args.lr,
    )
    print(report)
    if args.out_json:
        payload = {
            "episodes": args.episodes,
            "max_steps": args.max_steps,
            "seed": args.seed,
            "lr": args.lr,
            "summary_text": report,
        }
        with open(args.out_json, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2)
        print(f"\nWrote {args.out_json}")


if __name__ == "__main__":
    main()