import argparse
import random
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import matplotlib.pyplot as plt
from pat.utils import eef_euler_to_6d, abs_6d_2_abs_euler
from pat.models.pat_tokenizer import PATTokenizer, TokenizerType


def _pick_random_episode(dataset_root: Path) -> Path:
    data_dir = dataset_root / "data"
    if not data_dir.is_dir():
        raise FileNotFoundError(f"Missing data dir: {data_dir}")
    chunk_dirs = sorted(p for p in data_dir.iterdir() if p.is_dir() and p.name.startswith("chunk-"))
    if not chunk_dirs:
        raise FileNotFoundError(f"No chunk-* dirs in {data_dir}")
    random.shuffle(chunk_dirs)
    for chunk in chunk_dirs:
        episodes = sorted(chunk.glob("episode_*.parquet"))
        if episodes:
            return random.choice(episodes)
    raise FileNotFoundError(f"No episode parquet files under {data_dir}")


def _load_eepose(parquet_path: Path, column: str) -> np.ndarray:
    table = pq.read_table(parquet_path.as_posix(), columns=[column])
    data = table.column(0).to_pylist()
    arr = np.asarray(data, dtype=np.float32)
    if arr.ndim != 2:
        raise ValueError(f"Expected 2D array from {column}, got shape {arr.shape}")
    return arr


def _ensure_batch_first(arr: np.ndarray) -> np.ndarray:
    if arr.ndim == 2:
        return arr[None, ...]
    if arr.ndim == 3:
        return arr
    raise ValueError(f"Expected 2D or 3D array, got shape {arr.shape}")


def _convert_states_to_6d(euler_states: np.ndarray) -> np.ndarray:
    states_6d = []
    for state_14d in euler_states:
        state_20d = eef_euler_to_6d(state_14d[0:7].tolist(), state_14d[7:14].tolist())
        states_6d.append(state_20d)
    return np.stack(states_6d, axis=0)


def _convert_6d_to_euler(states_6d: np.ndarray) -> np.ndarray:
    states_euler = []
    for state_20d in states_6d:
        state_14d = abs_6d_2_abs_euler(state_20d)
        states_euler.append(state_14d)
    return np.stack(states_euler, axis=0)


def _chunk_sequence(
    data: np.ndarray, chunk_size: int, drop_last: bool = False
) -> tuple[list[np.ndarray], list[int]]:
    chunks = []
    boundaries = []
    for start in range(0, len(data), chunk_size):
        chunk = data[start : start + chunk_size]
        if len(chunk) == 0:
            continue
        if drop_last and len(chunk) < chunk_size:
            break
        if start != 0:
            boundaries.append(start)
        chunks.append(chunk)
    return chunks, boundaries


def _plot_left_right(
    orig: np.ndarray,
    recon: np.ndarray,
    out_path: Path,
    title: str,
    left_labels,
    right_labels,
    chunk_boundaries=None,
):
    if orig.shape[1] != recon.shape[1]:
        raise ValueError(f"Dim mismatch: orig {orig.shape[1]} vs recon {recon.shape[1]}")
    if len(left_labels) != len(right_labels):
        raise ValueError("Left/right label count mismatch")

    n_rows = len(left_labels)
    fig, axes = plt.subplots(n_rows, 2, figsize=(14, 2.2 * n_rows), sharex=True)
    axes = np.asarray(axes).reshape(-1)

    for r in range(n_rows):
        left_idx = r
        right_idx = r + n_rows

        ax_l = axes[r * 2]
        ax_l.plot(orig[:, left_idx], label="original", linewidth=1.0)
        ax_l.plot(recon[:, left_idx], label="decoded", linewidth=1.0)
        if chunk_boundaries:
            ax_l.scatter(
                chunk_boundaries,
                orig[chunk_boundaries, left_idx],
                color="red",
                s=10,
                zorder=5,
            )
        ax_l.set_title(f"L {left_labels[r]}")
        if r == 0:
            ax_l.legend(loc="upper right", fontsize=8)

        ax_r = axes[r * 2 + 1]
        ax_r.plot(orig[:, right_idx], label="original", linewidth=1.0)
        ax_r.plot(recon[:, right_idx], label="decoded", linewidth=1.0)
        if chunk_boundaries:
            ax_r.scatter(
                chunk_boundaries,
                orig[chunk_boundaries, right_idx],
                color="red",
                s=10,
                zorder=5,
            )
        ax_r.set_title(f"R {right_labels[r]}")

    fig.suptitle(title)
    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path.as_posix(), dpi=150)


def main():
    parser = argparse.ArgumentParser(description="Test fast tokenizer on lerobot eepose data")
    parser.add_argument("--tokenizer", default="fast", choices=["fast", "oat"], help="Tokenizer type")
    parser.add_argument("--model", default=None, help="Path to tokenizer model (fast or oat)")
    parser.add_argument("--dataset_root", default="data/batchdata", help="Root dir with lerobot datasets")
    parser.add_argument("--dataset", default=None, help="Dataset folder name under dataset_root")
    parser.add_argument("--episode", default=None, help="Explicit episode parquet path")
    parser.add_argument("--column", default="observation.state_eepose", help="Column to evaluate")
    parser.add_argument("--out_dir", default="tests", help="Output dir for plots")
    parser.add_argument("--chunk_size", default=42, type=int, help="Chunk size for chunked 6D encode/decode")
    args = parser.parse_args()

    dataset_root = Path(args.dataset_root)
    if args.episode:
        episode_path = Path(args.episode)
    else:
        if args.dataset:
            dataset_dir = dataset_root / args.dataset
        else:
            datasets = sorted(
                p
                for p in dataset_root.iterdir()
                if p.is_dir() and not p.name.startswith(".") and (p / "data").is_dir()
            )
            if not datasets:
                raise FileNotFoundError(f"No datasets found under {dataset_root}")
            dataset_dir = random.choice(datasets)
        episode_path = _pick_random_episode(dataset_dir)

    print(f"Using episode: {episode_path}")
    data_euler = _load_eepose(episode_path, args.column)
    print(f"Loaded {data_euler.shape[0]} timesteps, dim={data_euler.shape[1]}")

    if args.model is None:
        args.model = "output/fast" if args.tokenizer == "fast" else "output/oat_cb1000_h42"

    if args.tokenizer == "fast":
        tokenizer = PATTokenizer.from_pretrained(
            args.model, tokenizer_type=TokenizerType.FAST, device="cpu"
        )
    else:
        tokenizer = PATTokenizer.from_pretrained(args.model, device="cpu")

    out_dir = Path(args.out_dir)
    euler_labels = ["x", "y", "z", "rx", "ry", "rz", "gripper"]

    orig_euler = None
    recon_euler = None
    orig_6d = None
    recon_6d = None

    if args.tokenizer == "fast":
        # ===== Euler pipeline (14D) =====
        batch_euler = _ensure_batch_first(data_euler)
        tokens_euler = tokenizer.encode(batch_euler)
        decoded_euler, _ = tokenizer.decode(
            tokens=tokens_euler,
            time_horizon=batch_euler.shape[1],
            action_dim=batch_euler.shape[2],
        )
        decoded_euler = np.asarray(decoded_euler)
        if decoded_euler.ndim == 3:
            decoded_euler = decoded_euler[0]
        if decoded_euler.ndim != 2:
            raise ValueError(f"Decoded euler result is not 2D: shape {decoded_euler.shape}")

        min_len = min(data_euler.shape[0], decoded_euler.shape[0])
        orig_euler = data_euler[:min_len]
        recon_euler = decoded_euler[:min_len]
        mse_euler = float(np.mean((orig_euler - recon_euler) ** 2))
        print(f"Euler MSE (all dims): {mse_euler:.8f}")

        euler_plot = out_dir / "fast_tokenizer_compare_euler.png"
        _plot_left_right(
            orig_euler,
            recon_euler,
            euler_plot,
            f"{args.column} original vs decoded (euler)",
            euler_labels,
            euler_labels,
        )
        print(f"Saved plot: {euler_plot}")

    # ===== 6D pipeline (20D) =====
    data_6d = _convert_states_to_6d(data_euler)
    batch_6d = _ensure_batch_first(data_6d)

    if args.tokenizer == "oat":
        # OAT expects fixed horizon; use chunked path below for full sequence
        pass
    else:
        tokens_6d = tokenizer.encode(batch_6d)
        decoded_6d, _ = tokenizer.decode(
            tokens=tokens_6d,
            time_horizon=batch_6d.shape[1],
            action_dim=batch_6d.shape[2],
        )
        decoded_6d = np.asarray(decoded_6d)
        if decoded_6d.ndim == 3:
            decoded_6d = decoded_6d[0]
        if decoded_6d.ndim != 2:
            raise ValueError(f"Decoded 6D result is not 2D: shape {decoded_6d.shape}")

        min_len_6d = min(data_6d.shape[0], decoded_6d.shape[0])
        orig_6d = data_6d[:min_len_6d]
        recon_6d = decoded_6d[:min_len_6d]
        mse_6d = float(np.mean((orig_6d - recon_6d) ** 2))
        print(f"6D MSE (all dims): {mse_6d:.8f}")

        sixd_labels = ["x", "y", "z", "r6d_0", "r6d_1", "r6d_2", "r6d_3", "r6d_4", "r6d_5", "gripper"]
        sixd_plot = out_dir / "fast_tokenizer_compare_6d.png"
        _plot_left_right(
            orig_6d,
            recon_6d,
            sixd_plot,
            f"{args.column} original vs decoded (6d)",
            sixd_labels,
            sixd_labels,
        )
        print(f"Saved plot: {sixd_plot}")

    # ===== Chunked 6D pipeline, plot in euler =====
    if args.tokenizer == "oat":
        chunk_size = tokenizer.action_space.horizon
    else:
        chunk_size = args.chunk_size

    chunks_6d, boundaries = _chunk_sequence(
        data_6d, chunk_size, drop_last=(args.tokenizer == "oat")
    )
    recon_chunks_6d = []
    for chunk in chunks_6d:
        batch_chunk = _ensure_batch_first(chunk)
        if args.tokenizer == "oat":
            latents = tokenizer.encode(batch_chunk)
            decoded_chunk, _ = tokenizer.decode(latents=latents)
        else:
            tokens_chunk = tokenizer.encode(batch_chunk)
            decoded_chunk, _ = tokenizer.decode(
                tokens=tokens_chunk,
                time_horizon=batch_chunk.shape[1],
                action_dim=batch_chunk.shape[2],
            )
        decoded_chunk = np.asarray(decoded_chunk)
        if decoded_chunk.ndim == 3:
            decoded_chunk = decoded_chunk[0]
        if decoded_chunk.ndim != 2:
            raise ValueError(f"Decoded chunk result is not 2D: shape {decoded_chunk.shape}")
        recon_chunks_6d.append(decoded_chunk)

    recon_6d_chunked = np.concatenate(recon_chunks_6d, axis=0)
    recon_euler_chunked = _convert_6d_to_euler(recon_6d_chunked)
    min_len_chunked = min(data_euler.shape[0], recon_euler_chunked.shape[0])
    orig_euler_chunked = data_euler[:min_len_chunked]
    recon_euler_chunked = recon_euler_chunked[:min_len_chunked]
    orig_6d_chunked = data_6d[:min_len_chunked]
    recon_6d_chunked = recon_6d_chunked[:min_len_chunked]
    boundaries = [b for b in boundaries if b < min_len_chunked]
    mse_chunked_6d = float(np.mean((orig_6d_chunked - recon_6d_chunked) ** 2))
    print(f"Chunked 6D MSE (all dims): {mse_chunked_6d:.8f}")
    mse_chunked_euler = float(np.mean((orig_euler_chunked - recon_euler_chunked) ** 2))
    print(f"Chunked 6D->Euler MSE (all dims): {mse_chunked_euler:.8f}")

    chunked_plot = out_dir / "fast_tokenizer_compare_chunked_6d_to_euler.png"
    _plot_left_right(
        orig_euler_chunked,
        recon_euler_chunked,
        chunked_plot,
        f"{args.column} original vs decoded (chunked 6d -> euler, chunk={chunk_size})",
        euler_labels,
        euler_labels,
        chunk_boundaries=boundaries,
    )
    print(f"Saved plot: {chunked_plot}")

    # Ensure all three plots exist for OAT (derived from chunked recon)
    if args.tokenizer == "oat":
        orig_euler = orig_euler_chunked
        recon_euler = recon_euler_chunked
        orig_6d = orig_6d_chunked
        recon_6d = recon_6d_chunked

        euler_plot = out_dir / "fast_tokenizer_compare_euler.png"
        _plot_left_right(
            orig_euler,
            recon_euler,
            euler_plot,
            f"{args.column} original vs decoded (euler, from chunked 6d)",
            euler_labels,
            euler_labels,
        )
        print(f"Saved plot: {euler_plot}")

        sixd_labels = ["x", "y", "z", "r6d_0", "r6d_1", "r6d_2", "r6d_3", "r6d_4", "r6d_5", "gripper"]
        sixd_plot = out_dir / "fast_tokenizer_compare_6d.png"
        _plot_left_right(
            orig_6d,
            recon_6d,
            sixd_plot,
            f"{args.column} original vs decoded (6d, chunked)",
            sixd_labels,
            sixd_labels,
        )
        print(f"Saved plot: {sixd_plot}")


if __name__ == "__main__":
    main()
