
# Copyright 2020-2022 OpenDR European Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import cv2
import time
import pandas as pd
import torch
from typing import Dict
import os
SELECTED_GPUS = [0]
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(gpu_number) for gpu_number in SELECTED_GPUS])

# opendr imports
import argparse
from opendr.perception.facial_expression_recognition import ProgressiveSpatioTemporalBLNLearner
#from opendr.perception.facial_expression_recognition.landmark_based_facial_expression_recognition.\
#    algorithm.datasets.landmark_extractor import landmark_extractor
#from opendr.perception.facial_expression_recognition.landmark_based_facial_expression_recognition.\
#    algorithm.datasets.gen_facial_muscles_data import gen_muscle_data


def tile(a, dim, n_tile):
    a = torch.from_numpy(a)
    init_dim = a.size(dim)
    repeat_idx = [1] * a.dim()
    repeat_idx[dim] = n_tile
    a = a.repeat(*repeat_idx)
    order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))
    tiled = torch.index_select(a, dim, order_index)
    return tiled.numpy()


def data_normalization(data):
    data = torch.from_numpy(data)
    N, V, C, T, M = data.size()
    data = data.permute(0, 2, 3, 1, 4).contiguous().view(N, C, T, V, M)
    # remove the first 17 points
    data = data[:, :, :, 17:, :]
    N, C, T, V, M = data.size()
    # normalization
    for n in range(N):
        for t in range(T):
            for v in range(V):
                data[n, :, t, v, :] = data[n, :, t, v, :] - data[n, :, t, 16, :]
    return data.numpy()


class VideoReader(object):
    def __init__(self, file_name):
        self.file_name = file_name
        try:  # OpenCV needs int to read from webcam
            self.file_name = int(file_name)
        except ValueError:
            pass

    def __iter__(self):
        self.cap = cv2.VideoCapture(self.file_name)
        if not self.cap.isOpened():
            raise IOError('Video {} cannot be opened'.format(self.file_name))
        return self

    def __next__(self):
        was_read, img = self.cap.read()
        if not was_read:
            raise StopIteration
        return img


def tile(a, dim, n_tile):
    a = torch.from_numpy(a)
    init_dim = a.size(dim)
    repeat_idx = [1] * a.dim()
    repeat_idx[dim] = n_tile
    a = a.repeat(*repeat_idx)
    order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))
    tiled_a = torch.index_select(a, dim, order_index)
    return tiled_a.numpy()


def lndmrk2numpy(args, num_current_frames, landmarks_list, num_frames):
    C = 3
    T = num_frames
    V = 68
    data_numpy = np.zeros((1, V, 2, num_current_frames, 1))
    seq = np.zeros((1, V, 2, T, 1))
    for t in range(num_current_frames):
        data_numpy[0, :, :, t, 0] = landmarks_list[t][:, :]

    # if we have less than num_frames, repeat frames to reach num_frames
    diff = T - num_current_frames
    if diff == 0:
        seq = data_numpy
    while diff > 0:
        num_tiles = int(diff / num_current_frames)
        if num_tiles > 0:
            data_numpy = tile(data_numpy, 3, num_tiles+1)
            num_current_frames = data_numpy.shape[3]
            diff = T - num_current_frames
        elif num_tiles == 0:
            seq[:, :, :, :num_current_frames, :] = data_numpy
            for j in range(diff):
                seq[:, :, :, num_current_frames+j, :] = data_numpy[:, :, :, -1, :]
            break
    return seq


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", help="Device to use (cpu, cuda)", type=str, default="cuda")
    parser.add_argument('--video', default=0,
                        help='path to video file or camera id')
    parser.add_argument('--model_name', type=str, default='pstbln_afew',
                        help='action detection method')
    parser.add_argument('--num_frames', type=int, default=300,
                        help='number of frames to be processed for each action')
    parser.add_argument('--fps', type=int, default=30,
                        help='number of frames per second to be processed by pose estimator and action detector')

    args = parser.parse_args()
    device = args.device
    if args.model_name == 'pstbln_ck+':
        num_point = 303
        num_class = 7
        num_frames = 5
    elif args.model_name == 'pstbln_casia':
        num_point = 309
        num_class = 6
        num_frames = 5
    elif args.model_name == 'pstbln_afew':
        num_point = 312
        num_class = 7
        num_frames = 150


    expression_classifier = ProgressiveSpatioTemporalBLNLearner(device="cuda", dataset_name='afew',
                                                                num_class=num_class,
                                                                num_point=num_point, num_person=1, in_channels=2,
                                                                blocksize=5, topology=[15, 10, 15, 5, 5, 10])

    expression_classifier.init_model()

    dev = f"cuda:0"
    torch.cuda.reset_peak_memory_stats(device=dev)

    '''image_provider = VideoReader('./1.mp4')  # loading a video or get the camera id 0

    counter, avg_fps = 0, 0
    landmarks_list = []
    window = int(30/args.fps)
    f_ind = 0
    for img in image_provider:
        if f_ind % window == 0:
            start_time = time.perf_counter()
            cv2.imwrite(os.path.join('./frames', str(counter) + ".jpg"), img)
            landmarks = landmark_extractor(os.path.join('./frames', str(counter) + ".jpg"), './temp', './shape_predictor_68_face_landmarks.dat')
            counter += 1
            print(counter)
            landmarks_list.append(landmarks)
            if counter > num_frames:
                landmarks_list.pop(0)
                counter = num_frames
            if counter > 0:
                seq = lndmrk2numpy(args, counter, landmarks_list, num_frames)
                norm_data = data_normalization(seq)
            muscle_data = np.zeros((1, 2, num_frames, num_point, 1)) #gen_muscle_data(norm_data, './temp')
            prediction = expression_classifier.infer(muscle_data, monte_carlo_dropout=False)
            end_time = time.perf_counter()
            fps = 1.0 / (end_time - start_time)
            avg_fps = 0.8 * fps + 0.2 * fps
        f_ind += 1
        key = cv2.waitKey(1)
        if key == ord('q'):
            break
    print("Average inference fps: ", avg_fps)
    cv2.destroyAllWindows()'''


    start_time = time.perf_counter()
    muscle_data = np.zeros((1, 2, num_frames, num_point, 1))  # gen_muscle_data(norm_data, './temp')
    prediction = expression_classifier.infer(muscle_data, monte_carlo_dropout=False)
    end_time = time.perf_counter()
    fps = 1.0 / (end_time - start_time)
    avg_fps = 0.8 * fps + 0.2 * fps
    print("Average inference fps: ", avg_fps)

    memory_summary = torch.cuda.memory_summary(device=dev, abbreviated=True)
    stats = {"memory_allocated": torch.cuda.memory_allocated(device=dev), "max_memory_allocated": torch.cuda.max_memory_allocated(device=dev),}
    print("memory_allocated", torch.cuda.memory_allocated(device=dev))
    print("max_memory_allocated", torch.cuda.max_memory_allocated(device=dev))

