import json
import os
import time
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OUTPUT = (
    REPO_ROOT
    / "evaluation"
    / "outputs"
    / "train_replay_6tasks_30clips"
)
OUTPUT = Path(
    os.environ.get("FULL_EVAL_OUTPUT", str(DEFAULT_OUTPUT))
).expanduser().resolve()
OUTPUT.mkdir(parents=True, exist_ok=True)

os.environ.setdefault(
    "SMOKE_OUTPUT",
    str(OUTPUT / "_smoke_import"),
)

import offline_replay_smoke as smoke

NUM_EPISODES_PER_TASK = 5
ANCHOR_RATIOS = (0.20, 0.35, 0.50, 0.65, 0.80)
EXECUTE_STEPS = 5
ACTION_HORIZON = smoke.ACTION_HORIZON
HISTORY_OFFSETS = smoke.HISTORY_OFFSETS
RANDOM_SEED = 20260916

JOINT_NAMES = (
    "left_joint_1",
    "left_joint_2",
    "left_joint_3",
    "left_joint_4",
    "left_joint_5",
    "left_joint_6",
    "left_gripper",
    "right_joint_1",
    "right_joint_2",
    "right_joint_3",
    "right_joint_4",
    "right_joint_5",
    "right_joint_6",
    "right_gripper",
)

ARM_DIMS = smoke.ARM_DIMS
GRIPPER_DIMS = smoke.GRIPPER_DIMS


def load_all_episode_index():
    metadata_files = sorted(
        (smoke.DATASET / "meta" / "episodes").rglob("*.parquet")
    )
    if not metadata_files:
        raise FileNotFoundError(
            smoke.DATASET / "meta" / "episodes"
        )

    tables = [
        pq.read_table(
            path,
            columns=["episode_index", "tasks", "length"],
        )
        for path in metadata_files
    ]
    table = (
        pa.concat_tables(tables)
        if len(tables) > 1
        else tables[0]
    )

    groups = {}
    for row_index in range(table.num_rows):
        episode_index = int(
            table["episode_index"][row_index].as_py()
        )
        tasks = table["tasks"][row_index].as_py()
        length = int(table["length"][row_index].as_py())

        if not tasks:
            raise RuntimeError(
                f"Episode {episode_index} has no task"
            )

        task = str(tasks[0])
        groups.setdefault(task, []).append(
            {
                "episode_index": episode_index,
                "length": length,
            }
        )

    ordered_groups = sorted(
        groups.items(),
        key=lambda item: min(
            row["episode_index"] for row in item[1]
        ),
    )

    selections = []

    for task_id, (task, episodes) in enumerate(ordered_groups):
        episodes = sorted(
            episodes,
            key=lambda row: row["episode_index"],
        )

        if len(episodes) < NUM_EPISODES_PER_TASK:
            raise RuntimeError(
                f"Task {task!r} only has {len(episodes)} episodes"
            )

        positions = np.rint(
            np.linspace(
                0,
                len(episodes) - 1,
                NUM_EPISODES_PER_TASK,
            )
        ).astype(int)

        for sample_id, (position, anchor_ratio) in enumerate(
            zip(positions, ANCHOR_RATIOS, strict=True)
        ):
            selected = episodes[int(position)]
            selections.append(
                {
                    "task_id": task_id,
                    "task": task,
                    "sample_id": sample_id,
                    "episode_index": selected["episode_index"],
                    "length": selected["length"],
                    "anchor_ratio": float(anchor_ratio),
                }
            )

    if len(ordered_groups) != 6:
        raise RuntimeError(
            f"Expected 6 tasks, got {len(ordered_groups)}"
        )

    if len(selections) != 30:
        raise RuntimeError(
            f"Expected 30 selections, got {len(selections)}"
        )

    return selections


def create_policy():
    config = smoke.train_config.get_config(
        smoke.CONFIG_NAME
    )

    data_config = config.data
    assets_config = getattr(data_config, "assets", None)

    if assets_config is not None:
        assets_config = smoke.replace_dataclass_fields(
            assets_config,
            assets_dir=str(smoke.CKPT / "assets"),
            asset_id="real_piper_x_6tasks",
        )

    data_config = smoke.replace_dataclass_fields(
        data_config,
        repo_id=str(smoke.DATASET),
        assets=assets_config,
    )
    config = smoke.replace_dataclass_fields(
        config,
        data=data_config,
    )

    print("===== CONFIG =====", flush=True)
    print(f"config={config.name}", flush=True)
    print(f"model={config.model}", flush=True)
    print(f"checkpoint={smoke.CKPT}", flush=True)
    print(f"dataset={config.data.repo_id}", flush=True)
    print(
        f"assets_dir={config.data.assets.assets_dir}",
        flush=True,
    )

    start = time.monotonic()
    policy = smoke.policy_config.create_trained_policy(
        config,
        smoke.CKPT,
    )
    load_seconds = time.monotonic() - start

    print(
        f"policy_load_seconds={load_seconds:.3f}",
        flush=True,
    )

    return config, policy, load_seconds


def array_errors(prediction, target, hold, steps):
    prediction = prediction[:steps]
    target = target[:steps]
    hold = hold[:steps]

    model_difference = prediction - target
    hold_difference = hold - target

    result = {}

    for name, dims in (
        ("all", None),
        ("arm", ARM_DIMS),
        ("gripper", GRIPPER_DIMS),
    ):
        if dims is None:
            model_selected = model_difference
            hold_selected = hold_difference
        else:
            model_selected = model_difference[:, dims]
            hold_selected = hold_difference[:, dims]

        result[f"model_mae_{name}"] = float(
            np.mean(np.abs(model_selected))
        )
        result[f"model_rmse_{name}"] = float(
            np.sqrt(np.mean(np.square(model_selected)))
        )
        result[f"hold_mae_{name}"] = float(
            np.mean(np.abs(hold_selected))
        )
        result[f"hold_rmse_{name}"] = float(
            np.sqrt(np.mean(np.square(hold_selected)))
        )

    result["motion_mae_all"] = float(
        np.mean(np.abs(target - hold))
    )
    result["model_beats_hold_all"] = bool(
        result["model_mae_all"] < result["hold_mae_all"]
    )

    return result


def add_prefixed_metrics(row, prefix, metrics):
    for key, value in metrics.items():
        row[f"{prefix}_{key}"] = value


def decode_clip_images(metadata, frame_indices):
    images_by_frame = {
        frame_index: {}
        for frame_index in frame_indices
    }
    decode_debug = {}

    for camera in smoke.CAMERAS:
        video_path, from_timestamp, to_timestamp = (
            smoke.video_path_and_range(metadata, camera)
        )

        target_timestamps = [
            from_timestamp + frame_index / smoke.FPS
            for frame_index in frame_indices
        ]

        images, decoded_timestamps, errors = (
            smoke.decode_nearest_frames(
                video_path,
                target_timestamps,
            )
        )

        for position, frame_index in enumerate(frame_indices):
            images_by_frame[frame_index][camera] = (
                smoke.chw_uint8(images[position])
            )

        decode_debug[camera] = {
            "video_path": str(video_path),
            "from_timestamp": from_timestamp,
            "to_timestamp": to_timestamp,
            "targets": target_timestamps,
            "decoded": decoded_timestamps,
            "errors": errors,
        }

    return images_by_frame, decode_debug


def aggregate_rows(frame):
    result = {}

    for scope in ("exec5", "full10"):
        for source in ("model", "hold"):
            for group in ("all", "arm", "gripper"):
                column = f"{scope}_{source}_mae_{group}"
                result[column] = float(frame[column].mean())

        result[f"{scope}_win_rate_all"] = float(
            frame[f"{scope}_model_beats_hold_all"].mean()
        )
        result[f"{scope}_motion_mae_all"] = float(
            frame[f"{scope}_motion_mae_all"].mean()
        )
        result[f"{scope}_improvement_all"] = (
            result[f"{scope}_hold_mae_all"]
            - result[f"{scope}_model_mae_all"]
        )

    return result


def main():
    selections = load_all_episode_index()

    print("===== EVALUATION SELECTION =====", flush=True)
    print(
        pd.DataFrame(selections).to_string(index=False),
        flush=True,
    )

    config, policy, load_seconds = create_policy()
    rng = np.random.default_rng(RANDOM_SEED)

    records = []
    predictions = []
    ground_truths = []
    hold_predictions = []
    anchor_states = []
    trajectory_rows = []
    decode_records = []

    total_clips = len(selections)
    global_call_index = 0

    for clip_index, selection in enumerate(selections, start=1):
        episode_index = int(selection["episode_index"])
        smoke.EPISODE_INDEX = episode_index

        metadata = smoke.load_episode_metadata()
        episode = smoke.load_episode_data(metadata)

        states = episode["states"]
        actions = episode["actions"]
        episode_length = len(states)

        task_list = metadata["tasks"]
        if not task_list:
            raise RuntimeError(
                f"Episode {episode_index} has no task prompt"
            )
        prompt = str(task_list[0])

        if prompt != selection["task"]:
            raise RuntimeError(
                f"Task mismatch for episode {episode_index}: "
                f"{prompt!r} != {selection['task']!r}"
            )

        lower_anchor = abs(min(HISTORY_OFFSETS))
        upper_anchor = episode_length - ACTION_HORIZON

        if upper_anchor < lower_anchor:
            raise RuntimeError(
                f"Episode {episode_index} too short: "
                f"{episode_length}"
            )

        anchor = int(
            round(
                lower_anchor
                + selection["anchor_ratio"]
                * (upper_anchor - lower_anchor)
            )
        )

        history_indices = [
            anchor + offset
            for offset in HISTORY_OFFSETS
        ]

        images_by_frame, decode_debug = decode_clip_images(
            metadata,
            history_indices,
        )

        decode_records.append(
            {
                "episode_index": episode_index,
                "anchor": anchor,
                "history_indices": history_indices,
                "cameras": decode_debug,
            }
        )

        policy._model.reset_memory(policy.memory)

        print(
            f"[clip {clip_index:02d}/{total_clips}] "
            f"task={selection['task_id']} "
            f"episode={episode_index} "
            f"length={episode_length} "
            f"ratio={selection['anchor_ratio']:.2f} "
            f"frames={history_indices}",
            flush=True,
        )

        for memory_phase, frame_index in enumerate(history_indices):
            observation = {
                "state": states[frame_index].astype(
                    np.float32,
                    copy=False,
                ),
                "images": {
                    camera: images_by_frame[frame_index][camera]
                    for camera in smoke.CAMERAS
                },
                "prompt": prompt,
                "step": memory_phase,
            }

            noise = rng.standard_normal(
                (
                    ACTION_HORIZON,
                    int(config.model.action_dim),
                )
            ).astype(np.float32)

            start = time.monotonic()
            output = policy.infer(
                observation,
                noise=noise,
            )
            wall_ms = (
                time.monotonic() - start
            ) * 1000.0

            prediction = np.asarray(
                output["actions"],
                dtype=np.float32,
            )
            target = actions[
                frame_index : frame_index + ACTION_HORIZON
            ].astype(np.float32)
            hold = np.repeat(
                states[frame_index][None, :],
                ACTION_HORIZON,
                axis=0,
            ).astype(np.float32)

            expected_shape = (ACTION_HORIZON, 14)
            if prediction.shape != expected_shape:
                raise RuntimeError(
                    f"Prediction shape {prediction.shape}, "
                    f"expected {expected_shape}"
                )
            if target.shape != expected_shape:
                raise RuntimeError(
                    f"Target shape {target.shape}, "
                    f"expected {expected_shape}"
                )
            if not np.isfinite(prediction).all():
                raise RuntimeError(
                    f"Non-finite prediction at "
                    f"episode={episode_index}, frame={frame_index}"
                )

            exec5_metrics = array_errors(
                prediction,
                target,
                hold,
                EXECUTE_STEPS,
            )
            full10_metrics = array_errors(
                prediction,
                target,
                hold,
                ACTION_HORIZON,
            )

            policy_infer_ms = float(
                output.get(
                    "policy_timing",
                    {},
                ).get("infer_ms", float("nan"))
            )

            row = {
                "global_call_index": global_call_index,
                "clip_index": clip_index - 1,
                "task_id": int(selection["task_id"]),
                "task": prompt,
                "sample_id": int(selection["sample_id"]),
                "episode_index": episode_index,
                "episode_length": episode_length,
                "anchor_ratio": float(
                    selection["anchor_ratio"]
                ),
                "anchor_frame": anchor,
                "memory_phase": memory_phase,
                "frame_index": frame_index,
                "policy_infer_ms": policy_infer_ms,
                "wall_ms": wall_ms,
            }

            add_prefixed_metrics(
                row,
                "exec5",
                exec5_metrics,
            )
            add_prefixed_metrics(
                row,
                "full10",
                full10_metrics,
            )

            records.append(row)
            predictions.append(prediction)
            ground_truths.append(target)
            hold_predictions.append(hold)
            anchor_states.append(states[frame_index])

            for horizon_index in range(ACTION_HORIZON):
                for joint_index, joint_name in enumerate(
                    JOINT_NAMES
                ):
                    predicted_value = float(
                        prediction[horizon_index, joint_index]
                    )
                    target_value = float(
                        target[horizon_index, joint_index]
                    )
                    hold_value = float(
                        hold[horizon_index, joint_index]
                    )
                    signed_error = (
                        predicted_value - target_value
                    )

                    trajectory_rows.append(
                        {
                            "global_call_index": global_call_index,
                            "task_id": int(
                                selection["task_id"]
                            ),
                            "task": prompt,
                            "episode_index": episode_index,
                            "memory_phase": memory_phase,
                            "frame_index": frame_index,
                            "horizon_index": horizon_index,
                            "is_executed_step": (
                                horizon_index < EXECUTE_STEPS
                            ),
                            "joint_index": joint_index,
                            "joint_name": joint_name,
                            "joint_group": (
                                "gripper"
                                if joint_index in (6, 13)
                                else "arm"
                            ),
                            "prediction": predicted_value,
                            "ground_truth": target_value,
                            "hold_baseline": hold_value,
                            "signed_error": signed_error,
                            "absolute_error": abs(signed_error),
                        }
                    )

            print(
                f"  phase={memory_phase} "
                f"frame={frame_index} "
                f"exec5_mae={row['exec5_model_mae_all']:.6f} "
                f"hold={row['exec5_hold_mae_all']:.6f} "
                f"full10_mae={row['full10_model_mae_all']:.6f} "
                f"infer_ms={policy_infer_ms:.1f}",
                flush=True,
            )

            global_call_index += 1

    records_frame = pd.DataFrame(records)
    trajectory_frame = pd.DataFrame(trajectory_rows)

    prediction_array = np.stack(predictions)
    target_array = np.stack(ground_truths)
    hold_array = np.stack(hold_predictions)
    state_array = np.stack(anchor_states)

    records_path = OUTPUT / "records.csv"
    trajectory_path = OUTPUT / "trajectory_long.csv"
    prediction_path = OUTPUT / "predictions.npz"
    decode_path = OUTPUT / "video_decode.json"

    records_frame.to_csv(records_path, index=False)
    trajectory_frame.to_csv(trajectory_path, index=False)

    np.savez_compressed(
        prediction_path,
        predictions=prediction_array,
        ground_truth=target_array,
        hold_baseline=hold_array,
        states=state_array,
        episode_indices=records_frame[
            "episode_index"
        ].to_numpy(dtype=np.int64),
        frame_indices=records_frame[
            "frame_index"
        ].to_numpy(dtype=np.int64),
        memory_phases=records_frame[
            "memory_phase"
        ].to_numpy(dtype=np.int64),
        task_ids=records_frame[
            "task_id"
        ].to_numpy(dtype=np.int64),
    )

    decode_path.write_text(
        json.dumps(
            decode_records,
            indent=2,
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )

    full_history_frame = records_frame[
        records_frame["memory_phase"] == 2
    ].copy()

    overall_all_calls = aggregate_rows(records_frame)
    overall_full_history = aggregate_rows(
        full_history_frame
    )

    per_task_rows = []
    for (task_id, task), group in full_history_frame.groupby(
        ["task_id", "task"],
        sort=True,
    ):
        row = {
            "task_id": int(task_id),
            "task": task,
            "num_clips": int(len(group)),
        }
        row.update(aggregate_rows(group))
        per_task_rows.append(row)

    per_task_frame = pd.DataFrame(per_task_rows)
    per_task_path = OUTPUT / "per_task.csv"
    per_task_frame.to_csv(per_task_path, index=False)

    per_phase_rows = []
    for phase, group in records_frame.groupby(
        "memory_phase",
        sort=True,
    ):
        row = {
            "memory_phase": int(phase),
            "num_calls": int(len(group)),
        }
        row.update(aggregate_rows(group))
        per_phase_rows.append(row)

    per_phase_frame = pd.DataFrame(per_phase_rows)
    per_phase_path = OUTPUT / "per_memory_phase.csv"
    per_phase_frame.to_csv(per_phase_path, index=False)

    main_indices = full_history_frame.index.to_numpy(
        dtype=np.int64
    )
    main_prediction = prediction_array[main_indices]
    main_target = target_array[main_indices]
    main_hold = hold_array[main_indices]

    per_horizon_rows = []
    for horizon_index in range(ACTION_HORIZON):
        model_difference = (
            main_prediction[:, horizon_index]
            - main_target[:, horizon_index]
        )
        hold_difference = (
            main_hold[:, horizon_index]
            - main_target[:, horizon_index]
        )

        per_horizon_rows.append(
            {
                "horizon_index": horizon_index,
                "is_executed_step": (
                    horizon_index < EXECUTE_STEPS
                ),
                "model_mae_all": float(
                    np.mean(np.abs(model_difference))
                ),
                "hold_mae_all": float(
                    np.mean(np.abs(hold_difference))
                ),
                "model_mae_arm": float(
                    np.mean(
                        np.abs(
                            model_difference[:, ARM_DIMS]
                        )
                    )
                ),
                "hold_mae_arm": float(
                    np.mean(
                        np.abs(
                            hold_difference[:, ARM_DIMS]
                        )
                    )
                ),
                "model_mae_gripper": float(
                    np.mean(
                        np.abs(
                            model_difference[:, GRIPPER_DIMS]
                        )
                    )
                ),
                "hold_mae_gripper": float(
                    np.mean(
                        np.abs(
                            hold_difference[:, GRIPPER_DIMS]
                        )
                    )
                ),
            }
        )

    per_horizon_frame = pd.DataFrame(per_horizon_rows)
    per_horizon_path = OUTPUT / "per_horizon.csv"
    per_horizon_frame.to_csv(
        per_horizon_path,
        index=False,
    )

    per_dimension_rows = []
    for joint_index, joint_name in enumerate(JOINT_NAMES):
        model_difference = (
            main_prediction[:, :, joint_index]
            - main_target[:, :, joint_index]
        )
        hold_difference = (
            main_hold[:, :, joint_index]
            - main_target[:, :, joint_index]
        )

        per_dimension_rows.append(
            {
                "joint_index": joint_index,
                "joint_name": joint_name,
                "joint_group": (
                    "gripper"
                    if joint_index in (6, 13)
                    else "arm"
                ),
                "model_mae_full10": float(
                    np.mean(np.abs(model_difference))
                ),
                "hold_mae_full10": float(
                    np.mean(np.abs(hold_difference))
                ),
                "model_rmse_full10": float(
                    np.sqrt(
                        np.mean(np.square(model_difference))
                    )
                ),
                "hold_rmse_full10": float(
                    np.sqrt(
                        np.mean(np.square(hold_difference))
                    )
                ),
                "model_mae_exec5": float(
                    np.mean(
                        np.abs(
                            model_difference[:, :EXECUTE_STEPS]
                        )
                    )
                ),
                "hold_mae_exec5": float(
                    np.mean(
                        np.abs(
                            hold_difference[:, :EXECUTE_STEPS]
                        )
                    )
                ),
            }
        )

    per_dimension_frame = pd.DataFrame(
        per_dimension_rows
    )
    per_dimension_path = OUTPUT / "per_dimension.csv"
    per_dimension_frame.to_csv(
        per_dimension_path,
        index=False,
    )

    motion_threshold = float(
        full_history_frame[
            "full10_motion_mae_all"
        ].quantile(0.75)
    )
    high_motion_frame = full_history_frame[
        full_history_frame[
            "full10_motion_mae_all"
        ] >= motion_threshold
    ]
    high_motion_metrics = aggregate_rows(
        high_motion_frame
    )

    warm_timing_frame = records_frame[
        records_frame["global_call_index"] >= 3
    ]

    timing_summary = {
        "cold_first_three_wall_ms": records_frame.iloc[:3][
            "wall_ms"
        ].tolist(),
        "warm_num_calls": int(len(warm_timing_frame)),
        "warm_wall_ms_mean": float(
            warm_timing_frame["wall_ms"].mean()
        ),
        "warm_wall_ms_median": float(
            warm_timing_frame["wall_ms"].median()
        ),
        "warm_wall_ms_p95": float(
            warm_timing_frame["wall_ms"].quantile(0.95)
        ),
        "warm_policy_infer_ms_mean": float(
            warm_timing_frame["policy_infer_ms"].mean()
        ),
        "warm_policy_infer_ms_median": float(
            warm_timing_frame[
                "policy_infer_ms"
            ].median()
        ),
    }

    summary = {
        "evaluation_type": (
            "teacher-forced offline training-set replay"
        ),
        "closed_loop": False,
        "checkpoint": str(smoke.CKPT),
        "checkpoint_step": smoke.CHECKPOINT_STEP,
        "config_name": smoke.CONFIG_NAME,
        "dataset": str(smoke.DATASET),
        "num_tasks": 6,
        "num_clips": int(len(selections)),
        "calls_per_clip": 3,
        "num_inference_calls": int(len(records_frame)),
        "action_horizon": ACTION_HORIZON,
        "execute_steps": EXECUTE_STEPS,
        "history_offsets": list(HISTORY_OFFSETS),
        "random_seed": RANDOM_SEED,
        "policy_load_seconds": load_seconds,
        "overall_all_memory_phases": overall_all_calls,
        "overall_full_history_phase_2": overall_full_history,
        "high_motion_full_history": {
            "threshold": motion_threshold,
            "num_clips": int(len(high_motion_frame)),
            "metrics": high_motion_metrics,
        },
        "timing": timing_summary,
        "artifacts": {
            "records": str(records_path),
            "trajectory_long": str(trajectory_path),
            "predictions": str(prediction_path),
            "video_decode": str(decode_path),
            "per_task": str(per_task_path),
            "per_memory_phase": str(per_phase_path),
            "per_horizon": str(per_horizon_path),
            "per_dimension": str(per_dimension_path),
        },
    }

    summary_path = OUTPUT / "summary.json"
    summary_path.write_text(
        json.dumps(
            summary,
            indent=2,
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )

    figure, axes = plt.subplots(
        2,
        2,
        figsize=(16, 11),
    )

    x_task = np.arange(len(per_task_frame))
    width = 0.38
    axes[0, 0].bar(
        x_task - width / 2,
        per_task_frame["exec5_model_mae_all"],
        width,
        label="StreamPI",
    )
    axes[0, 0].bar(
        x_task + width / 2,
        per_task_frame["exec5_hold_mae_all"],
        width,
        label="Hold baseline",
    )
    axes[0, 0].set_xticks(x_task)
    axes[0, 0].set_xticklabels(
        [
            f"Task {task_id}"
            for task_id in per_task_frame["task_id"]
        ]
    )
    axes[0, 0].set_title(
        "Executed 5-step MAE by task (memory phase 2)"
    )
    axes[0, 0].set_ylabel("MAE")
    axes[0, 0].legend()
    axes[0, 0].grid(alpha=0.25)

    axes[0, 1].plot(
        per_horizon_frame["horizon_index"],
        per_horizon_frame["model_mae_all"],
        marker="o",
        label="StreamPI",
    )
    axes[0, 1].plot(
        per_horizon_frame["horizon_index"],
        per_horizon_frame["hold_mae_all"],
        marker="o",
        label="Hold baseline",
    )
    axes[0, 1].axvline(
        EXECUTE_STEPS - 0.5,
        color="black",
        linestyle="--",
        alpha=0.5,
        label="Execute boundary",
    )
    axes[0, 1].set_title(
        "MAE across action horizon (memory phase 2)"
    )
    axes[0, 1].set_xlabel("Horizon index")
    axes[0, 1].set_ylabel("MAE")
    axes[0, 1].legend()
    axes[0, 1].grid(alpha=0.25)

    x_joint = np.arange(len(per_dimension_frame))
    axes[1, 0].bar(
        x_joint - width / 2,
        per_dimension_frame["model_mae_full10"],
        width,
        label="StreamPI",
    )
    axes[1, 0].bar(
        x_joint + width / 2,
        per_dimension_frame["hold_mae_full10"],
        width,
        label="Hold baseline",
    )
    axes[1, 0].set_xticks(x_joint)
    axes[1, 0].set_xticklabels(
        per_dimension_frame["joint_name"],
        rotation=55,
        ha="right",
    )
    axes[1, 0].set_title(
        "10-step MAE by action dimension"
    )
    axes[1, 0].set_ylabel("MAE")
    axes[1, 0].legend()
    axes[1, 0].grid(alpha=0.25)

    axes[1, 1].plot(
        records_frame["global_call_index"],
        records_frame["wall_ms"],
        linewidth=1.0,
    )
    axes[1, 1].axvline(
        2.5,
        color="black",
        linestyle="--",
        alpha=0.5,
        label="Cold compile boundary",
    )
    axes[1, 1].set_title("Inference wall time")
    axes[1, 1].set_xlabel("Inference call")
    axes[1, 1].set_ylabel("Milliseconds")
    axes[1, 1].legend()
    axes[1, 1].grid(alpha=0.25)

    figure.suptitle(
        "StreamPI 10k ␔ training-set offline replay",
        fontsize=15,
    )
    figure.tight_layout()
    plot_path = OUTPUT / "evaluation_overview.png"
    figure.savefig(
        plot_path,
        dpi=180,
        bbox_inches="tight",
    )
    plt.close(figure)

    summary["artifacts"]["plot"] = str(plot_path)
    summary_path.write_text(
        json.dumps(
            summary,
            indent=2,
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )

    print("===== FULL-HISTORY MAIN METRICS =====")
    print(
        json.dumps(
            overall_full_history,
            indent=2,
            ensure_ascii=False,
        )
    )

    print("===== HIGH-MOTION METRICS =====")
    print(
        json.dumps(
            {
                "threshold": motion_threshold,
                "num_clips": len(high_motion_frame),
                "metrics": high_motion_metrics,
            },
            indent=2,
            ensure_ascii=False,
        )
    )

    print("===== TIMING =====")
    print(
        json.dumps(
            timing_summary,
            indent=2,
            ensure_ascii=False,
        )
    )

    print("===== OUTPUTS =====")
    for name, path in summary["artifacts"].items():
        print(f"{name}={path}")
    print(f"summary={summary_path}")
    print("FULL_EVAL_COMPLETE")


if __name__ == "__main__":
    main()