'''
vis_cnn_predictions.py visualises the predictions from the CNN models
'''


#TODO BROKEN CODE.... NEED TO FIX FOR NEW CODE!!!
import math
import os
import time
from collections import deque
import pickle
import numpy as np

import keras_code.models as models
from keras_code import predict
import utilities.paths as paths


DRIVE = paths.get_drive()

# def save_timelines(out_img_size,labels,dq,colours):
#
#     frame_img = np.ones((1, out_img_size * 2 + 5 + 10 + 5, 3), dtype=np.uint8) * 255
#     top_ = 0
#     bottom_ = 0
#     for class_num in range(len(labels)):
#         bottom_ = top_
#         top_ += (dq[len(dq) - 1][class_num]) * out_img_size
#
#         frame_img[:, int(round(bottom_)):max(int(round(top_)), out_img_size), :] = colours[class_num]
#
#     frame_img[:, (out_img_size + 5):(out_img_size + 15), :] = colours[np.argmax(dq[len(dq) - 1])]
#
#     frame_img[:, out_img_size + 5 + 15:out_img_size * 2 + 5 + 15, :] = colours[np.argmax(clip['class_scores'])]  # TODO make these framewise scores? in case clips don't all have saem score? requires work in dataset reader for new frame dict format
#
#     frame_img[:, out_img_size + 1 + 15:out_img_size + 4 + 15, :] = [40, 40, 40]
#
#     if cnt % 250 == 0:
#         mid_cnt += 35
#         frame_img[:, out_img_size + 1 + 15:out_img_size + 4 + 15, :] = [200, 200, 200]
#
#     if cnt % 1500 == 0:
#         mid_cnt = 40
#
#     # frame_img[:, out_img_size + 1 + 15:out_img_size+4 + 15, :] = [40,40,40]#[mid_cnt,mid_cnt,mid_cnt]
#
#     cnt += 1
#     out_img = np.append(out_img, frame_img, axis=0)
#     out_min_img = np.append(out_min_img, frame_img, axis=0)
#
#     if cnt % 1500 == 0:
#         cv2.imwrite(time_img_name + '_all.png', np.swapaxes(out_img, 0, 1))
#
#         cv2.imwrite('%s_m%02d.png' % (time_img_name, int(round(cnt / 1500))), np.swapaxes(out_min_img, 0, 1))
#         out_min_img = np.ones((1, out_img_size * 2 + 5 + 10 + 5, 3), dtype=np.uint8) * 20


def visualise(model_id, identifier, model_path, split, batch_size, load_epoch):
    layer_name = 'pred'

    save_video = False
    show_video = False
    save_timeline = True
    show_commentary = False

    if show_video:
        save_video = True

    if save_video:
        # only import opencv if saving video
        import cv2
        from cv2 import __version__
        cv_ver = __version__.split('.')

        NEW_OCV = True
        if int(cv_ver[0]) < 3:
            NEW_OCV = False



    model_path = model_path + model_id + '_' + identifier
    if not os.path.exists(model_path):
        os.makedirs(model_path)

    if load_epoch is not None:
        model, output_classes = models.get_model_from_id(model_id, identifier, batch_size, load_epoch=load_epoch)
    else:
        print('WARNING: No load_epoch specified')

    assert (model is not None), "model is None"

    # get data
    dataset = models.get_dataset(model_id, split, force_noneq=True, batch_size=batch_size)

    dataset.frame_order()

    # Check if preditions saved, if not, save them
    pred_path = paths.get_drive(2) + 'DATASETS/VIDEO/TENNIS/FEATURES/'+ model_id + '_' + identifier + '-e' + str(load_epoch)+'/'+ layer_name + '/pkl/' + split + '.pkl'
    if not os.path.exists(pred_path):
        print('Need to run predictions')
        predict(model_id, identifier, model_path, split, batch_size, load_epoch, True, layer_name,
                save_path=paths.get_drive(2) + 'DATASETS/VIDEO/TENNIS/FEATURES/',
                dataset=dataset)
    predictions = pickle.load(open(pred_path, 'rb'), encoding='latin1')




    vid_dest_path = dataset.get_dataset_path() + 'VISUALISATIONS/' + model_id + '_' + str(load_epoch) + '_' + identifier + '/PREDICTIONS/'
    if not os.path.exists(vid_dest_path):
        os.makedirs(vid_dest_path)

    labels = dataset.get_labels()

    if save_video:
        if NEW_OCV:
            video = cv2.VideoWriter(
                vid_dest_path + model_id + '_' + str(load_epoch) + '_' + identifier + '_' + split + '_hard.avi',
                cv2.VideoWriter_fourcc('F', 'M', 'P', '4'), 25, (1728, 972))
        else:
            video = cv2.VideoWriter(
                vid_dest_path + model_id + '_' + str(load_epoch) + '_' + identifier + '_' + split + '_hard.avi',
                cv2.cv.CV_FOURCC('F', 'M', 'P', '4'), 25, (1728, 972))


    if save_timeline:
        out_img_size = 40
        out_img = np.ones((1,out_img_size*2 +5+10+5,3), dtype=np.uint8)*255
        out_min_img = np.ones((1,out_img_size*2 +5+10+5,3), dtype=np.uint8)*255
        cnt = 0
        mid_cnt = 40
        time_img_name = vid_dest_path + model_id + '_' + str(load_epoch) + '_' + identifier  + '_' + split + '_hard'

    colours = [[200, 200, 200], [35, 118, 211], [34, 190, 204], [191, 185, 31], [79, 180, 30], [31, 31, 191],
               [160, 26, 26], [100, 200, 0], [200, 100, 0], [200, 0, 100], [0, 200, 100], [100, 0, 200], [0, 100, 200]]

    dq = deque()
    dqgt = deque()
    for k in dataset.keys:
        sid = dataset.get_data_key(k)

        # get points data for commentary
        # if show_commentary:
        #     with open(DRIVE + 'DATASETS/VIDEO/TENNIS/SPLITS/' + split_id + '/CLASSES/Point_splits.txt') as f:
        #         data = f.readlines()
        #     data = [line.rstrip().split() for line in data]
        #
        #     if model_id == 'MVK_10_02' and (load_epoch == 6 or load_epoch == 10):
        #         sent_model, best_epoch = 'MSK_01_30',220 # 128 dim, no extra feats
        #     elif model_id == 'MQK_46_02' and load_epoch == 10:
        #         sent_model, best_epoch = 'MSK_04_20', 260  # 256 dim, no extra feats
        #     elif model_id == 'MQK_46_01' and load_epoch == 4:
        #         sent_model, best_epoch = 'MSK_03_20', 280  # 256 dim, no extra feats
        #     elif model_id == 'MQK_12_01' and load_epoch == 1:
        #         sent_model, best_epoch = 'MSK_03_20', 280  # 256 dim, no extra feats,
        #     else:
        #         print 'ERROR: CANT FIND SENTENCE MODEL!!'
        #
        #     with open(DRIVE + 'MODELS/SENTENCE/KERAS/'+sent_model+'/RESULTS/'+split_name+'_'+str(best_epoch)+'.txt') as f:
        #         dataC = f.readlines()
        #     dataC = [line.rstrip() for line in dataC]
        #
        #     li = 0
        #     pid_tmp = ''
        #     gt_tmp = ''
        #     p_dict = {}
        #     for line in dataC:
        #         if line[0] == 'P':
        #             li = 2
        #             pid_tmp = line
        #         elif li == 2:
        #             li = 3
        #             gt_tmp = line
        #         elif li == 3:
        #             pred_tmp = line
        #             p_dict[pid_tmp] = [gt_tmp, pred_tmp]
        #             pid_tmp = ''
        #             gt_tmp = ''
        #         else:
        #             li = 0
        #
        #     comm_data = []
        #     for line in data:
        #         pid = line[0]
        #         atts = line[1].split(',')
        #
        #         if atts[0] == split_name.lower() and atts[1] == video_id:
        #             comm_data.append([pid, atts[2], atts[3], p_dict[pid][0], p_dict[pid][1]]) # point_id, start, finish, gt_sent, pred_sent

        done = 0

        dqgt.append(predictions[sid]['gt'])
        if len(dqgt) > 1500:
            dqgt.popleft()

        dq.append(predictions[sid][layer_name])

        if len(dq)>1500:
           dq.popleft()

        if save_video:
            display_frame = np.ones((972, 1728, 3), dtype=np.uint8) * 30
            frame = dataset.get_rgb_frame(k, False)
            frame = frame[:, ::-1, :, :]  # 'RGB'->'BGR'
            frame = np.swapaxes(frame,1,2)
            frame = np.swapaxes(frame,2,3)
            display_frame[-512:,:910,:] = cv2.resize(frame[0],(910,512))

            # Hardcoded crop to focus on tennis court

            # for i in range(256):
            #     for j in range(256):
            #         display_frame[2*i:2*i+2,920+2*j:920+2*j+2,:] = frame2[i,j,:]
            # display_frame[-512:,1000:1512,:] = np.repeat(np.repeat(cv2.resize(frame,(256,256)),2, axis=0), 2, axis=1)
            # if is_flow:
            #     display_frame[-512:,1000:1512,:] = cv2.resize(cv2.imread(),(512,512))
            # else:
            #     display_frame[-512:,1000:1512,:] = cv2.resize(frame,(512,512))

            dim = 224
            img = dataset.get_rgb_frame(k, False)

            img = img[:, ::-1, :, :]  # 'RGB'->'BGR'
            img = np.swapaxes(img,1,2)
            img = np.swapaxes(img,2,3)
            display_frame[-512+int((512-dim)/2.0):-int((512-dim)/2.0),1000+int((512-dim)/2.0):1512-int((512-dim)/2.0),:] = \
                cv2.resize(img[0],(dim,dim))

            display_frame[-519:-512,:,:] = 200

            extra = 0

            # Write class names to frame
            vert_slices = int(math.floor((900-512)/(len(dataset.get_labels())+1+1+extra)))
            vs_wp = vert_slices-10
            for i in range(len(dataset.get_labels())):
                bottom = (i+1)*vert_slices
                for c in range(len(dq)):
                    top = (1-dq[len(dq)-c-1][i])*vs_wp+10+i*vert_slices
                    left = 1500-c
                    right = 1501-c
                    display_frame[int(min(top,bottom-1)):int(bottom),left:right,:] = colours[i]


                display_frame[bottom-vs_wp:bottom,1510:1515,:] = colours[i]
                cv2.putText(display_frame,labels[i].rstrip(), (1520,int(bottom)-15), cv2.FONT_HERSHEY_SIMPLEX, .4, (200,200,200))

            bottom = (len(dataset.get_labels())+1)*vert_slices

            # if save_timelines:  # write out image file
            #     save_timelines()

            # if show_commentary:
            #     for comm in comm_data:
            #         if int(comm[1]) <= fn <= int(comm[2]):
            #             # this frame is part of a point so print out sent
            #             cv2.putText(display_frame, 'GT:    '+comm[3], (10, int(bottom) + 60), cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 200, 0))
            #             cv2.putText(display_frame, 'PRED:  '+comm[4], (10, int(bottom) + 90), cv2.FONT_HERSHEY_SIMPLEX, .5, (0, 200, 200))

            for c in range(len(dq)):
                best_ind = np.argmax(dq[len(dq)-c-1])
                gt_ind = np.argmax(dqgt[len(dqgt)-c-1])
                top = (1-dq[len(dq)-c-1][best_ind])*vs_wp+10+len(dataset.get_labels())*vert_slices
                left = 1500-c
                right = 1501-c
                display_frame[int(min(top,bottom-1)):int(bottom),left:right,:] = colours[best_ind]

                display_frame[bottom+5:bottom+10,left:right,:] = colours[best_ind]

                display_frame[bottom+20:bottom+25,left:right,:] = colours[gt_ind]

                display_frame[bottom+30:bottom+35,left:right,:] = colours[gt_ind]


            cv2.putText(display_frame,'COMBINED', (1510,int(bottom)-15), cv2.FONT_HERSHEY_SIMPLEX, .4, (200,200,200))


            for rep in range(dataset.framegap):
                video.write(display_frame)
            done += 1


            if show_video:
                cv2.imshow('Display',display_frame)
                # Press the q key to exit
                if cv2.waitKey(1) & 0xFF == ord('q'):
                   break

    if save_video:
        video.release()
        cv2.destroyAllWindows()

    if save_timeline:
        cv2.imwrite(time_img_name+'_all.png', np.swapaxes(out_img, 0, 1))


if __name__ == '__main__':

    model_id = 'MVK_50_61'
    identifier = '00001'
    model_path = DRIVE + 'MODELS/VISUAL/KERAS/'
    splits = ['train', 'val', 'test']
    force_noneq = True
    batch_size = 16
    load_epoch = 4
    layers = ['pred']  # can only do one atm
    save_path = None
    save_path = paths.get_drive(2) + 'DATASETS/VIDEO/TENNIS/FEATURES/'

    for split in splits:
        visualise(model_id, identifier, model_path, split, batch_size, load_epoch)