import torch
import numpy as np

try:
    import decord
    from decord import cpu
    USE_DECORD = True
except:
    import av
    USE_DECORD = False

def read_video(video_path):
    """Read video frames into list of HxWxC uint8 arrays for VideoMAE."""
    if USE_DECORD:
        vr = decord.VideoReader(video_path, ctx=cpu(0))
        frames = vr.get_batch(range(len(vr)))  # (T, H, W, C)
        frames = frames.asnumpy().astype(np.uint8)  # ensure uint8 for VideoMAE
        frames_list = [frame for frame in frames]  # list of T frames
    else:
        container = av.open(video_path)
        frames = []
        for frame in container.decode(video=0):
            frames.append(frame.to_ndarray(format="rgb24").astype(np.uint8))
        return frames
    return frames_list


def resample_video_idx(num_frames, original_fps, new_fps):
    """Return frame indices to match new fps"""
    step = float(original_fps) / new_fps
    if step.is_integer():
        step = int(step)
        return slice(None, None, step)
    idxs = torch.arange(num_frames, dtype=torch.float32) * step
    idxs = idxs.floor().to(torch.int64)
    return idxs

def process_frames(frames, target_num_frames, input_fps, target_fps=None):
    """
    frames: list of np arrays (H, W, C)
    target_num_frames: int
    Returns: list of np arrays (H, W, C) ready for processor
    """
    target_fps = target_fps or input_fps
    num_frames = len(frames)
    duration = num_frames / input_fps

    # Too short → resample with new fps
    if num_frames < target_num_frames:
        new_fps = np.ceil(target_num_frames / duration)
        idxs = resample_video_idx(target_num_frames, input_fps, new_fps)
        idxs = np.clip(idxs, 0, num_frames - 1)
        frames = [frames[i] for i in idxs]

    # Too long → uniform sampling
    elif num_frames > target_num_frames:
        idxs = np.linspace(0, num_frames - 1, target_num_frames).astype(int)
        frames = [frames[i] for i in idxs]

    # Pad if still short
    if len(frames) < target_num_frames:
        pad = target_num_frames - len(frames)
        frames.extend([frames[-1]] * pad)

    return frames




def build_transform(config, mode="train"):
    import torchvision.transforms as T
    import torchvision.transforms.functional as F
    import random
    from torchvision.transforms import Compose, RandomCrop, RandomHorizontalFlip, ColorJitter, RandomErasing, Resize, Lambda

    frame_height, frame_width = config.DATA.frame_size
    augmentations = config.DATA.augmentations

    def transform_fn(frame: np.ndarray):
        """
        frame: np.ndarray, HxWxC, dtype=uint8
        returns: np.ndarray, HxWxC, dtype=float32, normalized [0,1]
        """
        frame = frame.astype(np.float32) / 255.0  # normalize to [0,1]

        # Training augmentations
        if mode == "train":
            if getattr(augmentations, "random_crop", False):
                # Random crop
                i, j, h, w = T.RandomCrop.get_params(
                    torch.from_numpy(frame).permute(2,0,1), 
                    output_size=(frame_height, frame_width)
                )
                frame = torch.from_numpy(frame).permute(2,0,1)  # C,H,W
                frame = F.crop(frame, i, j, h, w)
                frame = frame.permute(1,2,0).numpy()  # back to H,W,C
            else:
                # Resize
                frame = F.resize(torch.from_numpy(frame).permute(2,0,1), [frame_height, frame_width]).permute(1,2,0).numpy()

            if getattr(augmentations, "random_horizontal_flip", False):
                if random.random() < getattr(augmentations, "flip_prob", 0.5):
                    frame = np.ascontiguousarray(frame[:, ::-1, :])  # horizontal flip

            if getattr(augmentations, "color_jitter", False):
                brightness, contrast, saturation, hue = getattr(augmentations, "jitter_params", [0.4,0.4,0.4,0.1])
                cj = T.ColorJitter(brightness=brightness, contrast=contrast, saturation=saturation, hue=hue)
                frame = cj(torch.from_numpy(frame).permute(2,0,1)).permute(1,2,0).numpy()

            if getattr(augmentations, "random_erasing", False):
                re = T.RandomErasing()
                frame = re(torch.from_numpy(frame).permute(2,0,1)).permute(1,2,0).numpy()

        else:  # validation / test
            frame = F.resize(torch.from_numpy(frame).permute(2,0,1), [frame_height, frame_width]).permute(1,2,0).numpy()

        return frame.astype(np.float32)

    return transform_fn

# import torch
# import numpy as np
# import decord
# from decord import cpu
# import torchvision.transforms as T

# def read_video(video_path):
#     """Read video frames into tensor (T, C, H, W)"""
#     vr = decord.VideoReader(video_path, ctx=cpu(0))
#     frames = vr.get_batch(range(len(vr)))  # (T, H, W, C)
#     frames = torch.from_numpy(frames.asnumpy()).permute(0, 3, 1, 2)  # (T, C, H, W)
#     return frames

# def resample_video_idx(num_frames, original_fps, new_fps):
#     """Return frame indices to match new fps"""
#     step = float(original_fps) / new_fps
#     if step.is_integer():
#         step = int(step)
#         return slice(None, None, step)
#     idxs = torch.arange(num_frames, dtype=torch.float32) * step
#     idxs = idxs.floor().to(torch.int64)
#     return idxs

# def process_frames(video_tensor, target_num_frames, input_fps, target_fps=None):
#     """Uniform sampling, padding, FPS adjustment"""
#     target_fps = target_fps or input_fps
#     num_frames = video_tensor.size(0)
#     duration = num_frames / input_fps

#     # Too short → resample
#     if num_frames < target_num_frames:
#         new_fps = np.ceil(target_num_frames / duration)
#         idxs = resample_video_idx(target_num_frames, input_fps, new_fps)
#         idxs = np.clip(idxs, 0, num_frames - 1)
#         video_tensor = video_tensor[idxs]

#     # Too long → uniform sample
#     elif num_frames > target_num_frames:
#         idxs = np.linspace(0, num_frames - 1, target_num_frames)
#         video_tensor = video_tensor[idxs.astype(int)]

#     # Pad if still short
#     if video_tensor.size(0) < target_num_frames:
#         pad = target_num_frames - video_tensor.size(0)
#         last_frame = video_tensor[-1:].repeat(pad, 1, 1, 1)
#         video_tensor = torch.cat([video_tensor, last_frame], dim=0)

#     return video_tensor

# def build_transform(frame_size=(224, 224), mean=None, std=None):
#     """Return default frame transform"""
#     mean = mean or [0.485, 0.456, 0.406]
#     std = std or [0.229, 0.224, 0.225]
#     return T.Compose([
#         T.ConvertImageDtype(torch.float32),
#         T.Resize(frame_size),
#         T.Normalize(mean, std),
#     ])

