"""标准库：小车环境、规划器和图像/注意力/跨动作一致性算例。"""
from argparse import ArgumentParser
from hashlib import sha256
from itertools import product
from pathlib import Path
import json

HERE = Path(__file__).resolve().parent
PLAN = HERE / "data/world-models-plan.json"
RECORD = HERE / "data/world-models-fit.json"
SMALL = HERE / "data/world-models-small.json"
MODES = ("position_action", "history_no_action", "history_action")
NAMES = ("单位置＋动作", "两位置，无动作", "两位置＋动作")
DT = 0.2


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def save_new(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("x") as f:
        json.dump(data, f, ensure_ascii=False, indent=2, allow_nan=False)
        f.write("\n")


def step(x, velocity, action, retention=0.8):
    """x 米、velocity 米/秒；命令 action 无单位，执行器在 ±1 饱和。"""
    applied = max(-1.0, min(1.0, action))
    next_velocity = retention * velocity + 0.4 * applied
    return x + DT * next_velocity, next_velocity


def features(mode, previous, current, action):
    if mode == "position_action":
        return [current, action, 1.0]
    if mode == "history_no_action":
        return [current, previous, 1.0]
    if mode == "history_action":
        return [current, previous, action, 1.0]
    raise ValueError(mode)


def predict(model, previous, current, action):
    if model["mode"] == "oracle":
        return step(current, (current - previous) / DT, action)[0]
    return sum(w * x for w, x in zip(
        model["weights"], features(model["mode"], previous, current, action),
    ))


def rollout(model, previous, current, actions):
    result = []
    for action in actions:
        following = predict(model, previous, current, action)
        result.append(following)
        previous, current = current, following
    return result


def choose(model, previous, current, goal, actions=(0, -1, 1), horizon=4):
    """穷举全部候选；目标仅用预测，不访问未来环境真值。"""
    best = None
    for index, sequence in enumerate(product(actions, repeat=horizon)):
        imagined = rollout(model, previous, current, sequence)
        velocity = (imagined[-1] - imagined[-2]) / DT
        cost = ((imagined[-1] - goal) ** 2 + 0.1 * velocity ** 2
                + 0.01 * sum(a * a for a in sequence))
        key = (round(cost, 12), index)  # 仅稳定近似相等的候选次序。
        if best is None or key < best[0]:
            best = (key, sequence, imagined, cost)
    return {"actions": list(best[1]), "imagined": best[2], "cost": best[3],
            "candidate_count": len(actions) ** horizon}


def control(model, case, actions=(0, -1, 1)):
    x, velocity, goal = case["x"], case["velocity"], case["goal"]
    previous = x - DT * velocity
    trace = []
    for t in range(24):
        plan = choose(model, previous, x, goal, actions)
        action = plan["actions"][0]
        new_x, new_velocity = step(x, velocity, action)
        trace.append({"t": t, "previous": previous, "x": x, "velocity": velocity,
                      "action": action, "next_x": new_x, "next_velocity": new_velocity,
                      "plan": plan})
        previous, x, velocity = x, new_x, new_velocity
    return {"case": case, "trace": trace, "final_error": abs(x - goal),
            "final_speed": abs(velocity),
            "success": abs(x - goal) < 0.15 and abs(velocity) < 0.2}


def image_toy():
    """两个等概率未来；每个未来只有一个 2×2 亮块。"""
    frames = []
    for column in (10, 22):
        frame = [[0.0] * 32 for _ in range(32)]
        for y in (15, 16):
            for x in (column, column + 1):
                frame[y][x] = 1.0
        frames.append(frame)
    mean = [[(frames[0][y][x] + frames[1][y][x]) / 2
             for x in range(32)] for y in range(32)]
    blank = [[0.0] * 32 for _ in range(32)]
    def expected_mse(prediction):
        return sum(sum((prediction[y][x] - actual[y][x]) ** 2
                       for y in range(32) for x in range(32)) / 1024
                   for actual in frames) / 2
    return {"futures": frames, "mean": mean,
            "expected_mse": {"mean": expected_mse(mean),
                             "one_valid_future": expected_mse(frames[0]),
                             "blank": expected_mse(blank)}}


def mixed_mask():
    labels = ["文0", "BOI", "图0", "图1", "EOI",
              "文1", "BOI", "图2", "图3", "EOI"]
    blocks = [None, None, 0, 0, None, None, None, 1, 1, None]
    mask = [[int(j <= i or (blocks[i] is not None and blocks[i] == blocks[j]))
             for j in range(len(labels))] for i in range(len(labels))]
    return {"labels": labels, "allowed": mask}


def shared_mechanism():
    """手工构造的两条冲突预测；不是已训练视频模型的输出。"""
    initial_velocity = 1.0
    actions, predicted_velocities = [1.0, -1.0], [1.2, 0.1]
    separate = [(v - 0.4 * a) / initial_velocity
                for a, v in zip(actions, predicted_velocities)]
    shared = sum(separate) / len(separate)
    error = sum((shared * initial_velocity + 0.4 * a - v) ** 2
                for a, v in zip(actions, predicted_velocities)) / 2
    return {"initial_velocity": initial_velocity, "actions": actions,
            "predictions": predicted_velocities, "separate_retentions": separate,
            "shared_retention": shared, "separate_mse": 0.0, "shared_mse": error}


def small_experiments():
    return {"image": image_toy(), "mask": mixed_mask(), "shared": shared_mechanism()}


def summary(record):
    rows = []
    for regime in ("standard", "strong_commands", "changed_retention"):
        rows.append(f"{regime}：单步 / 第 16 步位置 RMSE（米）")
        for mode, name in zip(MODES, NAMES):
            values = [r["evaluation"][regime][mode] for r in record["runs"]]
            one = sum(v["one_step_rmse"] for v in values) / len(values)
            last = sum(v["rollout_rmse"][-1] for v in values) / len(values)
            rows.append(f"  {name}：{one:.4f} / {last:.4f}")
    rows.append("闭环控制成功次数（3 份拟合 × 18 个固定案例）：")
    for mode in record["control_methods"]:
        cases = [c for run in record["runs"] for c in run["control"][mode]]
        rows.append(f"  {mode}：{sum(c['success'] for c in cases)}/{len(cases)}")
    toy = small_experiments()
    rows.append(f"图像均值的期望 MSE：{toy['image']['expected_mse']['mean']:.6f}")
    rows.append(f"共享惯性系数的拟合 MSE：{toy['shared']['shared_mse']:.4f}")
    return "\n".join(rows)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--small", type=Path)
    args = parser.parse_args()
    if args.small:
        save_new(args.small, small_experiments())
        print(args.small)
        return
    record = json.loads(RECORD.read_text())
    assert record["complete"]
    print(summary(record))


if __name__ == "__main__":
    main()