import numpy as np
import pandas as pd
import cv2
import torch
import warnings
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
import ffmpeg
import pytorchvideo
from pytorchvideo.transforms.functional import (
    uniform_temporal_subsample,
    short_side_scale_with_boxes,
    clip_boxes_to_image
)
from torchvision.transforms._functional_video import normalize
from pytorchvideo.data.ava import AvaLabeledVideoFramePaths
from pytorchvideo.models.hub import slowfast_r50_detection # Another option is slow_r50_detection
from visualization import VideoVisualizer


# This method takes in an image and generates the bounding boxes for people in the image.
def get_person_bboxes(inp_img, predictor):
    predictions = predictor(inp_img.cpu().detach().numpy())['instances'].to('cpu')
    boxes = predictions.pred_boxes if predictions.has("pred_boxes") else None
    scores = predictions.scores if predictions.has("scores") else None
    classes = np.array(predictions.pred_classes.tolist() if predictions.has("pred_classes") else None)
    predicted_boxes = boxes[np.logical_and(classes==0, scores>0.75 )].tensor.cpu() # only person
    return predicted_boxes


def ava_inference_transform(
    clip,
    boxes,
    num_frames = 32, # 4 if using slowfast_r50_detection, change this to 32
    crop_size = 256,
    data_mean = [0.45, 0.45, 0.45],
    data_std = [0.225, 0.225, 0.225],
    slow_fast_alpha = 4, # if using slowfast_r50_detection, change None to 4
    device = 'cpu'): 

    boxes = np.array(boxes)
    ori_boxes = boxes.copy()

    # Image [0, 255] -> [0, 1].
    clip = uniform_temporal_subsample(clip, num_frames)
    clip = clip.float()
    clip = clip / 255.0

    height, width = clip.shape[2], clip.shape[3]
    # The format of boxes is [x1, y1, x2, y2]. The input boxes are in the
    # range of [0, width] for x and [0,height] for y
    boxes = clip_boxes_to_image(boxes, height, width)

    # Resize short side to crop_size. Non-local and STRG uses 256.
    clip, boxes = short_side_scale_with_boxes(clip, size=crop_size, boxes=boxes)

    # Normalize images by mean and std.
    clip = normalize(clip, np.array(data_mean, dtype=np.float32), np.array(data_std, dtype=np.float32))

    boxes = clip_boxes_to_image(boxes, clip.shape[2],  clip.shape[3])

    # Incase of slowfast, generate both pathways
    if slow_fast_alpha is not None:
        fast_pathway = clip
        # Perform temporal sampling from the fast pathway.
        slow_pathway = torch.index_select(clip, 1, torch.linspace(
            0, clip.shape[1] - 1, clip.shape[1] // slow_fast_alpha).long())
        clip = [slow_pathway.unsqueeze(0).to(device), fast_pathway.unsqueeze(0).to(device)]

    return clip, torch.from_numpy(boxes), ori_boxes

# get video info
def with_opencv(filename):
    video = cv2.VideoCapture(filename)
    frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT)
    fps = video.get(cv2.CAP_PROP_FPS)
    s = round(frame_count / fps)
    video.release()
    return int(s), fps


def slow_fast_train(file_path, gpu=False):
    device = 'cuda' if gpu else 'cpu'
    top_k = 1

    video_model = slowfast_r50_detection(True) # Another option is slow_r50_detection(True) 
    video_model = video_model.eval().to(device)
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.55  # set threshold for this model
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
    cfg.MODEL.DEVICE = device
    predictor = DefaultPredictor(cfg)
    # Create an id to label name mapping
    label_map, allowed_class_ids = AvaLabeledVideoFramePaths.read_label_map('ava_action_list.pbtxt')
    # Create a video visualizer that can plot bounding boxes and visualize actions on bboxes.
    video_visualizer = VideoVisualizer(81, label_map, top_k=top_k, mode="thres",thres=0.5) #get top3 predictions show in each bounding box

    #preprocess video data
    encoded_vid = pytorchvideo.data.encoded_video.EncodedVideo.from_path(file_path)

    # Video predictions are generated each frame/second for the wholevideo.
    total_sec, fps = with_opencv(file_path)
    time_stamp_range = range(0, total_sec) # time stamps in video for which clip is sampled
    clip_duration = 1.0 # Duration of clip used for each inference step.
    gif_imgs = []
    xleft, ytop, xright, ybottom = [], [], [], []
    labels = []
    time_frame = []
    scores = []

    for time_stamp in time_stamp_range:

        # Generate clip around the designated time stamps
        inp_imgs = encoded_vid.get_clip(
            time_stamp - clip_duration/2.0,
            time_stamp + clip_duration/2.0)  
        inp_imgs = inp_imgs['video']

        #if time_stamp % 15 == 0:
            # Generate people bbox predictions using Detectron2's off the self pre-trained predictor
            # We use the the middle image in each clip to generate the bounding boxes.
        inp_img = inp_imgs[:,inp_imgs.shape[1]//2,:,:]
        inp_img = inp_img.permute(1,2,0)
    
        # Predicted boxes are of the form List[(x_1, y_1, x_2, y_2)]
        predicted_boxes = get_person_bboxes(inp_img, predictor)
        if len(predicted_boxes) == 0:
            print("Skipping clip no frames detected at time stamp: ", time_stamp)
            continue
        
        # Preprocess clip and bounding boxes for video action recognition.
        inputs, inp_boxes, _ = ava_inference_transform(inp_imgs, predicted_boxes.numpy(), device=device)
        # Prepend data sample id for each bounding box.
        # For more details refere to the RoIAlign in Detectron2
        inp_boxes = torch.cat([torch.zeros(inp_boxes.shape[0],1), inp_boxes], dim=1)
        
        # Generate actions predictions for the bounding boxes in the clip.
        # The model here takes in the pre-processed video clip and the detected bounding boxes.
        preds = video_model(inputs, inp_boxes.to(device)) #change inputs to inputs.unsqueeze(0).to(device) if using slow_r50
        
        preds = preds.to('cpu')
        # The model is trained on AVA and AVA labels are 1 indexed so, prepend 0 to convert to 0 index.
        preds = torch.cat([torch.zeros(preds.shape[0],1), preds], dim=1)

        # Plot predictions on the video and save for later visualization.
        inp_imgs = inp_imgs.permute(1,2,3,0)
        inp_imgs = inp_imgs/255.0
        out_img_pred = video_visualizer.draw_clip_range(inp_imgs, preds, predicted_boxes)
        gif_imgs += out_img_pred

        #format of bboxes(x_left, y_top, x_right, y_bottom)
        predicted_boxes_lst = predicted_boxes.tolist()
        topscores, topclasses = torch.topk(preds, k=1)
        topscores, topclasses = topscores.tolist(), topclasses.tolist()
        topclasses = np.concatenate(topclasses)
        topscores = np.concatenate(topscores)
                 
        #add top 1 prediction of behaviors in each time step
        for i in range(len(predicted_boxes_lst)):
            xleft.append(predicted_boxes_lst[i][0])
            ytop.append(predicted_boxes_lst[i][1])
            xright.append(predicted_boxes_lst[i][2])
            ybottom.append(predicted_boxes_lst[i][3])
            labels.append(label_map.get(topclasses[i]))
            time_frame.append(time_stamp)
            scores.append(topscores[i])

    print("Finished generating predictions.")
    # Generate Metadata file
    metadata = pd.DataFrame()
    metadata['frame'] = time_frame
    metadata['x_left'] = xleft
    metadata['y_top'] = ytop
    metadata['x_right'] = xright
    metadata['y_bottom'] = ybottom
    metadata['label'] = labels
    metadata['confidence'] = scores
            
    height, width = gif_imgs[0].shape[0], gif_imgs[0].shape[1]
    video_save_path = 'activity_recognition.mp4'
    video = cv2.VideoWriter(video_save_path, cv2.VideoWriter_fourcc(*'mp4v'), int(fps), (width, height))
    
    for image in gif_imgs:
        img = (255*image).astype(np.uint8)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        video.write(img)
    video.release()

    return video_save_path, metadata