import cv2
import os, random
import numpy as np
from keras.applications.vgg16 import VGG16, preprocess_input
from keras.preprocessing.image import img_to_array
from keras.optimizers import SGD

MAX_NB_CLASSES = 101


def extract_vgg16_features_live(model, video_input_file_path, video_model, expected_frames, labels_idx2word, label, correct_count, count1):
    print('Extracting frames from video: ', video_input_file_path)
    vidcap = cv2.VideoCapture(video_input_file_path)
    fps    = vidcap.get(cv2.CAP_PROP_FPS)
    print(fps)
    success, image = vidcap.read()
    features = []
    success = True
    count = 0
    text = ""
    while success:
        #vidcap.set(cv2.CAP_PROP_POS_MSEC, (count * 1000))  # added this line
        success, image = vidcap.read()
        # print('Read a new frame: ', success)
        if success:
            img = cv2.resize(image, (224, 224), interpolation=cv2.INTER_AREA)
            input = img_to_array(img)
            input = np.expand_dims(input, axis=0)
            input = preprocess_input(input)
            feature = model.predict(input).ravel()
            features.append(feature)
            count = count + 1

            font = cv2.FONT_HERSHEY_SIMPLEX
            cv2.putText(img, text, (5, 20), font, 0.5, (random.randint(0,255),random.randint(0,255),random.randint(0,255)), 2)
            cv2.imshow('Frame', img)

            if cv2.waitKey(25) & 0xFF == ord('q'):
                break
        if(count%20 == 0):
            features = np.array(features)
            frames = features.shape[0]
            if frames == 0:
                continue
            if frames > expected_frames:
                features = features[0:expected_frames, :]
            elif frames < expected_frames:
                print(features.shape)
                temp = np.zeros(shape=(expected_frames, features.shape[1]))
                temp[0:frames, :] = features
                features = temp
            predicted_class = np.argmax(video_model.predict(np.array([features]))[0])
            predicted_label = labels_idx2word[predicted_class]

            text = predicted_label + " - " + label
            """correct_count = correct_count + 1 if label == predicted_label else correct_count
            count1 += 1
            accuracy = correct_count / count1
            print('accuracy: ', accuracy)"""
            print(text)
            features = []
    #print(count)
    #unscaled_features = np.array(features)
    return correct_count, count1

def extract_vgg16_features(model, video_input_file_path, feature_output_file_path):
    if os.path.exists(feature_output_file_path):
        return np.load(feature_output_file_path)
    count = 0
    print('Extracting frames from video: ', video_input_file_path)
    vidcap = cv2.VideoCapture(video_input_file_path)
    success, image = vidcap.read()
    features = []
    success = True
    while success:
        vidcap.set(cv2.CAP_PROP_POS_MSEC, (count * 1000))  # added this line
        success, image = vidcap.read()
        # print('Read a new frame: ', success)
        if success:
            img = cv2.resize(image, (224, 224), interpolation=cv2.INTER_AREA)
            input = img_to_array(img)
            input = np.expand_dims(input, axis=0)
            input = preprocess_input(input)
            feature = model.predict(input).ravel()
            features.append(feature)
            count = count + 1
    unscaled_features = np.array(features)
    np.save(feature_output_file_path, unscaled_features)
    return unscaled_features


def scan_and_extract_vgg16_features(data_dir_path, output_dir_path, model=None, data_set_name=None):
    if data_set_name is None:
        data_set_name = 'UCF-101'
    #print(data_set_name)
    input_data_dir_path = data_dir_path + '/' + data_set_name
    output_feature_data_dir_path = data_dir_path + '/' + output_dir_path

    if model is None:
        model = VGG16(include_top=True, weights='imagenet')
        model.compile(optimizer=SGD(), loss='categorical_crossentropy', metrics=['accuracy'])
    
    if not os.path.exists(output_feature_data_dir_path):
        os.makedirs(output_feature_data_dir_path)

    y_samples = []
    x_samples = []

    dir_count = 0
    for f in os.listdir(input_data_dir_path):
        file_path = input_data_dir_path + os.path.sep + f
        #print(file_path)
        if not os.path.isfile(file_path):
            #print('yes')
            output_dir_name = f
            output_dir_path = output_feature_data_dir_path + os.path.sep + output_dir_name
            if not os.path.exists(output_dir_path):
                print(output_dir_path)
                os.makedirs(output_dir_path)
            dir_count += 1
            for ff in os.listdir(file_path):
                video_file_path = file_path + os.path.sep + ff
                output_feature_file_path = output_dir_path + os.path.sep + ff.split('.')[0] + '.npy'
                x = extract_vgg16_features(model, video_file_path, output_feature_file_path)
                y = f
                y_samples.append(y)
                x_samples.append(x)

        if dir_count == MAX_NB_CLASSES:
            break

    return x_samples, y_samples

