import argparse
import numpy as np
import os
import torch
import h5py

from video_depth_anything.video_depth import VideoDepthAnything
from utils.dc_utils import read_video_frames, save_video
from tqdm import tqdm
from pathlib import Path


def _to_numpy_batch(batch):
    """Convert a decord/torch batch object to a numpy array."""
    if hasattr(batch, "asnumpy"):
        return batch.asnumpy()
    if hasattr(batch, "detach"):
        return batch.detach().cpu().numpy()
    if hasattr(batch, "numpy"):
        return batch.numpy()
    return np.asarray(batch)


def read_kp_from_h5(kp_path):
    with h5py.File(kp_path, 'r') as hf:
        print('Attributes:')
        print(f"Video path: {hf.attrs['video_path']}")
        print(f"FPS: {hf.attrs['fps']}")
        print(f"Number of frames: {hf.attrs['num_frames']}")
        print(f"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
                }

    return output

def sample2DdepthFromDepthMap(depths, points_2d):
    """
    depths: (B, H, W) np.ndarray
    points_2d: (B, N, 2) - Can be Tensor or Numpy
    """
    is_tensor = torch.is_tensor(points_2d)
    if is_tensor:
        points_np = points_2d.detach().cpu().numpy()
    else:
        points_np = points_2d

    B, H, W = depths.shape
    assert points_np.shape[0] == B, "Batch size of points_2d must match batch size of depths"
    assert points_2d.ndim == 3 and points_2d.shape[2] == 2, "points_2d must have shape (B, N, 2)"
    
    # 2. Check Bounds (Validity Mask)
    x = points_np[..., 0]
    y = points_np[..., 1]
    
    valid_mask = ((x >= 0) & (x < W - 0.5) & (y >= 0) & (y < H - 0.5)).astype(np.float32)
    
    # 3. Safe Sampling with Rounding (Nearest Neighbor)
    x_clamped = np.clip(np.rint(x), 0, W - 1).astype(int)
    y_clamped = np.clip(np.rint(y), 0, H - 1).astype(int)
    
    sampled_depths = depths[np.arange(B)[:, None], y_clamped, x_clamped]
    if is_tensor:
        return torch.from_numpy(sampled_depths), torch.from_numpy(valid_mask)
    
    return sampled_depths, valid_mask

from decord import VideoReader, cpu

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Video Depth Anything')
    parser.add_argument('--input_video', type=str, default='./assets/example_videos/davis_rollercoaster.mp4')
    parser.add_argument('--output_path', type=str, required=True)
    parser.add_argument('--input_size', type=int, default=518)
    parser.add_argument('--max_res', type=int, default=10000)
    parser.add_argument('--encoder', type=str, default='vitl', choices=['vits', 'vitb', 'vitl'])
    parser.add_argument('--metric', action='store_true', help='use metric model')
    parser.add_argument('--fp32', action='store_true', help='model infer with torch.float32, default is torch.float16')
    parser.add_argument('--grayscale', action='store_true', help='do not apply colorful palette')
    parser.add_argument('--save_npz', action='store_true', help='save depths as npz')
    parser.add_argument('--save_exr', action='store_true', help='save depths as exr')
    parser.add_argument('--focal-length-x', default=470.4, type=float,
                        help='Focal length along the x-axis.')
    parser.add_argument('--focal-length-y', default=470.4, type=float,
                        help='Focal length along the y-axis.')

    parser.add_argument('--h5_keypoints', type=str, required=True, help='Path to h5 file containing keypoints')


    args = parser.parse_args()

    # if os.path.exists(args.output_path):
    #     print('Output already exists!', args.output_path)
    #     exit(0)

    assert not args.metric, "not needed"
    output_kps = read_kp_from_h5(args.h5_keypoints)
    assert not output_kps['filter']

    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

    model_configs = {
        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
    }
    checkpoint_name = 'metric_video_depth_anything' if args.metric else 'video_depth_anything'

    video_depth_anything = VideoDepthAnything(**model_configs[args.encoder], metric=args.metric)
    video_depth_anything.load_state_dict(torch.load(f'./checkpoints/{checkpoint_name}_{args.encoder}.pth', map_location='cpu'), strict=True)
    video_depth_anything = video_depth_anything.to(DEVICE).eval()


    vid = VideoReader(args.input_video, ctx=cpu(0))
    original_height, original_width = vid.get_batch([0]).shape[1:3]
    height = original_height
    width = original_width
    if args.max_res > 0 and max(height, width) > args.max_res:
        assert False, "keep the original resolution to match keypoints!"
        scale = max_res / max(original_height, original_width)
        height = ensure_even(round(original_height * scale))
        width = ensure_even(round(original_width * scale))

    vid = VideoReader(args.input_video, ctx=cpu(0), width=width, height=height)

    fps = vid.get_avg_fps()
    stride = round(vid.get_avg_fps() / fps)
    stride = max(stride, 1)

    frames_idx = list(range(0, len(vid), stride))
    target_fps = fps

    frame_step = 2*22 # from VideoDepthAnything

    TEST = False
    depths_all = []
    depths_per_person = {key: [] for key in output_kps['people'].keys()}
    depths_mask_per_person = {key: [] for key in output_kps['people'].keys()}
    
    person_inv_frame_map = {}
    for person_id in output_kps['people'].keys():
        person_inv_frame_map[person_id] = {frame_idx: order_idx for order_idx, frame_idx in enumerate(output_kps['people'][person_id]['successful_frames'])}

    for i in tqdm(range(0, len(frames_idx), frame_step)):
        frames = _to_numpy_batch(vid.get_batch(frames_idx[i:i+frame_step]))
        depths = video_depth_anything.infer_video_depth(frames, target_fps, input_size=args.input_size, device=DEVICE, fp32=args.fp32)[0]
        if TEST:
            depths_all.append(depths) # for test
        for person_id in output_kps['people'].keys():
            kp_idxs, depth_idxs = [], []
            for ii, frame_idx in enumerate(frames_idx[i:i+frame_step]):
                if frame_idx in person_inv_frame_map[person_id]:
                    kp_idxs.append(person_inv_frame_map[person_id][frame_idx])
                    depth_idxs.append(ii)

            if len(kp_idxs) > 0:
                sampled_depths, sampled_depths_mask = sample2DdepthFromDepthMap(depths[depth_idxs], output_kps['people'][person_id]['keypoints'][kp_idxs,:,:2])
                depths_per_person[person_id].append(sampled_depths)
                depths_mask_per_person[person_id].append(sampled_depths_mask)

    if TEST:
        save_video(np.concatenate(depths_all), args.output_path.replace('.h5', '.mp4'), fps=fps, is_depths=True, grayscale=args.grayscale)

    for key in depths_per_person.keys():
        depths_per_person[key] = np.concatenate(depths_per_person[key])
        depths_mask_per_person[key] = np.concatenate(depths_mask_per_person[key])


    if '/' in args.output_path:
        Path(args.output_path).parent.mkdir(parents=True, exist_ok=True)

    assert args.output_path.endswith('.h5')
    with h5py.File(args.output_path, 'w') as hf:
        hf.attrs['video_path'] = args.input_video
        hf.attrs['h5_keypoints'] = args.h5_keypoints
        hf.attrs['encoder'] = args.encoder
        
        for person_id in depths_per_person.keys():
            person_group = hf.create_group(f'{person_id}')
            person_group.create_dataset('depths',
                                        data=np.array(depths_per_person[person_id], dtype=np.float32),
                                        compression="gzip",
                                        compression_opts=4)

            person_group.create_dataset('depths_mask',
                                        data=np.array(depths_mask_per_person[person_id], dtype=np.float32),
                                        compression="gzip",
                                        compression_opts=4)

            person_group.create_dataset('successful_frames',
                                        data=np.array(output_kps['people'][person_id]['successful_frames'], dtype=np.int32),
                                        compression="gzip",
                                        compression_opts=4)
