from tdcnn.config import cfg
from utils.blob import prep_im_for_blob
from tdcnn.test import video_detect
import numpy as np
from utils.timer import Timer
from tdcnn.nms_wrapper import nms
from test_log import TestLog

def image_prep(frame):
    random_idx = [int(cfg.TEST.FRAME_SIZE[1]-cfg.TEST.CROP_SIZE) / 2,
                  int(cfg.TEST.FRAME_SIZE[0]-cfg.TEST.CROP_SIZE) / 2]
    frame = prep_im_for_blob(frame, cfg.PIXEL_MEANS, tuple(cfg.TRAIN.FRAME_SIZE[::-1]),
                            cfg.TRAIN.CROP_SIZE, random_idx)
    return frame

def create_running_blob(image, prev_buffer, buffer_size):
    """
    Create running blob based on the current image and the previous buffer
    1. If previous buffer < buffer_size:
        Create the input blob with the image in the end of the previous buffer
        and pad with the last frame to create the remaining buffer in the blob
    2. else:
        Extract the last buffer_size-1 images from the previous buffer and
        append the image to it
    """
    current_frame = image_prep(image)[np.newaxis, :]
    if prev_buffer is None:
        prev_buffer = np.array(current_frame, copy=True)
        buffer_array = np.repeat(current_frame, buffer_size-1, axis=0)
        input_blob = np.vstack((current_frame, buffer_array))
    elif prev_buffer.shape[0] < buffer_size-1 :
        prev_buffer = np.vstack((prev_buffer, current_frame))
        buffer_array = np.repeat( current_frame, buffer_size-prev_buffer.shape[0],
                                 axis=0 )
        input_blob = np.vstack((prev_buffer, buffer_array))
    else:
        input_blob = np.vstack([prev_buffer[1:, :, :, :], current_frame])
        prev_buffer = np.array(input_blob, copy=True)

    channel_swap = (3, 0, 1, 2)
    input_blob = input_blob.transpose(channel_swap)
    return input_blob[np.newaxis, :], prev_buffer

def create_blob(image, buffer_size):
    frame_blob = np.zeros((1, cfg.TEST.LENGTH[0], cfg.TEST.CROP_SIZE,
                            cfg.TEST.CROP_SIZE, 3))
    for current_buffer_size in range(buffer_size):
        frame_blob[0, current_buffer_size, :, :, :] = \
                                        image_prep(image)
    channel_swap = (0, 4, 1, 2, 3)
    frame_blob = frame_blob.transpose(channel_swap)
    return frame_blob


def test_net_online(net, image, current_frame_index, stride,
                    previous_buffer=None,  buffer_size=768,
                    max_per_image=100, thresh=0.05, vis=False,
                    use_running_frames=False, dataset_id=None):
    current_test_log = TestLog()
    box_proposals = None

    if use_running_frames:
        if previous_buffer is None:
            start_index = current_frame_index
        else:
            start_index = current_frame_index - previous_buffer.shape[0]*stride
        video, previous_buffer = create_running_blob(image, previous_buffer,
                                                     buffer_size)
        end_index = current_frame_index
    else:
        start_index = current_frame_index - stride
        end_index = current_frame_index
        video = create_blob(image, buffer_size)
    current_test_log.frames.extend([start_index,
                                    end_index, stride])
    current_test_log.bg_name = dataset_id
    current_test_log.fg_name = dataset_id
    scores, wins = video_detect(net, video, box_proposals)
    # skip j = 0, because it's the background class
    for j in range(1, cfg.NUM_CLASSES):
        inds = np.where(scores[:, j] > thresh)[0]
        cls_scores = scores[inds, j]
        cls_wins = wins[inds, j*2:(j+1)*2]
        cls_dets = np.hstack((cls_wins, cls_scores[:, np.newaxis])) \
            .astype(np.float32, copy=False)
        keep = nms(cls_dets, cfg.TEST.NMS)  #0.3
        if len(keep) != 0:
          cls_dets = cls_dets[keep, :]
          current_test_log.activities[j] = cls_dets
    return current_test_log, previous_buffer

