import sys
import os
import csv
import cPickle as pickle

import cv2
import numpy as np
import caffe

from ACT_utils import *
from copy import deepcopy

from exp_config import experiment_config, expcfg_from_file
from virat_dataset import ViratDataset
import argparse
import json
import pickle as pkl

class ParticipatingActor:
    def __init__(self, start_frame, start_bbox, class_name, class_conf):
        self.start_frame = start_frame
        self.end_frame = start_frame + 1
        self.start_bbox = start_bbox
        self.current_bbox = start_bbox
        self.class_name = class_name
        self.class_conf = class_conf

    def merge_bbox(self, bbox, class_name, class_conf, frame_id):
        if _iou(self.current_bbox, bbox) > 0.5 and \
                self.class_name == class_name and \
                class_conf > 0.5 and \
                self.end_frame == frame_id:
            self.current_bbox = bbox
            self.end_frame = frame_id + 1
            return True
        else:
            return False

    def finish_tube(self, frame_id):
        if self.end_frame == frame_id:
            return True
        else:
            return False

def _get_image_name(image_number, buffer_size=5):
    return str(image_number).zfill(buffer_size) + "." + \
            experiment_config.data.img_extension

def extract_frames(virat_dataset, net_rgb, net_flow, video, resolution_array, h, w):
    video_dets = {}
    for image_index in xrange(1, 1 + virat_dataset.test_num_frames(video) - 
                                 - experiment_config.test.batch_size - 
                                 experiment_config.data.num_frames + 1,
                                 experiment_config.test.batch_size * \
                                         experiment_config.data.num_frames):
        kwargs_rgb  = {}
        kwargs_flo = {}
        for buffer_index in range(0, experiment_config.data.num_frames):
            image_batch = np.zeros([experiment_config.test.batch_size, 3,
                                    experiment_config.train.imgsize,
                                    experiment_config.train.imgsize])
            flow_batch = np.zeros([experiment_config.test.batch_size,
                                    3*experiment_config.test.number_flow,
                                    experiment_config.train.imgsize,
                                    experiment_config.train.imgsize])

            for batch_index in range(experiment_config.test.batch_size):
                im = cv2.imread(virat_dataset.imfile(video, 
                        _get_image_name(image_index + \
                                batch_index * experiment_config.data.num_frames + \
                                buffer_index)))
                if im is None:
                    print "Image {0} does not exist".\
                            format(virat_dataset.imfile(video, 
                                   _get_image_name(image_index+buffer_index)))
                    return video_dets
                imscale = cv2.resize(im, (experiment_config.train.imgsize,
                                         experiment_config.train.imgsize),
                                         interpolation=cv2.INTER_LINEAR)
                imscale = np.transpose(imscale - experiment_config.test.channel_mean,
                                        (2, 0, 1))
                imf = [cv2.imread(virat_dataset.flowfile(video, 
                        _get_image_name(min(virat_dataset.test_num_frames(video), 
                                        image_index + batch_index * \
                                        experiment_config.data.num_frames + buffer_index \
                                        + iflow)))) 
                        for iflow in xrange(experiment_config.test.number_flow)]
                if np.any(imf) is None:
                    print "Flow image does not exist"
                    return video_dets
                imscalef = [cv2.resize(im, (experiment_config.train.imgsize, 
                            experiment_config.train.imgsize), 
                            interpolation=cv2.INTER_LINEAR) for im in imf]

                timscale = [np.transpose(im-experiment_config.test.channel_mean, (2, 0, 1))
                                for im in imscalef]
                image_batch[batch_index, :, :, :] = imscale
                flow_batch[batch_index, :, :, :] = np.concatenate(timscale, axis=0)
                
            kwargs_rgb['data_stream' + str(buffer_index)] = image_batch

            kwargs_flo['data_stream' + str(buffer_index) + 'flow'] = \
                                            flow_batch
        # compute rgb and flow scores
        # two forward passes: one for the rgb and one for the flow 
        net_rgb.forward(end="mbox_conf_flatten", **kwargs_rgb) # forward of rgb with confidence and regression
        net_flow.forward(end="mbox_conf_flatten", **kwargs_flo) # forward of flow5 with confidence and regression
        
        # compute late fusion of rgb and flow scores (keep regression from rgb)
        # use net_rgb for standard detections, net_flo for having all boxes
        scores = 0.5 * (net_rgb.blobs['mbox_conf_flatten'].data + \
                net_flow.blobs['mbox_conf_flatten'].data)

        net_rgb.blobs['mbox_conf_flatten'].data[...] = scores
        net_flow.blobs['mbox_conf_flatten'].data[...] = scores
        net_flow.blobs['mbox_loc'].data[...] = net_rgb.blobs['mbox_loc'].data
        
        # two forward passes, only for the last layer 
        # dets is the detections after per-class NMS and thresholding (stardard)
        # dets_all contains all the scores and regressions for all tubelets 
        dets = net_rgb.forward(start='detection_out')['detection_out'][:, 0, :, 1:]
        dets_all = net_flow.forward(start='detection_out_full')['detection_out_full'][:, 0, :, 1:]
        if dets.shape[1]%experiment_config.test.batch_size != 0:
            dets = dets[:, :dets.shape[1]//experiment_config.test.batch_size*\
                    experiment_config.test.batch_size, :]

        dets = np.reshape(dets, [experiment_config.test.batch_size, 
                                 dets.shape[1]//experiment_config.test.batch_size, 
                                 dets.shape[2]])

        if dets_all.shape[1]%experiment_config.test.batch_size != 0:
            dets_all = dets_all[:, :dets_all.shape[1]//experiment_config.test.batch_size,
                                dets_all.shape[2], :]

        dets_all = np.reshape(dets_all, [experiment_config.test.batch_size, 
                                dets_all.shape[1]//experiment_config.test.batch_size, 
                                dets_all.shape[2]])
        # parse detections with per-class NMS
        if dets.shape[0] == experiment_config.test.batch_size and np.all(dets == -1):
            dets = np.empty((0, dets.shape[1]), dtype=np.float32)
        dets[:, :, 2:] *= resolution_array # network output was normalized in [0..1]
        dets[:, :, 2::2] = np.maximum(0, np.minimum(w, dets[:, :, 2::2]))
        dets[:, :, 3::2] = np.maximum(0, np.minimum(h, dets[:, :, 3::2]))
        # parse detections with global NMS at 0.7 (top 300)
        # coordinates were normalized in [0..1]
        dets_all[:, :, 0:4*experiment_config.data.num_frames] *= resolution_array 
        dets_all[:, :, 0:4*experiment_config.data.num_frames:2] = \
                np.maximum(0, np.minimum(w, dets_all[:, :, 0:4*experiment_config.data.num_frames:2]))
        dets_all[:, :, 1:4*experiment_config.data.num_frames:2] = np.maximum(0, np.minimum(h, dets_all[:, :, 1:4*experiment_config.data.num_frames:2]))
        for batch_index in range(experiment_config.test.batch_size):
            idx = nms_tubelets(np.concatenate((
                    dets_all[batch_index, :, :5*experiment_config.data.num_frames],
                    np.max(dets_all[batch_index, :, 4*experiment_config.data.num_frames+1:],
                        axis=1)[:, None]), axis=1), 0.7, 300)
            nms_dets_all = dets_all[batch_index, idx, :]
            video_dets[image_index+batch_index*experiment_config.data.num_frames] = nms_dets_all
        print "Completed: {0}/{1}".format(image_index + \
                experiment_config.test.batch_size*experiment_config.data.num_frames - 1,
                1 + virat_dataset.test_num_frames(video) - 
                                experiment_config.data.num_frames + 1)
        
    return video_dets

def _tubescore(tt):
    return np.mean(np.array([tt[i][1][-1] for i in xrange(len(tt))]))

def build_tubes(virat_dataset, video_detections, video):
    num_frames = virat_dataset.test_num_frames(video)
    result = {}
    for label in range(virat_dataset.nlabels):
        finished_tubes = []
        current_tubes = []
        for frame in xrange(1, virat_dataset.test_num_frames(video) + 2 - 
                experiment_config.data.num_frames, experiment_config.data.num_frames):
            
            # load boxes of the new frame and do nms while keeping Nkeep highest scored
            try:
                ltubelets = video_detections[frame][:,range(4*experiment_config.data.num_frames)
                    + [4*experiment_config.data.num_frames + label]] # Nx(4K+1) with (x1 y1 x2 y2)*K ilabel-score
            except Exception:
                continue
            idx = nms_tubelets(ltubelets, 0.3, top_k=10)
            ltubelets = ltubelets[idx,:]
            
            # just start new tubes
            if frame == 1:
                for i in xrange(ltubelets.shape[0]):
                    current_tubes.append( [(1,ltubelets[i,:])] )
                continue
            # sort current tubes according to average score
            avgscore = [_tubescore(t) for t in current_tubes]
            argsort = np.argsort(-np.array(avgscore))
            current_tubes = [current_tubes[i] for i in argsort]
            
            # loop over tubes
            finished = []
            for it, t in enumerate(current_tubes):
                # compute ious between the last box of t and ltubelets
                last_frame, last_tubelet = t[-1]
                ious = []
                offset = frame - last_frame
                if offset < experiment_config.data.num_frames:
                    nov = experiment_config.data.num_frames - offset
                    ious = sum([iou2d(ltubelets[:, 4*iov:4*iov+4], 
                                last_tubelet[4*(iov+offset):4*(iov+offset+1)]) 
                                for iov in xrange(nov)])/float(nov)
                else:
                    ious = iou2d(ltubelets[:, :4], 
                            last_tubelet[4*experiment_config.data.num_frames-4:
                                4*experiment_config.data.num_frames])
                
                valid = np.where(ious >= 0.2)[0]
                if valid.size>0:
                    # take the one with maximum score
                    idx = valid[ np.argmax(ltubelets[valid, -1])]
                    current_tubes[it].append((frame, ltubelets[idx,:]))
                    ltubelets = np.delete(ltubelets, idx, axis=0)
                else:
                    # skip
                    if offset>=experiment_config.data.num_frames:
                        finished.append(it)
            # finished tubes that are done
            for it in finished[::-1]: # process in reverse order to delete them with the right index
                finished_tubes.append( current_tubes[it][:])
                del current_tubes[it]
            
            # start new tubes
            for i in xrange(ltubelets.shape[0]):
                current_tubes.append([(frame,ltubelets[i,:])])

        # all tubes are not finished
        finished_tubes += current_tubes

        # TODO: Get rid of the hardcoded parameters and move them to exp_config
        # build real tubes
        output = []
        for t in finished_tubes:
            score = _tubescore(t)
            
            # just start new tubes
            if score< 0.1:
                continue
            
            beginframe = t[0][0]
            endframe = t[-1][0]+experiment_config.data.num_frames-1
            length = endframe+1-beginframe
            
            # delete tubes with short duration
            if length < 15:
                continue

            # build final tubes by average the tubelets
            out = np.zeros((length, experiment_config.data.num_frames), dtype=np.float32)
            out[:, 0] = np.arange(beginframe,endframe+1)
            n_per_frame = np.zeros((length, 1), dtype=np.int32)
            for i in xrange(len(t)):
                frame, box = t[i]
                for k in xrange(experiment_config.data.num_frames):
                    out[frame-beginframe+k, 1:5] += box[4*k:4*k+4]
                    out[frame-beginframe+k, -1] += box[-1]
                    n_per_frame[frame-beginframe+k ,0] += 1
            nz_out = out[np.where(n_per_frame>0)[0], :]
            n_per_frame = n_per_frame[n_per_frame>0][:, np.newaxis]
            nz_out[:, 1:] /= n_per_frame
            output.append((nz_out, score))
        result[label] = output
    return result

def find_files(current_directory):
    file_paths = []
    for file_name in os.listdir(current_directory):
        file_path = os.path.join(current_directory, file_name)
        if os.path.isdir(file_path):
            file_paths.extend(find_files(file_path))
        else:
            file_paths.extend([file_path])
    return file_paths

def _scale_bboxes(input_img_dimensions, output_img_dimensions, bbox):
    xmin, ymin, xmax, ymax = bbox[2:]
    xcenter = (xmin + xmax)/2.0
    ycenter = (ymin + ymax)/2.0
    width = xmax - xmin
    height = ymax - ymin
    inp_img_width, inp_img_height = input_img_dimensions
    op_img_width, op_img_height = output_img_dimensions
    scale_width_ratio = op_img_width/float(inp_img_width)
    scale_height_ratio = op_img_height/float(inp_img_height)
    scaled_width = width * scale_width_ratio
    scaled_height = height * scale_height_ratio
    shifted_xcenter = xcenter * scale_width_ratio
    shifted_ycenter = ycenter * scale_height_ratio
    shifted_xmin = shifted_xcenter - scaled_width/2
    shifted_xmax = shifted_xcenter + scaled_width/2
    shifted_ymin = shifted_ycenter - scaled_height/2
    shifted_ymax = shifted_ycenter + scaled_height/2
    # Corner cases
    if shifted_xmin < 0:
        shifted_xmin = 0
    if shifted_ymin < 0:
        shifted_ymin = 0
    if shifted_xmax > op_img_width:
        shifted_xmax = op_img_width
    if shifted_ymax > op_img_height:
        shifted_ymax = op_img_height

    bbox = (bbox[0], bbox[1], shifted_xmin, shifted_ymin, shifted_xmax, shifted_ymax)
    return bbox

def read_detections(detection_file):
    detection_dict = {}
    with open(detection_file, 'r') as det_f:
        detection_reader = csv.reader(det_f, delimiter=",")
        for file_index, detection in enumerate(detection_reader):
            if file_index == 0 or file_index == 1:
                continue
            frame_number, _,  xmin, ymin, xmax, ymax, conf, class_name, __= detection
            frame_number, xmin, ymin, xmax, ymax, conf = \
                    map(float, (frame_number, xmin, ymin, xmax, ymax, conf))
            bbox = (class_name, conf, xmin, ymin, xmax, ymax)
            bbox = _scale_bboxes([1024, 1024], [1920, 1080], bbox)
            if frame_number in detection_dict.keys():
                detection_dict[frame_number].append(bbox)
            else:
                detection_dict[frame_number] = [bbox]
    return detection_dict

def _iou(activity_bbox, object_bbox):
    act_xmin, act_ymin, act_xmax, act_ymax = activity_bbox
    obj_xmin, obj_ymin, obj_xmax, obj_ymax = object_bbox
    # xmin for intersection and union
    if act_xmin > obj_xmin:
        int_xmin = act_xmin
        un_xmin = obj_xmin
    else:
        int_xmin = obj_xmin
        un_xmin = act_xmin
    # ymin for intersection and union
    if act_ymin > obj_ymin:
        int_ymin = act_ymin
        un_ymin = obj_ymin
    else:
        int_ymin = obj_ymin
        un_ymin = act_ymin
    # xmax for intersection and union
    if act_xmax > obj_xmax:
        int_xmax = obj_xmax
        un_xmax = act_xmax
    else:
        int_xmax = act_xmax
        un_xmax = obj_xmax
    # ymax for intersection and union
    if act_ymax > obj_ymax:
        int_ymax = obj_ymax
        un_ymax = act_ymax
    else:
        int_ymax = act_ymax
        un_ymax = obj_ymax
    if int_xmin > int_xmax or int_ymin > int_ymax:
        intersection = 0
    else:
        intersection = (int_xmax - int_xmin) * (int_ymax - int_ymin)
    union = (act_xmax - act_xmin) * (act_ymax - act_ymin) + \
            (obj_xmax - obj_xmin) * (obj_ymax - obj_ymin) - \
            intersection
    return float(intersection)/union

def _create_object_annotation(participants, frame_id, object_id, video_name, 
                                is_last_frame=False):
    pruned_participants = []
    participant_annotation = []
    for participant in participants:
        if participant.finish_tube(frame_id) or is_last_frame:
            object_annotation = {
                "objectType": participant.class_name,
                "objectID": int(object_id),
                "localization": {
                        video_name: {
                            str(int(participant.start_frame)): {
                                    "presenceConf": float(participant.class_conf),
                                    "boundingBox": {
                                        "x": int(participant.start_bbox[0]),
                                        "y": int(participant.start_bbox[1]),
                                        "w": int(participant.start_bbox[2] - 
                                                participant.start_bbox[0]),
                                        "h": int(participant.start_bbox[3] -
                                                participant.start_bbox[1])
                                        }
                                },
                            str(int(participant.end_frame)): {}
                        }
                }
            }
            object_id += 1
            participant_annotation.append(object_annotation)
        else:
            pruned_participants.append(participant)
    return pruned_participants, participant_annotation, object_id

def compute_participating_objects(bboxes, darknet_detections, video_name, merge_participants):
    participant_annotation = []
    participants = []
    object_id = 1
    for box_index in range(bboxes.shape[0]):
        frame_id = bboxes[box_index, 0]
        activity_bbox = bboxes[box_index, 1:]
        activity_bbox = activity_bbox[0:4]
        if frame_id in darknet_detections.keys():
            object_detections = darknet_detections[frame_id]
        else:
            object_detections = []
        for object_detection in object_detections:
            class_name = object_detection[0]
            class_conf = object_detection[1]
            obj_bbox = object_detection[2:]
            if _iou(activity_bbox, obj_bbox) > 0.5 and class_conf>0.3:
                if merge_participants:
                    for participant in participants:
                        if participant.merge_bbox(obj_bbox, class_name, class_conf, frame_id):
                            is_merged = True
                            break
                else:
                    is_merged = False

                if not is_merged:
                    participants.append(ParticipatingActor(frame_id, 
                        obj_bbox, class_name, class_conf))
                participants, current_participant_annotation, object_id = \
                        _create_object_annotation(participants, frame_id, \
                                                    object_id, video_name)
                participant_annotation.extend(current_participant_annotation)

    _, current_participant_annotation, __ = \
            _create_object_annotation(participants, frame_id, object_id, video_name, 
                                        True)
    participant_annotation.extend(current_participant_annotation)
    return participant_annotation

def tubes_to_json(virat_dataset, all_tubes, json_path, all_detections, \
                    eval_aod, merge_participants):
    # convert tubes into nist specified format
    files_processed = []
    activities = []
    activity_id = 1
    index = 0
    for video, labeled_tubes in all_tubes.iteritems():
        video_name = os.path.basename(video)+ ".avi"
        if video not in files_processed:
            files_processed.append(video_name)

        for label, tubes in labeled_tubes.iteritems():
            for tube in tubes:
                bboxes, score = tube
                if label == 0:
                    continue
                elif float(score) < 0.1:
                    continue
                if eval_aod:
                    participating_objects = compute_participating_objects(
                                                bboxes, all_detections[video],
                                                video_name, merge_participants)
                else:
                    participating_objects = None
                upperbound_video = (virat_dataset.test_num_frames(
                        virat_dataset._get_video_prefix(video_name))//(
                                experiment_config.data.num_frames* \
                                        experiment_config.test.batch_size)) * \
                                experiment_config.data.num_frames * \
                                experiment_config.test.batch_size
                if bboxes[0][0] == 1 and \
                        bboxes[bboxes.shape[0]-1][0] >= upperbound_video:
                            continue
                class_label = virat_dataset.labels.keys()[virat_dataset.labels.\
                                    values().index(label)]
                if eval_aod:
                    if len(participating_objects) > 0:
                        current_activity = {"activity": class_label, 
                                            "activityID": int(activity_id),
                                            "presenceConf": float(score),
                                            "alertFrame": int(bboxes[0][0]),
                                            "localization": {
                                                video_name: {
                                                    str(int(bboxes[0][0])): 1,
                                                    str(int(bboxes[bboxes.shape[0]-1][0])): 0
                                                }    
                                                },
                                            "objects": participating_objects
                                            }
                    else:
                        # skip annotations with no paticipating objects
                        continue
                else:
                    current_activity = {"activity": class_label, 
                                        "activityID": int(activity_id),
                                        "presenceConf": float(score),
                                        "alertFrame": int(bboxes[0][0]),
                                        "localization": {
                                            video_name: {
                                                str(int(bboxes[0][0])): 1,
                                                str(int(bboxes[bboxes.shape[0]-1][0])): 0
                                            }    
                                            }
                                        }
                activities.append(current_activity)
                activity_id += 1
        print "Complete {0}/{1}".format(index+1, len(all_tubes))
        index += 1
    json_dict = {"filesProcessed": files_processed, "activities": activities}
    with open(json_path, 'w') as json_f:
        json.dump(json_dict, json_f, indent=4)

def main(args):
    virat_dataset = ViratDataset( experiment_config.data.frame_roots, 
                                  experiment_config.data.flow_roots,
                                  experiment_config.data.train_annotation_dirs,
                                  experiment_config.data.test_annotation_dirs,
                                  experiment_config.data.class_index, 
                                  experiment_config.data.save_directory,
                                  experiment_config.train.kpf_mode,
                                  experiment_config.train.json_mode,
                                  experiment_config.data.save_prefix,
                                  True )
    caffe.set_mode_gpu()
    caffe.set_device(experiment_config.gpu)
    model_dir = experiment_config.train.model_dir
    output_dir = model_dir
    
    rgb_proto = os.path.join(model_dir, "deploy_RGB.prototxt")
    rgb_model = os.path.join(model_dir, "virat_RGB_iter_" + str(args.model_iter) + \
                                        ".caffemodel")
    if not os.path.exists(rgb_model):
        raise Exception("rgb model not found")
    net_rgb = caffe.Net(rgb_proto, caffe.TEST, weights=rgb_model)
    
    caffe.set_device(experiment_config.gpu)
    # load the FLOW5 network
    flow_proto = os.path.join(model_dir, "deploy_FLOW5.prototxt")
    flow_model = os.path.join(model_dir, "virat_FLOW5_iter_" + str(args.model_iter) + \
                                        ".caffemodel")
    if not os.path.exists(flow_model):
        raise Exception("flow model not found")
    net_flow = caffe.Net(flow_proto, caffe.TEST, weights=flow_model)
    
    # Initialization for testing
    all_videos = virat_dataset.test_video_list
    all_tubes = {}
    all_detections = {}

    if not os.path.exists(experiment_config.test.cache_dir):
        os.makedirs(experiment_config.test.cache_dir)
    
    # Iterate over model
    for video_index, video in enumerate(all_videos):
        # extract frames
        print("Processing video {:d}/{:d}: {:s}".format( video_index+1, 
            len(all_videos), video))
        h, w = virat_dataset.test_video_resolution(video)
        resolution_array = np.array([w,h,w,h]*experiment_config.data.num_frames,
                                    dtype=np.float32)
        activity_cache_path = os.path.join( experiment_config.test.cache_dir, 
                                        "activity_detection_" + 
                                        str(os.path.basename(video)) + ".pkl" )
        if not os.path.exists( activity_cache_path ):
            video_detections = extract_frames(virat_dataset, net_rgb, net_flow,video,
                                            resolution_array, h, w)
            with open( activity_cache_path, "w" ) as f:
                pkl.dump(video_detections, f)
        else:
            print "Loading activity detection from: " + str(activity_cache_path)
            video_detections = pkl.load(open( activity_cache_path, "r"))
       
        print("Building tubes")
        # build tubes
        video_tubes = build_tubes(virat_dataset, video_detections, video)
        all_tubes[video] = video_tubes

        # Object detection
        if args.eval_aod:
            object_cache_path = os.path.join( experiment_config.test.cache_dir, 
                                        "object_detection_" + 
                                        str(os.path.basename(video)) + ".pkl" )
            if not os.path.exists( object_cache_path ):
                object_detection = read_detections( 
                                    os.path.join( experiment_config.test.detection_dir, 
                                    os.path.basename(video) + ".csv" ) )
                with open( object_cache_path, "w" ) as f:
                    pkl.dump(object_detection, f )
            else:
                print "Loading object detections from: " + str(object_cache_path )
                object_detection = pkl.load( open( object_cache_path, "r" ) )

            all_detections[video] = object_detection
    tubes_to_json(virat_dataset, all_tubes, args.json_path, all_detections,
                     args.eval_aod, args.merge_participants)

if __name__ == "__main__":    
    parser = argparse.ArgumentParser("Test ACT-detector")
    parser.add_argument("--exp", 
            help="yml configuration for experiment parameters")
    parser.add_argument("--json_path",
            help="path to save the json file (Includes the name of the file too)")
    parser.add_argument("--eval_aod", action="store_true",
            help="flag to use ACT as an AOD detector")
    parser.add_argument("--merge_participants", action="store_true",
            help="flag to merge aod object detections")
    parser.add_argument("--model_iter", type=int, default=110000,
            help="Model iteration to test")
    parser.set_defaults(eval_aod=False)
    parser.set_defaults(merge_participants=False)
    args = parser.parse_args()
    if args.exp is not None:
        expcfg_from_file(args.exp)
    main(args)

