"""Data loading from TEASER HDF5 and SMPL NLF files."""

from pathlib import Path

import h5py

try:
    import hdf5plugin  # noqa: F401 -- required for LZ4 decompression
except ImportError:
    pass

import numpy as np

from ..constants import EXPR_BLINK_INDICES, FPS
from ..models import TrackingQuality


def load_teaser_data(
    h5_path: str, person: str = "0", load_expression: bool = False, load_nms: bool = False
):
    """Load FLAME parameters and metadata from TEASER h5.

    Returns eyelid signal (from the active eye column), raw arrays, and
    a TrackingQuality assessment.
    """
    with h5py.File(h5_path, "r") as f:
        if person not in f:
            available = list(f.keys())
            raise KeyError(f"Person '{person}' not in {h5_path}. Available: {available}")

        grp = f[person]
        eyelid_both = grp["flame/eyelid_params"][:].astype(np.float64)
        successful_frames = grp["successful_frames"][:].astype(np.int64)

        boxes = grp["boxes"][:].astype(np.float32) if "boxes" in grp else None
        pose = grp["pose_params"][:].astype(np.float32) if "pose_params" in grp else None
        scores = grp["boxes_scores"][:].astype(np.float64) if "boxes_scores" in grp else None

        jaw_pose = None
        if load_nms and "flame/jaw_pose" in grp:
            jaw_pose = grp["flame/jaw_pose"][:].astype(np.float64)

        expression = None
        if (load_expression or load_nms) and "flame/expression" in grp:
            expr_full = grp["flame/expression"][:].astype(np.float64)
            if load_nms:
                expression = expr_full
            else:
                valid_indices = [i for i in EXPR_BLINK_INDICES if i < expr_full.shape[1]]
                expression = expr_full[:, valid_indices] if valid_indices else None

    n = len(eyelid_both)

    has_left = eyelid_both[:, 0] > 0
    has_right = (eyelid_both[:, 1] > 0) if eyelid_both.shape[1] > 1 else np.zeros(n, dtype=bool)
    both_eyes = has_left & has_right

    eyelid = np.where(
        both_eyes,
        np.maximum(eyelid_both[:, 0], eyelid_both[:, 1]),
        np.where(has_left, eyelid_both[:, 0], np.where(has_right, eyelid_both[:, 1], 0.0)),
    )

    nz_left = int(has_left.sum())
    nz_right = int(has_right.sum())
    if nz_left > 0 and nz_right > 0:
        active_eye = 2
    elif nz_left > 0:
        active_eye = 0
    elif nz_right > 0:
        active_eye = 1
    else:
        active_eye = -1

    video_span = int(successful_frames.max() - successful_frames.min() + 1) if n > 0 else 0
    eyelid_valid = int(np.count_nonzero(eyelid > 0))

    quality = TrackingQuality(
        total_data_frames=n,
        total_video_frames=video_span,
        detection_rate=n / video_span if video_span > 0 else 0.0,
        eyelid_valid_frames=eyelid_valid,
        eyelid_valid_rate=eyelid_valid / n if n > 0 else 0.0,
        score_mean=float(np.mean(scores)) if scores is not None else 0.0,
        active_eye=active_eye,
        usable_time_sec=eyelid_valid / FPS,
    )

    return (eyelid, successful_frames, boxes, pose, scores, expression, jaw_pose, quality)


def resolve_smpl_path(h5_path: str, smpl_root: Path | None = None) -> Path | None:
    """Derive the SMPL NLF file path from a TEASER HDF5 path.

    Looks for '{stem}_nlf.h5' in the same directory or in smpl_root.
    """
    basename = Path(h5_path).stem
    # First check same directory
    same_dir = Path(h5_path).parent / f"{basename}_nlf.h5"
    if same_dir.exists():
        return same_dir
    # Then check smpl_root if provided
    if smpl_root is not None:
        smpl_path = smpl_root / f"{basename}_nlf.h5"
        if smpl_path.exists():
            return smpl_path
    return None


def load_smpl_eye_data(
    smpl_path: Path,
    teaser_successful_frames: np.ndarray,
    teaser_boxes: np.ndarray | None = None,
    verbose: bool = False,
):
    """Load eye gaze data from an SMPL NLF file, aligned to TEASER frames."""
    n_teaser = len(teaser_successful_frames)
    leye_out = np.full((n_teaser, 3), np.nan, dtype=np.float64)
    reye_out = np.full((n_teaser, 3), np.nan, dtype=np.float64)

    try:
        with h5py.File(str(smpl_path), "r") as f:
            frame_idx = f["frame_idx"][:]
            smpl_boxes = f["boxes"][:]
            leye_raw = f["smplx/leye_pose"][:]
            reye_raw = f["smplx/reye_pose"][:]
    except (KeyError, OSError) as e:
        if verbose:
            print(f"  Warning: Cannot load SMPL eye data from {smpl_path}: {e}")
        return None, None

    if teaser_boxes is not None and len(teaser_boxes) > 0:
        teaser_cx = (teaser_boxes[:, 0] + teaser_boxes[:, 2]) / 2.0
        target_cx = float(np.nanmedian(teaser_cx))
    else:
        target_cx = None

    n_smpl_frames = len(frame_idx)
    matched = 0

    for i in range(n_teaser):
        video_frame = int(teaser_successful_frames[i])

        if video_frame >= n_smpl_frames:
            continue

        det_start, n_det = int(frame_idx[video_frame, 0]), int(frame_idx[video_frame, 1])
        if n_det == 0:
            continue

        if n_det == 1 or target_cx is None:
            best = det_start
        else:
            best = det_start
            best_dist = float("inf")
            for d in range(n_det):
                idx = det_start + d
                cx = (smpl_boxes[idx, 0] + smpl_boxes[idx, 2]) / 2.0
                dist = abs(cx - target_cx)
                if dist < best_dist:
                    best_dist = dist
                    best = idx

        leye_out[i] = leye_raw[best]
        reye_out[i] = reye_raw[best]
        matched += 1

    if verbose and n_teaser > 0:
        print(
            f"  SMPL eye data: {matched}/{n_teaser} frames matched "
            f"({100 * matched / n_teaser:.0f}%)"
        )

    return leye_out, reye_out
