"""Pilot sweep: slope x cable stiffness x gait strategy, greedy gait search.

Writes results/pilot.json, which the web console (/app/) replays.
Run: .venv/bin/python experiments/run_pilot.py
"""
from __future__ import annotations

import itertools
import json
import multiprocessing
import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
from tensegrity.gait import (  # noqa: E402
    Robot, greedy_gait, pair_actions, single_actions)
from tensegrity.model import bar_pairs, cable_pairs, node_positions  # noqa: E402

SLOPES = [0, 5, 10, 15, 20]
STIFFNESSES = [100, 300, 800]
STRATEGIES = {
    # name: (candidate actions, release seconds)
    "single_release": ("single", 1.0),
    "single_hold": ("single", 0.0),
    "pair_hold": ("pair", 0.0),
}
STEPS = 6
RATIO = 0.4
FRAME_EVERY = 50  # 50 x 2 ms = one frame every 0.1 s


def run_one(config: tuple[int, int, str]) -> dict:
    slope, stiffness, strategy = config
    kind, release = STRATEGIES[strategy]
    candidates = single_actions() if kind == "single" else single_actions() + pair_actions()
    robot = Robot(slope_deg=slope, stiffness=stiffness, frame_every=FRAME_EVERY)
    start_x = float(robot.com()[0])
    frames: list = []
    robot.frames = frames
    log = greedy_gait(robot, steps=STEPS, ratio=RATIO, candidates=candidates, release=release)
    return {
        "slope": slope, "stiffness": stiffness, "strategy": strategy,
        "start_x": start_x, "net_dx": log[-1]["x"] - start_x,
        "steps": log, "frames_per_step": len(frames) // STEPS, "frames": frames,
    }


def main() -> None:
    configs = list(itertools.product(SLOPES, STIFFNESSES, STRATEGIES))
    with multiprocessing.Pool() as pool:
        runs = pool.map(run_one, configs)
    for run in runs:
        print(f"slope {run['slope']:>2}°  k={run['stiffness']:>3}  {run['strategy']:<15} net {run['net_dx']:+.2f} m")
    out = pathlib.Path(__file__).resolve().parents[1] / "results" / "pilot.json"
    out.parent.mkdir(exist_ok=True)
    out.write_text(json.dumps({
        "bars": bar_pairs(), "cables": cable_pairs(node_positions(1.0)),
        "steps": STEPS, "ratio": RATIO, "frame_dt": FRAME_EVERY * 0.002, "runs": runs,
    }, separators=(",", ":")))
    print("wrote", out, out.stat().st_size // 1024, "KB")


if __name__ == "__main__":
    main()