from __future__ import annotations

import argparse
import json
import pickle
from glob import glob
from pathlib import Path

import numpy as np
import joblib

try:
    import torch
except Exception:
    torch = None

GIT_GETUP_FOCUS_CLIP = "fallAndGetUp2_subject2_1200_1370"


def _load_payload(path: Path):
    if path.suffix.lower() in (".pt", ".pth"):
        if torch is None:
            raise RuntimeError("Loading .pt/.pth requires torch to be installed.")
        return torch.load(str(path), map_location="cpu")
    if path.suffix.lower() in (".pkl",):
        try:
            with path.open("rb") as f:
                return pickle.load(f)
        except Exception:
            return joblib.load(str(path))
    if path.suffix.lower() in (".json",):
        with path.open("r", encoding="utf-8") as f:
            return json.load(f)
    raise ValueError(f"Unsupported file type: {path.suffix}")


def _to_numpy(payload, key_hint: str = "") -> np.ndarray:
    if torch is not None and isinstance(payload, torch.Tensor):
        return payload.detach().cpu().numpy().astype(np.float32)
    if isinstance(payload, np.ndarray):
        return payload.astype(np.float32)
    if isinstance(payload, list):
        return np.asarray(payload, dtype=np.float32)
    if isinstance(payload, dict):
        candidate_keys = [key_hint] if key_hint else []
        candidate_keys += [
            "expert_features",
            "features",
            "observations",
            "obs",
            "joint_pos",
            "motion",
            "data",
        ]
        for key in candidate_keys:
            if key and key in payload:
                value = payload[key]
                if torch is not None and isinstance(value, torch.Tensor):
                    return value.detach().cpu().numpy().astype(np.float32)
                if isinstance(value, np.ndarray):
                    return value.astype(np.float32)
                if isinstance(value, list):
                    return np.asarray(value, dtype=np.float32)
        for value in payload.values():
            if torch is not None and isinstance(value, torch.Tensor):
                return value.detach().cpu().numpy().astype(np.float32)
            if isinstance(value, np.ndarray):
                return value.astype(np.float32)
            if isinstance(value, list):
                return np.asarray(value, dtype=np.float32)
    raise ValueError("Could not locate tensor-like payload. Provide --input_key when needed.")


def _ensure_2d(x: np.ndarray) -> np.ndarray:
    if x.ndim == 1:
        return x[None, :]
    if x.ndim != 2:
        raise ValueError(f"Expected 2D array [N, D], got shape={tuple(x.shape)}")
    return x.astype(np.float32)


def _from_legged_lab_motion_dict(payload: dict, target_dof: int) -> np.ndarray:
    """
    Convert legged_lab motion pickle payload to AMP feature tensor [N, D].
    Expected keys: fps, root_pos, dof_pos.
    """
    dof_pos = np.asarray(payload.get("dof_pos", None))
    root_pos = np.asarray(payload.get("root_pos", None))
    fps = float(payload.get("fps", 30.0))
    if dof_pos.ndim != 2:
        raise ValueError("legged_lab payload missing valid dof_pos [N, M].")
    if root_pos.ndim != 2 or root_pos.shape[0] != dof_pos.shape[0]:
        root_pos = np.zeros((dof_pos.shape[0], 3), dtype=np.float32)

    dt = 1.0 / max(fps, 1e-3)
    dof_pos = dof_pos.astype(np.float32)
    root_pos = root_pos.astype(np.float32)

    # Current get_up AMP feature dim expects 23 dof by default.
    target_dof = int(target_dof)
    if dof_pos.shape[1] >= target_dof:
        dof_pos = dof_pos[:, :target_dof]
    else:
        dof_pos = np.pad(dof_pos, ((0, 0), (0, target_dof - dof_pos.shape[1])), mode="constant")

    dof_vel = np.zeros_like(dof_pos, dtype=np.float32)
    if dof_pos.shape[0] > 1:
        dof_vel[1:] = (dof_pos[1:] - dof_pos[:-1]) / dt
        dof_vel[0] = dof_vel[1]

    root_lin = np.zeros((dof_pos.shape[0], 3), dtype=np.float32)
    if root_pos.shape[0] > 1:
        root_lin[1:] = (root_pos[1:] - root_pos[:-1]) / dt
        root_lin[0] = root_lin[1]
    root_ang = np.zeros((dof_pos.shape[0], 3), dtype=np.float32)
    gravity = np.zeros((dof_pos.shape[0], 3), dtype=np.float32)
    gravity[:, 2] = -1.0

    x = np.concatenate([dof_pos, dof_vel, root_lin, root_ang, gravity], axis=-1).astype(np.float32)
    return x


def _normalize_dim(x: np.ndarray, feature_dim: int) -> np.ndarray:
    if x.shape[1] == feature_dim:
        return x
    if x.shape[1] > feature_dim:
        return x[:, :feature_dim]
    pad = np.zeros((x.shape[0], feature_dim - x.shape[1]), dtype=np.float32)
    return np.concatenate([x, pad], axis=-1)


def _collect_input_files(input_path: Path, glob_pattern: str) -> list[Path]:
    if input_path.is_file():
        return [input_path]
    if input_path.is_dir():
        files = sorted([Path(p) for p in glob(str(input_path / glob_pattern))])
        if files:
            return files
    raise FileNotFoundError(f"No source files found for input={input_path} pattern={glob_pattern}")


def _build_clip_weights(clip_names: list[str], weight_mode: str) -> np.ndarray:
    if len(clip_names) == 0:
        return np.zeros((0,), dtype=np.float32)
    if weight_mode == "uniform":
        return np.ones((len(clip_names),), dtype=np.float32)
    if weight_mode == "git_getup_focus":
        weights = np.array([1.0 if GIT_GETUP_FOCUS_CLIP in n else 0.0 for n in clip_names], dtype=np.float32)
        if float(np.sum(weights)) <= 0.0:
            return np.ones((len(clip_names),), dtype=np.float32)
        return weights
    return np.ones((len(clip_names),), dtype=np.float32)


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Template converter: migrate legged_lab (or other) motion/expert data "
            "to AMP expert_features.pt for rl_game/get_up."
        )
    )
    parser.add_argument("--input", required=True, type=str, help="Path to source motion/expert file or directory.")
    parser.add_argument("--output", required=True, type=str, help="Output path for expert_features.pt.")
    parser.add_argument(
        "--input_key",
        type=str,
        default="",
        help="Optional key to locate tensor inside input payload dict.",
    )
    parser.add_argument(
        "--feature_dim",
        type=int,
        default=55,
        help="Target AMP feature dimension. For current get_up config, default=55.",
    )
    parser.add_argument(
        "--input_glob",
        type=str,
        default="*.pkl",
        help="Glob used when --input is a directory.",
    )
    parser.add_argument(
        "--target_dof",
        type=int,
        default=23,
        help="Target dof count when converting legged_lab pkl dof_pos.",
    )
    parser.add_argument(
        "--clip_weight_mode",
        type=str,
        default="git_getup_focus",
        choices=["git_getup_focus", "uniform"],
        help="Clip sampling weight mode for expert sequence training.",
    )
    parser.add_argument("--repeat", type=int, default=1, help="Repeat samples for small datasets.")
    args = parser.parse_args()

    in_path = Path(args.input).expanduser().resolve()
    input_files = _collect_input_files(in_path, args.input_glob)
    clip_arrays: list[np.ndarray] = []
    clip_names: list[str] = []
    for f in input_files:
        payload = _load_payload(f)
        if isinstance(payload, dict) and "dof_pos" in payload and f.suffix.lower() == ".pkl":
            x = _from_legged_lab_motion_dict(payload, target_dof=int(args.target_dof))
        else:
            x = _to_numpy(payload, key_hint=args.input_key)
            x = _ensure_2d(x)
        x = _normalize_dim(x.astype(np.float32), int(args.feature_dim))
        clip_arrays.append(x)
        clip_names.append(f.stem)

    repeat = max(int(args.repeat), 1)
    if repeat > 1:
        clip_arrays = clip_arrays * repeat
        clip_names = clip_names * repeat

    clip_weights = _build_clip_weights(clip_names, weight_mode=args.clip_weight_mode)
    x = np.concatenate(clip_arrays, axis=0).astype(np.float32)
    clip_lengths = [int(c.shape[0]) for c in clip_arrays]
    clip_offsets: list[int] = []
    cursor = 0
    for length in clip_lengths:
        clip_offsets.append(cursor)
        cursor += length
    repeat = max(int(args.repeat), 1)

    out_path = Path(args.output).expanduser().resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "expert_features": x,
        "expert_clips": clip_arrays,
        "clip_names": clip_names,
        "clip_weights": clip_weights,
        "clip_lengths": clip_lengths,
        "clip_offsets": clip_offsets,
        "meta": {
            "source": str(in_path),
            "source_count": len(input_files),
            "input_key": args.input_key,
            "feature_dim": int(args.feature_dim),
            "target_dof": int(args.target_dof),
            "input_glob": args.input_glob,
            "clip_weight_mode": args.clip_weight_mode,
            "git_getup_focus_clip": GIT_GETUP_FOCUS_CLIP,
            "repeat": int(repeat),
        },
    }
    if torch is not None:
        torch.save(payload, str(out_path))
    else:
        with out_path.open("wb") as f:
            pickle.dump(payload, f)
    print(f"[OK] Saved AMP expert features -> {out_path} shape={tuple(x.shape)}")


if __name__ == "__main__":
    main()
