import h5py
import hdf5plugin  # This library is needed to read compressed HDF5 files
import os
import torch
import numpy as np
import pytorch3d

from .maths_utils import axis_to_euler_pytorch3d, axis_to_6d_pytorch3d
from dataclasses import dataclass
from typing import Dict, Any, Union
from pathlib import Path
from .utils_smplx import flipPoseRot

def readH5FileFLAME(h5_file, median_shape=True, to_torch=True, jaw_pose='axis', conv='ZYX'):
    assert os.path.exists(h5_file), f"HDF5 file not found: {h5_file}"

    required_keys = ['cam', 'pose_params', 'successful_frames']
    out = {}
    with h5py.File(h5_file, 'r') as hf:
        out['meta'] = {key: hf.attrs[key] for key in hf.attrs.keys()}

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

        if 'kpts_2d' in hf:
            out['kpts_2d'] = 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] = hf['flame'][name][()]
            if to_torch:
                out['flame'][name] = torch.FloatTensor(out['flame'][name])
    
    if to_torch and 'kpts_2d' in out:
        out['kpts_2d'] = torch.FloatTensor(out['kpts_2d']).type(torch.float32)

    if median_shape:
        if to_torch:
            out['flame']['shape'] = out['flame']['shape'].median(0).values[None]
        else:
            out['flame']['shape'] = np.median(out['flame']['shape'], axis=0)[None]
    assert len(out['successful_frames']) == out['pose_params'].shape[0], "Mismatch in successful_frames and pose_params length"
    if jaw_pose == 'axis':
        pass # already in axis-angle format
    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


@dataclass
class H5Features:
    """Container for features loaded from an H5 file."""
    model_type: str
    file_path: Path
    data: Dict[str, np.ndarray]
    meta: Dict[str, Any]
    num_frames: int
    available_keys: list

    def __repr__(self):
        return (f"H5Features(model_type={self.model_type}, file_path={self.file_path}, "
                f"num_frames={self.num_frames}, available_keys={self.available_keys})")


def _clean(data, skip_keys={'_attrs'}):
    """Recursively clean numpy arrays to consistent dtypes."""
    if isinstance(data, np.ndarray):
        if data.dtype == np.bool_:
            return data
        if np.issubdtype(data.dtype, np.floating):
            data = data.astype(np.float32)
        elif np.issubdtype(data.dtype, np.integer):
            data = data.astype(np.int32)
        return data
    elif isinstance(data, dict):
        return {k: (data[k] if k in skip_keys else _clean(v, skip_keys))
                for k, v in data.items()}
    else:
        return data
    
def _read_h5_group(group: Union[h5py.Group, h5py.File], out_dict: Dict) -> Dict:
    if group.attrs:
        out_dict["_attrs"] = {k: v for k, v in group.attrs.items()}
    for key in group.keys():
        item = group[key]
        if isinstance(item, h5py.Dataset):
            out_dict[key] = item[()]
        elif isinstance(item, h5py.Group):
            out_dict[key] = _read_h5_group(item, {})
    return out_dict


def h5_load_smpl(path: str, angle_format='axis', conv='ZYX', to_torch=True) -> H5Features:
    with h5py.File(path, "r") as hf:
        data = _read_h5_group(hf, {})
    data = _clean(data)
    if "smplx" in data:
        angular_keys = ["body_pose", "global_orient", "left_hand_pose", "right_hand_pose", "jaw_pose", "leye_pose", "reye_pose"]
        for angular_key in angular_keys:
            if angular_key not in data['smplx']:
                continue
            if angle_format == 'axis':
                data['smplx'][angular_key] = torch.from_numpy(data['smplx'][angular_key]).type(torch.float32).reshape(len(data['smplx'][angular_key]), -1)
            elif angle_format == '6d':
                data['smplx'][angular_key] = axis_to_6d_pytorch3d(torch.from_numpy(data['smplx'][angular_key]).float()).reshape(len(data['smplx'][angular_key]), -1)
            elif angle_format == 'euler':
                data['smplx'][angular_key] = axis_to_euler_pytorch3d(torch.from_numpy(data['smplx'][angular_key]).float(), conv).reshape(len(data['smplx'][angular_key]), -1)
            elif angle_format == 'mat':
                data['smplx'][angular_key] = pytorch3d.transforms.axis_angle_to_matrix(torch.from_numpy(data['smplx'][angular_key]).float().reshape(len(data['smplx'][angular_key]), -1, 3)).flatten(1, -1)
            else:
                raise NotImplementedError(f"Unsupported angle_format: {angle_format}")

    for key in data['smplx'].keys():
        if to_torch and (not torch.is_tensor(data['smplx'][key])):
            data['smplx'][key] = torch.from_numpy(data['smplx'][key]).type(torch.float32)
        elif (not to_torch) and torch.is_tensor(data['smplx'][key]):
            data['smplx'][key] = data['smplx'][key].detach().cpu().numpy()

    data['smplx']['transl'] = data['smplx'].pop('trans', None)
    return H5Features(
        model_type="smpl",
        file_path=str(path),
        data=data,
        meta=data.get("_attrs", {}),
        num_frames=data["frame_idx"].shape[0] if "frame_idx" in data else 0,
        available_keys=list(data.keys())
    )


def h5_load_mano(path: str, angle_format='axis', conv='ZYX', to_torch=True) -> H5Features:
    with h5py.File(path, "r") as hf:
        data = _read_h5_group(hf, {})
    data = _clean(data)
    if "mano" in data:
        num_frames = data['_attrs']['num_frames']
        frame_idx = data['frame_idx']  # [N, 2] Gives start detection and number of frames
        data_mano = {
            "betas_r": np.zeros((num_frames, 10)),
            "betas_l": np.zeros((num_frames, 10)),
            "global_orient_r": np.eye(3)[None].repeat(num_frames, 0),
            "global_orient_l": np.eye(3)[None].repeat(num_frames, 0),
            "right_hand_pose": np.eye(3)[None,None].repeat(num_frames, 0).repeat(15, 1),
            "left_hand_pose": np.eye(3)[None,None].repeat(num_frames, 0).repeat(15, 1),
            'success_r': np.zeros((num_frames,), dtype=bool),
            'success_l': np.zeros((num_frames,), dtype=bool),
        }
        assert len(frame_idx) == num_frames, "frame_idx length should match num_frames"
        for frame_idx, (start, count) in enumerate(frame_idx):
            if count == 0:
                continue # no detections
            if count == 1:
                is_right = data["right"][start]
                data_mano['success_r' if is_right else 'success_l'][frame_idx] = True
                data_mano['betas_r' if is_right else 'betas_l'][frame_idx] = data['mano']["betas"][start]
                data_mano['global_orient_r' if is_right else 'global_orient_l'][frame_idx] = data['mano']["global_orient"][start]
                data_mano['right_hand_pose' if is_right else 'left_hand_pose'][frame_idx] = data['mano']["hand_pose"][start]
            else:
                assert count == 2, f"Expected count to be 1 or 2, but got {count}"
                data_mano['success_r'][frame_idx] = data_mano['success_l'][frame_idx] = True
                is_right = data["right"][start:start+2]
                if is_right[0]:
                    assert not is_right[1], "Expected only one detection to be right hand"
                    right_idx, left_idx = start, start+1
                else:
                    assert is_right[1], "Expected only one detection to be right hand"
                    right_idx, left_idx = start+1, start

                data_mano['betas_r'][frame_idx] = data['mano']["betas"][right_idx]
                data_mano['betas_l'][frame_idx] = data['mano']["betas"][left_idx]

                data_mano['global_orient_r'][frame_idx] = data['mano']["global_orient"][right_idx]
                data_mano['global_orient_l'][frame_idx] = data['mano']["global_orient"][left_idx]

                data_mano['right_hand_pose'][frame_idx] = data['mano']["hand_pose"][right_idx]
                data_mano['left_hand_pose'][frame_idx] = data['mano']["hand_pose"][left_idx]

        data['mano'] = data_mano
        for key in ['global_orient_l', 'left_hand_pose']:
            data['mano'][key] = flipPoseRot(torch.from_numpy(data['mano'][key]).type(torch.float32))

        assert angle_format == 'mat', "Only 'mat' angle_format is supported."
 
    for key in data['mano'].keys():
        if key != 'success_r' and key != 'success_l': # these are boolean masks, keep as is
            if to_torch and (not torch.is_tensor(data['mano'][key])):
                data['mano'][key] = torch.from_numpy(data['mano'][key]).type(torch.float32)
            elif (not to_torch) and torch.is_tensor(data['mano'][key]):
                data['mano'][key] = data['mano'][key].detach().cpu().numpy()

    return H5Features(
        model_type="mano",
        file_path=str(path),
        data=data,
        meta=data.get("_attrs", {}),
        num_frames=data["frame_idx"].shape[0] if "frame_idx" in data else 0,
        available_keys=list(data.keys())
    )

def readDepthH5(filename, to_torch=True):
    with h5py.File(filename, 'r') as hf:
        output = {}
        for person_id in hf.keys():
            output[person_id] = {}
            output[person_id]['depths'] = hf[person_id]['depths'][:]
            output[person_id]['depths_mask'] = hf[person_id]['depths_mask'][:]
            if to_torch:
                output[person_id]['depths'] = torch.from_numpy(output[person_id]['depths']).type(torch.float32)
                output[person_id]['depths_mask'] = torch.from_numpy(output[person_id]['depths_mask']).type(torch.float32)
    return output

def readRTMDetections(filename, zero_weights_out_frame=True):
    with h5py.File(filename, 'r') as hf:
        resolution = hf.attrs['resolution']

        if hf.attrs['filter']:
            all_kp = hf['keypoints'][:]
            all_bboxes = hf['bboxes'][:]

            output = {'filter': True, 'keypoints': all_kp, 'bboxes': all_bboxes, 'video_path': hf.attrs['video_path']}
            assert len(all_kp) == len(all_bboxes)

        else:
            output = {'filter': False, 'video_path': hf.attrs['video_path'], 'people': {}}
            for person_id in hf.keys():
                person_group = hf[person_id]
                person_kp = person_group['keypoints'][:]
                person_bboxes = person_group['bboxes'][:]
                successful_frames = person_group['successful_frames'][:]

                
                output['people'][person_id] = {
                    'video_path': hf.attrs['video_path'],
                    'filter': False,
                    'keypoints': person_kp,
                    'bboxes': person_bboxes,
                    'successful_frames': successful_frames
                }
    if zero_weights_out_frame:
        for person_id in output['people']:
            mask_out_of_frame = (output['people'][person_id]['keypoints'][...,0] < 0) | (output['people'][person_id]['keypoints'][...,0] >= resolution[1]) | (output['people'][person_id]['keypoints'][...,1] < 0) | (output['people'][person_id]['keypoints'][...,1] >= resolution[0])
            output['people'][person_id]['keypoints'][mask_out_of_frame,2] = 0.0

    return output
