"""Visualize phase annotations: plot gripper signal with subgoal color bands."""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import pandas as pd


# Color palette for subgoals
COLORS = [
    "#4C72B0", "#DD8452", "#55A868", "#C44E52",
    "#8172B3", "#937860", "#DA8BC3", "#8C8C8C",
    "#CCB974", "#64B5CD", "#E377C2", "#7F7F7F",
]


def plot_episode(
    parquet_path: Path,
    output_path: Path | None = None,
    signal_source: str = "action",
    figsize: tuple = (16, 4),
):
    """Plot phase annotation for a single episode."""
    df = pd.read_parquet(parquet_path)

    # Extract signals
    state = np.array(df["observation.state"].tolist())
    action = np.array(df["action"].tolist())
    gripper_state = state[:, 7]   # modality: gripper dim
    gripper_action = action[:, 6]

    if signal_source == "action":
        gripper_signal = gripper_action
        ylabel = "gripper action (1=open, 0=closed)"
    else:
        gripper_signal = gripper_state
        ylabel = "gripper state (dim 7)"

    T = len(df)
    frames = np.arange(T)

    subgoals = df["phase.subgoal"].tolist()
    progress = df["phase.progress"].tolist()

    ep_idx = int(df["episode_index"].iloc[0])
    task_idx = int(df["task_index"].iloc[0])

    # Assign colors to unique subgoals (preserving order)
    unique_sgs = []
    for sg in subgoals:
        if sg not in unique_sgs:
            unique_sgs.append(sg)
    sg_to_color = {sg: COLORS[i % len(COLORS)] for i, sg in enumerate(unique_sgs)}

    # Find subgoal spans
    spans = []
    span_start = 0
    current_sg = subgoals[0]
    for t in range(1, T):
        if subgoals[t] != current_sg:
            spans.append((span_start, t, current_sg))
            current_sg = subgoals[t]
            span_start = t
    spans.append((span_start, T, current_sg))

    # Find early/late boundaries within each subgoal span
    early_late_boundaries = []
    for s, e, sg in spans:
        for t in range(s + 1, e):
            if progress[t] != progress[t - 1]:
                early_late_boundaries.append(t)

    # Find gripper transitions (on action signal)
    close_events = []
    open_events = []
    for t in range(1, T):
        if gripper_action[t - 1] > 0.5 and gripper_action[t] < 0.5:
            close_events.append(t)
        elif gripper_action[t - 1] < 0.5 and gripper_action[t] > 0.5:
            open_events.append(t)

    # --- Plot ---
    fig, ax = plt.subplots(figsize=figsize)

    # Background color bands for subgoals
    for s, e, sg in spans:
        ax.axvspan(s, e, alpha=0.2, color=sg_to_color[sg])

    # Gripper signal line
    ax.plot(frames, gripper_signal, color="black", linewidth=1.0, alpha=0.8)

    # Early/late dividers
    for b in early_late_boundaries:
        ax.axvline(b, color="gray", linestyle="--", linewidth=0.7, alpha=0.6)

    # Gripper events
    y_range = gripper_signal.max() - gripper_signal.min()
    y_top = gripper_signal.max() + 0.05 * y_range
    for t in close_events:
        ax.plot(t, y_top, marker="v", color="red", markersize=6)
    for t in open_events:
        ax.plot(t, y_top, marker="^", color="green", markersize=6)

    # Labels
    ax.set_xlabel("Frame index")
    ax.set_ylabel(ylabel)
    ax.set_title(f"Episode {ep_idx} (task {task_idx})")
    ax.set_xlim(0, T)

    # Legend
    patches = [mpatches.Patch(color=sg_to_color[sg], alpha=0.3, label=sg) for sg in unique_sgs]
    patches.append(plt.Line2D([], [], color="gray", linestyle="--", label="early/late boundary"))
    patches.append(plt.Line2D([], [], marker="v", color="red", linestyle="", label="close event"))
    patches.append(plt.Line2D([], [], marker="^", color="green", linestyle="", label="open event"))
    ax.legend(handles=patches, loc="upper right", fontsize=7, ncol=2)

    plt.tight_layout()
    if output_path:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        print(f"Saved to {output_path}")
    else:
        plt.show()
    plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description="Visualize phase annotations for episodes.")
    parser.add_argument("--data_dir", type=str, required=True,
                        help="Annotated dataset directory (output of phase_annotate)")
    parser.add_argument("--episode", type=int, nargs="+", default=None,
                        help="Episode indices to plot (default: first 5)")
    parser.add_argument("--output_dir", type=str, default=None,
                        help="Save plots to this directory (default: show interactively)")
    parser.add_argument("--signal_source", type=str, default="action",
                        choices=["action", "state"])
    args = parser.parse_args()

    data_dir = Path(args.data_dir)
    parquets = sorted((data_dir / "data").rglob("episode_*.parquet"))

    if not parquets:
        print(f"No episode parquets found in {data_dir}")
        return

    # Select episodes
    if args.episode is not None:
        selected = []
        for ep_idx in args.episode:
            name = f"episode_{ep_idx:06d}.parquet"
            matches = [p for p in parquets if p.name == name]
            if matches:
                selected.append(matches[0])
            else:
                print(f"Episode {ep_idx} not found")
        parquets = selected
    else:
        parquets = parquets[:5]

    output_dir = Path(args.output_dir) if args.output_dir else None

    for p in parquets:
        out_path = None
        if output_dir:
            out_path = output_dir / f"{p.stem}.png"
        plot_episode(p, out_path, signal_source=args.signal_source)


if __name__ == "__main__":
    main()