"""Adapters to convert VLTK-format H5 files to the optimizer's expected input format.

Each loader function reads a VLTK extractor H5 and returns a dict compatible
with the SMPLFX optimisation pipeline (``reoptimize_smplfx.py`` conventions).
"""

from __future__ import annotations

import logging
from pathlib import Path

import h5py
import hdf5plugin  # noqa: F401 – required for compressed datasets
import numpy as np
import torch

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# TEASER (face_teaser.h5) – nearly identical to original readH5FileFLAME
# ---------------------------------------------------------------------------

def load_teaser_for_fitting(
    h5_path: str | Path,
    median_shape: bool = True,
    jaw_pose: str = "euler",
    conv: str = "ZYX",
) -> dict:
    """Load TEASER FLAME H5 in the format expected by the optimizer.

    Args:
        h5_path: Path to ``face_teaser.h5``.
        median_shape: Collapse shape to a single median vector.
        jaw_pose: Target rotation format for jaw_pose (``'axis'`` or ``'euler'``).
        conv: Euler angle convention when ``jaw_pose='euler'``.

    Returns:
        Dict with keys ``cam``, ``pose_params``, ``successful_frames``,
        ``flame`` (sub-dict with ``shape``, ``expression``, ``jaw_pose``,
        ``eyelid_params``), and optionally ``kpts_2d`` / ``boxes``.
    """
    from .maths_utils import axis_to_euler_pytorch3d

    h5_path = str(h5_path)
    out: dict = {}

    with h5py.File(h5_path, "r") as hf:
        out["meta"] = {k: hf.attrs[k] for k in hf.attrs.keys()}

        for name in ("cam", "pose_params", "successful_frames"):
            assert name in hf, f"Missing dataset: {name}"
            out[name] = hf[name][()]

        if "kpts_2d" in hf:
            out["kpts_2d"] = torch.FloatTensor(hf["kpts_2d"][()])
        else:
            out["boxes"] = hf["boxes"][()]
            out["boxes_scores"] = hf["boxes_scores"][()]

        assert "flame" in hf, "Missing group: flame"
        out["flame"] = {}
        for name in ("shape", "expression", "jaw_pose", "eyelid_params"):
            assert name in hf["flame"], f"Missing dataset in flame group: {name}"
            out["flame"][name] = torch.FloatTensor(hf["flame"][name][()])

    if median_shape:
        out["flame"]["shape"] = out["flame"]["shape"].median(0).values[None]

    assert len(out["successful_frames"]) == out["pose_params"].shape[0]

    if jaw_pose == "axis":
        pass
    elif jaw_pose == "euler":
        out["flame"]["jaw_pose"] = axis_to_euler_pytorch3d(
            out["flame"]["jaw_pose"], conv
        ).squeeze(1)
    else:
        raise NotImplementedError(f"Unsupported jaw_pose format: {jaw_pose}")

    return out


# ---------------------------------------------------------------------------
# NLF (body_nlf.h5) – adapts VLTK keys to optimizer expectations
# ---------------------------------------------------------------------------

def load_nlf_for_fitting(
    h5_path: str | Path,
    angle_format: str = "euler",
    conv: str = "ZYX",
) -> dict:
    """Load NLF SMPL-X H5 in the format expected by the optimizer.

    Key differences from the original ``h5_load_smpl``:
    - VLTK uses ``img_idx`` (not ``frame_idx``) and ``trans`` (not ``transl``).
    - VLTK does not write an ``_attrs`` group; metadata is in HDF5 attributes.

    Args:
        h5_path: Path to ``body_nlf.h5``.
        angle_format: Rotation representation (``'euler'``, ``'axis'``, ``'6d'``, ``'mat'``).
        conv: Euler angle convention.

    Returns:
        Wrapper dict with ``smplx`` sub-dict, ``_attrs``, ``frame_idx``, etc.
    """
    from .maths_utils import axis_to_euler_pytorch3d, axis_to_6d_pytorch3d
    import pytorch3d

    with h5py.File(str(h5_path), "r") as hf:
        attrs = {k: v for k, v in hf.attrs.items()}

        smplx_data: dict = {}
        if "smplx" in hf:
            for key in hf["smplx"].keys():
                smplx_data[key] = hf["smplx"][key][()]
        else:
            raise ValueError(f"NLF H5 missing 'smplx' group: {h5_path}")

        # Build img_idx → frame_idx mapping
        if "img_idx" in hf:
            img_idx = hf["img_idx"][()]
        elif "frame_idx" in hf:
            img_idx = hf["frame_idx"][()]
        else:
            img_idx = np.arange(smplx_data["body_pose"].shape[0])

        extra: dict = {}
        for key in ("joints2d", "joints3d"):
            if key in hf:
                extra[key] = hf[key][()]

    # Convert angular keys
    angular_keys = [
        "body_pose", "global_orient", "left_hand_pose",
        "right_hand_pose", "jaw_pose", "leye_pose", "reye_pose",
    ]
    for akey in angular_keys:
        if akey not in smplx_data:
            continue
        raw = torch.from_numpy(smplx_data[akey]).float()
        if angle_format == "axis":
            smplx_data[akey] = raw.reshape(len(raw), -1)
        elif angle_format == "6d":
            smplx_data[akey] = axis_to_6d_pytorch3d(raw).reshape(len(raw), -1)
        elif angle_format == "euler":
            smplx_data[akey] = axis_to_euler_pytorch3d(raw, conv).reshape(len(raw), -1)
        elif angle_format == "mat":
            smplx_data[akey] = pytorch3d.transforms.axis_angle_to_matrix(
                raw.reshape(len(raw), -1, 3)
            ).flatten(1, -1)
        else:
            raise NotImplementedError(f"Unsupported angle_format: {angle_format}")

    # Ensure float32 tensors
    for key in list(smplx_data.keys()):
        v = smplx_data[key]
        if not torch.is_tensor(v):
            smplx_data[key] = torch.from_numpy(v).float()

    # Rename trans → transl (optimizer expects 'transl')
    if "trans" in smplx_data and "transl" not in smplx_data:
        smplx_data["transl"] = smplx_data.pop("trans")

    data = {
        "smplx": smplx_data,
        "_attrs": attrs,
        "frame_idx": img_idx,
        **extra,
    }

    class _Container:
        """Minimal container mimicking H5Features."""
        pass

    result = _Container()
    result.data = data
    result.meta = attrs
    result.num_frames = len(img_idx)
    return result


# ---------------------------------------------------------------------------
# WiLoR (hands_wilor.h5) – restructures per-detection to per-frame
# ---------------------------------------------------------------------------

def load_wilor_for_fitting(
    h5_path: str | Path,
) -> dict:
    """Load WiLoR MANO H5 in the format expected by the optimizer.

    VLTK WiLoR H5 stores:
    - ``mano/betas``            (N_det, 10)
    - ``mano/global_orient``    (N_det, 1, 3, 3)
    - ``mano/hand_pose``        (N_det, 15, 3, 3)
    - ``right``                 (N_det,) boolean
    - ``img_idx``               (N_det, 2)  [start, count] per frame

    The optimizer expects per-frame dicts keyed by ``mano``.

    Returns:
        Container with ``.data`` dict compatible with ``h5_load_mano`` output.
    """
    from .utils_smplx import flipPoseRot

    with h5py.File(str(h5_path), "r") as hf:
        attrs = {k: v for k, v in hf.attrs.items()}
        num_frames = int(attrs["num_frames"])

        mano_betas = hf["mano"]["betas"][()]
        mano_global_orient = hf["mano"]["global_orient"][()]
        mano_hand_pose = hf["mano"]["hand_pose"][()]
        right_flags = hf["right"][()]

        if "img_idx" in hf:
            frame_idx = hf["img_idx"][()]
        elif "frame_idx" in hf:
            frame_idx = hf["frame_idx"][()]
        else:
            raise ValueError("WiLoR H5 missing 'img_idx' or 'frame_idx'")

    # Build per-frame arrays (same logic as original h5_load_mano)
    data_mano = {
        "betas_r": np.zeros((num_frames, 10), dtype=np.float32),
        "betas_l": np.zeros((num_frames, 10), dtype=np.float32),
        "global_orient_r": np.tile(np.eye(3, dtype=np.float32), (num_frames, 1, 1)),
        "global_orient_l": np.tile(np.eye(3, dtype=np.float32), (num_frames, 1, 1)),
        "right_hand_pose": np.tile(
            np.eye(3, dtype=np.float32)[None, None], (num_frames, 15, 1, 1)
        ),
        "left_hand_pose": np.tile(
            np.eye(3, dtype=np.float32)[None, None], (num_frames, 15, 1, 1)
        ),
        "success_r": np.zeros(num_frames, dtype=bool),
        "success_l": np.zeros(num_frames, dtype=bool),
    }

    # Handle both formats:
    #   - Flat per-detection: img_idx shape (N_det,) — each value is the frame
    #   - CSR-style: img_idx shape (num_frames, 2) — each row is [start, count]
    if frame_idx.ndim == 1:
        # Flat per-detection format: group detections by frame
        for det_i in range(len(frame_idx)):
            fi = int(frame_idx[det_i])
            is_right = bool(right_flags[det_i])
            side = "r" if is_right else "l"
            data_mano[f"success_{side}"][fi] = True
            data_mano[f"betas_{side}"][fi] = mano_betas[det_i]
            data_mano[f"global_orient_{side}"][fi] = mano_global_orient[det_i]
            pose_key = "right_hand_pose" if is_right else "left_hand_pose"
            data_mano[pose_key][fi] = mano_hand_pose[det_i]
    else:
        assert len(frame_idx) == num_frames
        for fi, (start, count) in enumerate(frame_idx):
            if count == 0:
                continue

            if count == 1:
                is_right = right_flags[start]
                side = "r" if is_right else "l"
                data_mano[f"success_{side}"][fi] = True
                data_mano[f"betas_{side}"][fi] = mano_betas[start]
                data_mano[f"global_orient_{side}"][fi] = mano_global_orient[start]
                pose_key = "right_hand_pose" if is_right else "left_hand_pose"
                data_mano[pose_key][fi] = mano_hand_pose[start]

            else:
                assert count == 2, f"Expected count 1 or 2, got {count}"
                data_mano["success_r"][fi] = True
                data_mano["success_l"][fi] = True

                is_right_pair = right_flags[start : start + 2]
                if is_right_pair[0]:
                    assert not is_right_pair[1]
                    ri, li = start, start + 1
                else:
                    assert is_right_pair[1]
                    ri, li = start + 1, start

                data_mano["betas_r"][fi] = mano_betas[ri]
                data_mano["betas_l"][fi] = mano_betas[li]
                data_mano["global_orient_r"][fi] = mano_global_orient[ri]
                data_mano["global_orient_l"][fi] = mano_global_orient[li]
                data_mano["right_hand_pose"][fi] = mano_hand_pose[ri]
                data_mano["left_hand_pose"][fi] = mano_hand_pose[li]

    # Flip left-hand rotations (mirror symmetry)
    for key in ("global_orient_l", "left_hand_pose"):
        data_mano[key] = flipPoseRot(torch.from_numpy(data_mano[key]).float())

    # Convert remaining arrays to tensors
    for key in list(data_mano.keys()):
        if key.startswith("success"):
            continue
        v = data_mano[key]
        if not torch.is_tensor(v):
            data_mano[key] = torch.from_numpy(v).float()

    class _Container:
        pass

    result = _Container()
    result.data = {
        "mano": data_mano,
        "_attrs": attrs,
        "frame_idx": frame_idx,
    }
    result.meta = attrs
    result.num_frames = num_frames
    return result


# ---------------------------------------------------------------------------
# RTMPose (wholebody_rtmpose.h5) – filtered mode
# ---------------------------------------------------------------------------

def load_rtm_for_fitting(
    h5_path: str | Path,
    zero_weights_out_frame: bool = True,
) -> dict:
    """Load RTMPose H5 in the format expected by the optimizer.

    VLTK RTMPose in filtered mode stores ``keypoints`` of shape
    ``(T, 1, 133, 3)`` where the last dimension is ``(x, y, confidence)``.
    The optimizer expects a ``people`` dict with ``person_0`` sub-dict.

    Returns:
        Dict compatible with ``readRTMDetections`` output.
    """
    with h5py.File(str(h5_path), "r") as hf:
        attrs = {k: v for k, v in hf.attrs.items()}
        resolution = attrs["resolution"]  # (H, W)
        is_filtered = bool(attrs.get("filter", True))
        video_path = str(attrs.get("video_path", ""))

        if is_filtered:
            kp = hf["keypoints"][()]  # (T, 1, 133, 3)
            # Squeeze person dimension → (T, 133, 3)
            if kp.ndim == 4 and kp.shape[1] == 1:
                kp = kp[:, 0]

            bboxes = hf["bboxes"][()] if "bboxes" in hf else None

            num_frames = kp.shape[0]
            output: dict = {
                "filter": False,
                "resolution": np.asarray(resolution),
                "video_path": video_path,
                "people": {
                    "person_0": {
                        "video_path": video_path,
                        "filter": False,
                        "keypoints": kp,
                        "bboxes": bboxes,
                        "successful_frames": np.arange(num_frames, dtype=np.int32),
                    },
                },
            }
        else:
            output = {
                "filter": False,
                "resolution": np.asarray(resolution),
                "video_path": video_path,
                "people": {},
            }
            for person_id in hf.keys():
                grp = hf[person_id]
                if not isinstance(grp, h5py.Group):
                    continue
                output["people"][person_id] = {
                    "video_path": video_path,
                    "filter": False,
                    "keypoints": grp["keypoints"][()],
                    "bboxes": grp["bboxes"][()],
                    "successful_frames": grp["successful_frames"][()],
                }

    if zero_weights_out_frame:
        h, w = int(resolution[0]), int(resolution[1])
        for pid in output["people"]:
            kp = output["people"][pid]["keypoints"]
            mask_oob = (
                (kp[..., 0] < 0) | (kp[..., 0] >= w)
                | (kp[..., 1] < 0) | (kp[..., 1] >= h)
            )
            kp[mask_oob, 2] = 0.0

    return output


# ---------------------------------------------------------------------------
# Depth (depth.h5)
# ---------------------------------------------------------------------------

def load_depth_for_fitting(
    h5_path: str | Path,
) -> dict:
    """Load depth H5 written by the depth extractor.

    Returns:
        Dict keyed by person_id, each with ``depths`` and ``depths_mask`` tensors.
    """
    output: dict = {}
    with h5py.File(str(h5_path), "r") as hf:
        output["attrs"] = dict(hf.attrs)
        for person_id in hf.keys():
            grp = hf[person_id]
            if not isinstance(grp, h5py.Group):
                continue
            output[person_id] = {
                "depths": torch.from_numpy(grp["depths"][()]).float(),
                "depths_mask": torch.from_numpy(grp["depths_mask"][()]).float(),
            }
    return output
