from Environment import Env
from BrainPG import PolicyGradient
from BrainAC import ActorCritic
from BrainDDPG import DDPG
import numpy as np
import random
from  datetime import*
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import copy
import os
import tensorflow as tf

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

class TikcetPlay():
    def __init__(self):
        self.order = []
        self.routeline = []
        # self.allRoute = self.readRoute("/home/pycharm_workspace/Ticket_bak/data/route")
        # self.allOrder = self.readOrder("/home/pycharm_workspace/Ticket_bak/data/order")
        self.allRoute = self.readRoute("./data/route")
        self.allOrder = self.readOrder("./data/order")
        self.epsilon = 0.9
        # print(self.allRoute[21])
        # print(self.allOrder[0])

    def transcate_PG(self):
        total_steps = 0  # 记录步数，一天是一步
        profit_list = []  # 记录每局总收益
        profitAdvanced_list = []
        actions = 2  # 行动个数
        brain = PolicyGradient(
            n_actions=2,
            n_features=87,
            learning_rate=0.1,
            reward_decay=1,
        )
        gameNum = 0 #记录游戏轮数
        ex_steps = 500 #探索衰减的轮数
        epsilon = self.epsilon
        last_remainder  = 0
        reward_list = [0] #存储每次的收益，来计算baseline
        Loss_list = [] #存储训练过程中的损失值
        wait_list = [] #记录等待天数
        gameSplit = 500 #每多少轮游戏画图
        while total_steps < 60000:
            # 初始化游戏
            # routeId = random.randrange(0, 49, 1)
            routeId = 21
            self.routeline = self.allRoute[routeId]
            # print(self.routeline)
            env = Env(self.routeline)
            gameNum += 1
            # state = env.getState()  # 以state[0]、state[1]方式访问
            today = env.getToday()

            terminal = False
            order_accepted = False
            isExploration = False
            create_date = 1
            end_date = 0
            stay_num = 0
            # 一局游戏
            # print("GAME#:",gameNum)
            baseline = 0
            tao_prob = []
            tao_reward = 0
            wait_day = []

            while today < self.routeline[-1] and terminal == False:
                # 有新订单产生 (当订单数已满10个时，此处不会收到新订单)
                if order_accepted == False:
                    self.orderSelect(self.routeline,60)
                    # print(self.order)
                    env.setOrder(self.order)
                    order_accepted = True
                    # print(self.order[1])
                # 遍历self.orders(即state[0])字典，对每一个订单操作
                state = env.getState()

                # 当前状态
                state_tf = np.mat(state)
                # print(state_tf,len(state_tf))
                # 由神经网络选择行动
                if random.random()<epsilon and isExploration == False:
                    isExploration = True
                    end_date = random.randrange(env.getTodayIndex(),87,1)
                    # end_date = 60

                if isExploration:
                    if env.getTodayIndex() == end_date:
                        action_model = 1
                        if ex_steps>0:
                            ex_steps -= 1
                    else:
                        action_model = 0
                else:
                    #action from learning
                    action_model,p = brain.choose_action(state_tf, env.getTodayIndex())
                    tao_prob.append(p)
                if action_model == 0:
                    action_finishOrder = [1, 0]
                else:
                    action_finishOrder = [0, 1]

                # 订单字典 历史曲线 reward

                reward = env.getReward(action_model)

                # 订单完成或者到最后一天
                terminal = env.isTerminal(action_model)
                if terminal:
                    tmp = reward
                    baseline = np.mean(reward_list)
                    profitAdvanced_list.append(baseline)
                    reward -= baseline
                    reward_list.append(tmp)
                    # print("END_REWARD:",reward,",reward_list:",reward_list)
                # 保存记录到记忆库
                # print("this is store arg:",state_tf,";", action_model,";", reward,";", env.getTodayIndex())
                brain.store_transition(state_tf, action_model, reward, env.getTodayIndex())
                # print(action_model)

                total_steps += 1
                if terminal:
                    loss,wait_day,tao_reward = brain.learn()
                    Loss_list.append(loss)
                    wait_list.append(wait_day[-1])
                    break

                # step 过一天加一
                env.nextStep()

            # 一局的总收益
            epsilon = self.epsilon*(ex_steps/500)
            print("epsilon:",epsilon)
            print("Baseline:",baseline)
            profit = env.getTotalReward()
            profit_list.append(profit)
            print("total_steps:",total_steps)
            print("profit_list", profit_list)
            print("profit:", profit, "profitAvg:", np.mean(profit_list))
            print("action-prob:",tao_prob)
            print("Reward:",tao_reward)
            print("wait_day:",wait_day)
            self.writeHistory('./picture/history.txt',
                              epsilon,baseline,total_steps,profit_list,profit,tao_prob,tao_reward,wait_day,gameNum)

            print("########################"+str(gameNum)+"###########################")
            if len(profit_list) >=gameSplit:
                plt.figure()
                plt.plot(profit_list, 'r-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_profit_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), profit_list)
                plt.savefig('./picture/' + str(gameNum) + 'scatter_profit_PG.jpg')
                plt.figure()
                plt.plot(profitAdvanced_list, 'g-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_advanced_PG.jpg')
                plt.figure()
                plt.plot(Loss_list, 'y-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_loss_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), wait_list,c='r')
                plt.savefig('./picture/' + str(gameNum) + 'scatter_waitDay_PG.jpg')
                profit_list.clear()
                wait_list.clear()
            # last_remainder = total_steps % 1000

    def transcate_AC(self):
        total_steps = 0  # 记录步数，一天是一步
        profit_list = []  # 记录每局总收益
        profitAdvanced_list = []
        actions = 2  # 行动个数
        brain = ActorCritic(
            n_actions=2,
            n_features=87,
            LR_A=0.001,
            LR_C=0.01,
            reward_decay=1.,
            prob_clip=0.,
        )
        gameNum = 0 #记录游戏轮数
        ex_steps = 500 #探索衰减的轮数
        epsilon = self.epsilon
        last_remainder  = 0
        reward_list = [0] #存储每次的收益，来计算baseline
        Loss_list = [] #存储训练过程中的损失值
        wait_list = [] #记录N轮游戏分别等待天数
        gameSplit = 500 #每多少轮游戏画图
        while total_steps < 60000:
            # 初始化游戏
            # routeId = random.randrange(0, 49, 1)
            routeId = 21
            self.routeline = self.allRoute[routeId]
            # print(self.routeline)
            env = Env(self.routeline)
            gameNum += 1
            # state = env.getState()  # 以state[0]、state[1]方式访问
            today = env.getToday()

            terminal = False
            order_accepted = False
            isExploration = False
            create_date = 1
            end_date = 0
            stay_num = 0
            # 一局游戏
            # print("GAME#:",gameNum)
            baseline = 0
            tao_prob = []
            tao_reward = []
            wait_day = []#记录一局游戏等待哪些天

            while today < self.routeline[-1] and terminal == False:
                # 有新订单产生 (当订单数已满10个时，此处不会收到新订单)
                if order_accepted == False:
                    self.orderSelect(self.routeline,60)
                    # print(self.order)
                    env.setOrder(self.order)
                    order_accepted = True
                # 遍历self.orders(即state[0])字典，对每一个订单操作
                state = env.getState()

                # 当前状态
                state_tf = np.mat(state)
                # print(state_tf,len(state_tf))
                # 由神经网络选择行动
                if random.random()<epsilon and isExploration == False:
                    isExploration = True
                    end_date = random.randrange(env.getTodayIndex(),87,1)
                    # end_date = 60

                if isExploration:
                    if env.getTodayIndex() == end_date:
                        action_model = 1
                        if ex_steps>0:
                            ex_steps -= 1
                    else:
                        action_model = 0
                else:
                    #action from learning
                    action_model,p = brain.choose_action(state_tf, env.getTodayIndex())
                    tao_prob.append(p)
                if action_model == 0:
                    action_finishOrder = [1, 0]
                else:
                    action_finishOrder = [0, 1]

                wait_day.append(env.getTodayIndex())
                # 订单字典 历史曲线 reward

                reward = env.getReward(action_model)
                tao_reward.append(reward)
                # 订单完成或者到最后一天
                terminal = env.isTerminal(action_model)
                state_ = env.getNextState(action_model)
                # print(state_tf)
                # print(state_)
                td_error = brain.criticLearn(state_tf, reward, state_)
                baseline = td_error
                profitAdvanced_list.append(td_error[0][0])
                loss = brain.actorLearn(state_tf, action_model, td_error)
                # print(loss)
                Loss_list.append(loss)
                # 保存记录到记忆库
                # print("this is store arg:",state_tf,";", action_model,";", reward,";", env.getTodayIndex())
                # brain.store_transition(state_tf, action_model, reward, env.getTodayIndex())
                # print(action_model)

                total_steps += 1
                if terminal:
                    wait_list.append(wait_day[-1])
                    break

                # step 过一天加一
                env.nextStep()

            # 一局的总收益
            # epsilon = self.epsilon*(ex_steps/500)
            print("epsilon:",epsilon)
            print("TD_Error:",baseline)
            profit = env.getTotalReward()
            profit_list.append(profit)
            print("total_steps:",total_steps)
            print("profit_list", profit_list)
            print("profit:", profit, "profitAvg:", np.mean(profit_list))
            print("action-prob:",tao_prob)
            print("Reward:",tao_reward)
            print("wait_day:",wait_day)
            self.writeHistory('./picture/history.txt',
                              epsilon,baseline,total_steps,profit_list,profit,tao_prob,tao_reward,wait_day,gameNum)

            print("########################"+str(gameNum)+"###########################")
            if len(profit_list) >=gameSplit:
                plt.figure()
                plt.plot(profit_list, 'r-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_profit_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), profit_list)
                plt.savefig('./picture/' + str(gameNum) + 'scatter_profit_PG.jpg')
                plt.figure()
                plt.plot(profitAdvanced_list, 'g-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_advanced_PG.jpg')
                plt.figure()
                plt.plot(Loss_list, 'y-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_loss_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), wait_list,c='r')
                plt.savefig('./picture/' + str(gameNum) + 'scatter_waitDay_PG.jpg')
                profit_list.clear()
                wait_list.clear()
            # last_remainder = total_steps % 1000

    def transcate_DDPG(self):
        BATCH_SIZE = 32
        total_steps = 0  # 记录步数，一天是一步
        profit_list = []  # 记录每局总收益
        profitAdvanced_list = []
        actions = 2  # 行动个数
        s_dim = 87
        a_dim = 1
        brain = DDPG(
            a_dim = a_dim,
            s_dim = s_dim,
            a_bound = 1.,
            LR_A = 0.001,
            LR_C = 0.001,
            GAMMA = .99,
            TAU = 0.01,
            # replacement=REPLACEMENT,
        )
        gameNum = 0 #记录游戏轮数
        ex_steps = 500 #探索衰减的轮数
        epsilon = self.epsilon
        last_remainder  = 0
        reward_list = [0] #存储每次的收益，来计算baseline
        Loss_list = [] #存储训练过程中的损失值
        wait_list = [] #记录N轮游戏分别等待天数
        gameSplit = 5000 #每多少轮游戏画图
        while total_steps < 60000:
            # 初始化游戏
            # routeId = random.randrange(0, 49, 1)
            routeId = 21
            self.routeline = self.allRoute[routeId]
            # print(self.routeline)
            env = Env(self.routeline)
            gameNum += 1
            # state = env.getState()  # 以state[0]、state[1]方式访问
            today = env.getToday()

            terminal = False
            order_accepted = False
            isExploration = False
            create_date = 1
            end_date = 0
            stay_num = 0
            # 一局游戏
            # print("GAME#:",gameNum)
            baseline = 0
            tao_prob = []
            tao_reward = []
            wait_day = []#记录一局游戏等待哪些天

            while today < self.routeline[-1] and terminal == False:
                # 有新订单产生 (当订单数已满10个时，此处不会收到新订单)
                if order_accepted == False:
                    self.orderSelect(self.routeline,60)
                    # print(self.order)
                    env.setOrder(self.order)
                    order_accepted = True
                # 遍历self.orders(即state[0])字典，对每一个订单操作
                state = env.getState()

                # 当前状态
                state_tf = np.mat(state)
                # print(state_tf,len(state_tf))
                # 由神经网络选择行动
                if random.random()<epsilon and isExploration == False:
                    isExploration = True
                    # end_date = random.randrange(env.getTodayIndex(),87,1)
                    end_date = 60

                if isExploration:
                    if env.getTodayIndex() == end_date:
                        action_model = 1
                        if ex_steps>0:
                            ex_steps -= 1
                    else:
                        action_model = 0
                else:
                    #action from learning
                    action_model = brain.choose_action(state_tf)
                    # print(action_model)

                wait_day.append(env.getTodayIndex())
                # 订单字典 历史曲线 reward

                reward = env.getReward(action_model)
                tao_reward.append(reward)
                # 订单完成或者到最后一天
                terminal = env.isTerminal(action_model)
                state_ = env.getNextState(action_model)
                if len(state_) == 1:
                    state_ = copy.deepcopy(state)
                brain.store_transition(state, action_model, reward, state_)
                # profitAdvanced_list.append(td_error[0][0])

                if brain.pointer > brain.MEMORY_CAPACITY :
                    # print(b_s_)
                    brain.learn()

                total_steps += 1
                if terminal:
                    # wait_list.append(wait_day[-1])
                    # loss = brain.learn()
                    # Loss_list.append(loss)
                    break

                # step 过一天加一
                env.nextStep()

            # 一局的总收益
            epsilon = self.epsilon*(ex_steps/500)
            print("epsilon:",epsilon)
            print("TD_Error:",baseline)
            profit = env.getTotalReward()
            profit_list.append(profit)
            print("total_steps:",total_steps)
            print("profit_list", profit_list)
            print("profit:", profit, "profitAvg:", np.mean(profit_list))
            print("action-prob:",tao_prob)
            print("Reward:",tao_reward)
            print("wait_day:",wait_day)
            self.writeHistory('./picture/history.txt',
                              epsilon,baseline,total_steps,profit_list,profit,tao_prob,tao_reward,wait_day,gameNum)

            print("########################"+str(gameNum)+"###########################")
            if len(profit_list) >=gameSplit:
                plt.figure()
                plt.plot(profit_list, 'r-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_profit_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), profit_list)
                plt.savefig('./picture/' + str(gameNum) + 'scatter_profit_PG.jpg')
                plt.figure()
                plt.plot(profitAdvanced_list, 'g-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_advanced_PG.jpg')
                plt.figure()
                plt.plot(Loss_list, 'y-')
                plt.savefig('./picture/' + str(gameNum) + 'liner_loss_PG.jpg')
                plt.figure()
                plt.scatter(np.arange(gameSplit), wait_list,c='r')
                plt.savefig('./picture/' + str(gameNum) + 'scatter_waitDay_PG.jpg')
            if len(profit_list) >= 500:
                profit_list.clear()
                wait_list.clear()
            # last_remainder = total_steps % 1000

    def orderGenerator(price_mean, price_cov, time_mean, time_cov, size=1):
        price_axis = np.random.multivariate_normal(mean=price_mean, cov=price_cov, size=size)
        time_axis = np.random.multivariate_normal(mean=time_mean, cov=time_cov, size=size)
        # 需要再组合下两个，形成一个list
        return price_axis, time_axis

    def orderSelect(self,routeline,day_num):
        while True:
            # orderId = random.randrange(0, 59, 1)
            orderId = 0
            self.order = copy.deepcopy(self.allOrder[orderId])
            self.order[1] = '2019-06-20'
            orderDateStr = self.order[1]
            routeDateStr = routeline[-2]
            order_date = datetime.strptime(orderDateStr, "%Y-%m-%d")
            create_date = datetime.strptime(routeDateStr, "%Y-%m-%d")

            # print(orderDateStr,routeDateStr)
            # if orderDateStr < routeDateStr or order_date > create_date+ timedelta(days=40):
            if order_date > create_date + timedelta(days=day_num) or order_date < create_date + timedelta(days=25):
                continue
            else:
                # self.order = copy.deepcopy(self.allOrder[orderId])
                # self.order[1] = '2019-06-27'
                # self.order[1] = '2019-07-09'
                return
#读取订单数据
    def readOrder(self,filename):
        with open(filename, 'r') as f:
            content = f.readlines()
            result = []
            for c in content:
                list = []
                c = c.strip("[]\n\t")
                c = c.split(",")
                list.append(float(c[0]))
                c[1] = c[1].replace("'", "")
                list.append(c[1].replace(" ",""))
                result.append(list)
            # print(result[-1])
            return result
#读取航班数据
    def readRoute(self,filename):
        with open(filename, 'r') as f:
            content = f.readlines()
            result = []
            for c in content:
                list = []
                c = c.strip("[]\n")
                c = c.split(",")
                for i in range(89):
                    if i < 87:
                        list.append(float(c[i]))
                    else:
                        c[i] = c[i].replace("'", "")
                        list.append(c[i].replace(" ",""))
                result.append(list)
            # print(result)
            return result
#存储训练过程
    def writeHistory(self, filename, epsilon, baseline, total_steps, profit_list, profit, tao_prob, tao_reward, wait_day, gameNum):
        f = open(filename,'a')
        f.write("epsilon:"+ str(epsilon) + "\n")
        f.write("Baseline:"+ str(baseline)+"\n")
        f.write("total_steps:"+ str(total_steps)+"\n")
        f.write("profit_list"+ str(profit_list)+"\n")
        f.write("profit:"+ str(profit)+"profitAvg:"+ str(np.mean(profit_list))+"\n")
        f.write("action-prob:"+ str(tao_prob)+"\n")
        f.write("Reward:"+ str(tao_reward)+"\n")
        f.write("wait_day:"+ str(wait_day)+"\n")
        f.write("########################"+str(gameNum)+"###########################\n")
        f.flush()

    def testLSTM(self):
        brain = PolicyGradient(
            n_actions=2,
            n_features=87,
            learning_rate=0.02,
            reward_decay=0.85,
        )
        while True:
            # N = random.randrange(1,87,1)
            N=1
            list = self.random_int_list(1,3000,N)
            while len(list)<87:
                list.append(0)
            # print(np.mat(list)/2000)
            action_model = brain.choose_action(np.mat(list), N)
            print("length:",N,"action:",action_model)

    def random_int_list(self,start, stop, length):
        start, stop = (int(start), int(stop)) if start <= stop else (int(stop), int(start))
        length = int(abs(length)) if length else 0
        random_list = []
        for i in range(length):
            random_list.append(random.randint(start, stop))
        return random_list

    def testRandomPolicy(self,day_num):
        game_num = 0
        profit_list = []
        while game_num < 5000:
            routeId = random.randrange(0, 49, 1)
            self.routeline = self.allRoute[routeId]
            self.orderSelect(self.routeline,60)
            buy_day = 1
            profit_list.append(self.order[0] - self.routeline[day_num-1])
            game_num += 1
        print("####################################")
        print("day_num",day_num)
        # print("Profit_List",profit_list)
        print("AVG_Profit",np.mean(profit_list))

if __name__ == '__main__':
    P = TikcetPlay()
    # for i in range(30,88):
    #     P.testRandomPolicy(i)
    P.transcate_DDPG()
    # P.testLSTM()


#!/usr/bin/env python3
# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.
"""
Run inference for pre-processed data with a trained model.
"""

import logging
import os
import random
import string
import sys
from vad import get_microphone_chunks
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import datetime as dt
from time import time
import datetime as dt

import sentencepiece as spm
import torch
import torchaudio
import numpy as np
from fairseq import options, progress_bar, utils, tasks
from fairseq.meters import StopwatchMeter, TimeMeter
from fairseq.utils import import_user_module

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

def add_asr_eval_argument(parser):
    parser.add_argument("--ctc", action="store_true",
                        help="decode a ctc model")
    parser.add_argument("--rnnt", default=False, help="decode a rnnt model")
    parser.add_argument("--kspmodel", default=None,
                        help="sentence piece model")
    parser.add_argument(
        "--wfstlm", default=None, help="wfstlm on dictonary output units"
    )
    parser.add_argument(
        "--rnnt_decoding_type",
        default="greedy",
        help="wfstlm on dictonary\
output units",
    )
    parser.add_argument(
        "--lm_weight",
        default=0.2,
        help="weight for wfstlm while interpolating\
with neural score",
    )
    parser.add_argument(
        "--rnnt_len_penalty", default=-0.5, help="rnnt length penalty on word level"
    )
    return parser

def check_args(args):
    assert args.path is not None, "--path required for generation!"
    assert (
        not args.sampling or args.nbest == args.beam
    ), "--sampling requires --nbest to be equal to --beam"
    assert (
        args.replace_unk is None or args.raw_text
    ), "--replace-unk requires a raw text dataset (--raw-text)"

def process_predictions(args, hypos, sp, tgt_dict):
    res = []
    for hypo in hypos[: min(len(hypos), args.nbest)]:
        hyp_pieces = tgt_dict.string(hypo["tokens"].int().cpu())
        hyp_words = sp.DecodePieces(hyp_pieces.split())
        res.append(hyp_words)
    return res

def optimize_models(args, use_cuda, models):
    """Optimize ensemble for generation
    """
    for model in models:
        model.make_generation_fast_(
            beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
            need_attn=args.print_alignment,
        )
        if args.fp16:
            model.half()
        if use_cuda:
            model.cuda()

def calc_mean_invstddev(feature):
    if len(feature.shape) != 2:
        raise ValueError("We expect the input feature to be 2-D tensor")
    mean = np.mean(feature, axis=0)
    var = np.var(feature, axis=0)
    # avoid division by ~zero
    if var.any() < sys.float_info.epsilon:
        return mean, 1.0 / (np.sqrt(var) + sys.float_info.epsilon)
    return mean, 1.0 / np.sqrt(var)

def calcMN(features):
    mean, invstddev = calc_mean_invstddev(features)
    res = (features - mean) * invstddev
    return res

def transcribe(waveform, args, task, generator, models, sp, tgt_dict):
    r"""
    CUDA_VISIBLE_DEVICES=0 python infer_asr.py /Users/jamarshon/Documents/downloads/ \
        --task speech_recognition --max-tokens 250000 --nbest 1 --path \
        /Users/jamarshon/Downloads/checkpoint_avg_60_80.pt --beam 20
    """
    num_features = 80
    output = torchaudio.compliance.kaldi.fbank(
        waveform, num_mel_bins=num_features)
    output_cmvn = calcMN(output.cpu().detach().numpy())

    # size (m, n)
    source = torch.tensor(output_cmvn)
    frames_lengths = torch.LongTensor([source.size(0)])

    # size (1, m, n). In general, if source is (x, m, n), then hypos is (x, ...)
    source.unsqueeze_(0)
    sample = {'net_input': {'src_tokens': source, 'src_lengths': frames_lengths}}

    hypos = task.inference_step(generator, models, sample)

    assert len(hypos) == 1
    transcription = []
    for i in range(len(hypos)):
        # Process top predictions
        hyp_words = process_predictions(args, hypos[i], sp, tgt_dict)
        transcription.append(hyp_words)
    return transcription

def main(args):
    check_args(args)
    import_user_module(args)

    if args.max_tokens is None and args.max_sentences is None:
        args.max_tokens = 30000
    logger.info(args)

    use_cuda = torch.cuda.is_available() and not args.cpu

    # Load dataset splits
    task = tasks.setup_task(args)

    # Set dictionary
    tgt_dict = task.target_dictionary

    # Load ensemble
    logger.info("| loading model(s) from {}".format(args.path))
    models, _model_args = utils.load_ensemble_for_inference(
        args.path.split(":"),
        task,
        model_arg_overrides=eval(args.model_overrides),  # noqa
    )
    optimize_models(args, use_cuda, models)

    # Initialize generator
    generator = task.build_generator(args)

    sp = spm.SentencePieceProcessor()
    sp.Load(os.path.join(args.data, 'spm.model'))

    print("READY!")
    for (waveform, sample_rate) in get_microphone_chunks():
        waveform = torchaudio.transforms.Resample(
            orig_freq=sample_rate, new_freq=16000)(waveform.reshape(1, -1))
        transcription = transcribe(waveform, args, task, generator, models, sp, tgt_dict)
        print("{}: {}".format(dt.datetime.now().strftime('%H:%M:%S'), transcription[0][0]))

def cli_main():
    parser = options.get_generation_parser()
    parser = add_asr_eval_argument(parser)
    args = options.parse_args_and_arch(parser)
    main(args)

if __name__ == "__main__":
    cli_main()


# -*- coding: utf-8 -*-
"""
Created on Sat Feb 15 14:15:00 2020
@author: asadl
"""

# -*- coding: utf-8 -*-
"""
Spyder Editor
This is a temporary script file.
"""

import multiprocessing
import numpy as np
import dpkt
import struct
import time
#from copy import deepcopy
#from time import gmtime, strftime
import pandas as pd
import cv2
import collision
import itertools

from Send_Alert import *

''' This block of code is mainly for reading raw packet of data received from a lidar sensor and convert it to arrays of distance measurements or frames, DO NOT edit it '''
#pedestrians
resolution = int(9*100)
radiusToCover = 22.5 #32.5 #3.75 # in meters
distancePixelRatio = resolution / (2 * radiusToCover)

# for vlp-32c
elevations = ((-25, -1, -1.667, -15.639, -11.31, 0, -0.667, -8.843, -7.254, 0.333, -0.333, -6.148, -5.333, 1.333, 0.667, -4, -4.667, 1.667, 1, -3.667, -3.333, 3.333, 2.333, -2.667, -3, 7, 4.667, -2.333, -2, 15, 10.333, -1.333))
sinElevations = np.array(np.sin(np.radians([[-25, -1, -1.667, -15.639, -11.31, 0, -0.667, -8.843, -7.254, 0.333, -0.333, -6.148, -5.333, 1.333, 0.667, -4, -4.667, 1.667, 1, -3.667, -3.333, 3.333, 2.333, -2.667, -3, 7, 4.667, -2.333, -2, 15, 10.333, -1.333]])))
cosElevations = np.array(np.cos(np.radians([[-25, -1, -1.667, -15.639, -11.31, 0, -0.667, -8.843, -7.254, 0.333, -0.333, -6.148, -5.333, 1.333, 0.667, -4, -4.667, 1.667, 1, -3.667, -3.333, 3.333, 2.333, -2.667, -3, 7, 4.667, -2.333, -2, 15, 10.333, -1.333]])))
deltas = np.array(np.radians([[1.4,-4.2,1.4,-1.4,1.4,-1.4,4.2,-1.4,1.4,-4.2,1.4,-1.4,4.2,-1.4,4.2,-1.4,1.4,-4.2,1.4,-4.2,4.2,-1.4,1.4,-1.4,1.4,-1.4,1.4,-4.2,4.2,-1.4,1.4,-1.4]]))
distanceMultiplier = 0.004

class PcapSniffer(multiprocessing.Process):
    def __init__(self,fileName,startTime=0.0,startPacket=0):
        self.fileName = fileName
        self.startTime = startTime
        self.startPacket = startPacket
        self.numberOfPackets = 0
        self.firstPacketTime = 0.0
        self.endPacketTime = 0.0
        self.maxPacketTime = 0.0
        self.hasPassedMaxTime = False
        self.frameOutQueue = multiprocessing.Queue()
        self.controlQueue = multiprocessing.Queue()
        self.measurePcap()
        multiprocessing.Process.__init__(self)
        self.start()

    def run(self):
        # loop to open the pcap after a control input is received
        while(True):
            try:
                self.startTime, self.startPacket = self.controlQueue.get(block=False)
                break
            except:
                pass
            f = open(self.fileName,'rb')
            pcap = dpkt.pcap.Reader(f)

            frameData = []
            lastAzimuth = 0.0
            timeStamp = 0.0

            # loop through packets in pcap
            for _,buf in pcap:
                lastFrameProcessTime = time.time()

                try:
                    self.startTime, self.startPacket = self.controlQueue.get(block=False)
                    if(self.startTime == 0.0 and self.startPacket != 0.0):
                        self.startTime = (timeStamp - (timeStamp %5.0)) + self.startPacket
                    if(self.startTime < timeStamp):
                        break
                except:
                    pass

                if(len(buf) == 1248):
                    data = np.frombuffer(buf, dtype=np.uint8)[42:]
                    if(timeStamp > (struct.unpack_from('I', data[-6:-2])[0] / 1000000)):
                        self.hasPassedMaxTime = True

                    timeStamp = (float(self.hasPassedMaxTime)*self.maxPacketTime) + struct.unpack_from('I', data[-6:-2])[0] / 1000000

                    if(timeStamp > self.startTime - 1.0):
                        # wait for last frame to be taken off the predictionsOutQueue
                        while (not self.frameOutQueue.empty()):
                            time.sleep(0.001)
                            pass

                        azimuth = ((data[3] * 256 + data[2]) * 0.01)

                        frameData.append(data)
                        if (lastAzimuth < 180.0 and azimuth > 180.0):
                            payloadsArray = np.array(frameData)
                            frameData = []
                            payloadsArray = payloadsArray[:, 0:1200]
                            payloadsArray = payloadsArray.reshape((-1, 100))
                            azimuths = np.radians((payloadsArray[:, 3] * 256 + payloadsArray[:, 2]) * 0.01).reshape(
                                (-1, 1)) + 0.0*np.pi

                            payloadsArray = payloadsArray[:, 4:]
                            payloadsArray = payloadsArray.reshape((-1, 3))
                            distances = ((payloadsArray[:, 1] * 256 + payloadsArray[:, 0]) * distanceMultiplier).reshape(
                                (-1, 32))
                            # reflectivities = (payloadsArray[:,2] / 255)

                            deltaAdjustedAzimuths = (
                                        np.repeat(azimuths, distances.shape[1], axis=1) + np.repeat(deltas, distances.shape[0],
                                                                                axis=0)).reshape((-1))
                            cosElevationsRepeated = (np.repeat(cosElevations, distances.shape[0], axis=0)).reshape((-1))
                            # sinElevationsRepeated = (np.repeat(sinElevations, distances.shape[0], axis=0)).reshape((-1))

                            distances = distances.reshape((-1))

                            Xs = (((radiusToCover - distances * cosElevationsRepeated * np.sin(
                                deltaAdjustedAzimuths)) - 0.0) * distancePixelRatio).astype(np.uint16).astype(np.int)
                            Ys = ((distances * cosElevationsRepeated * np.cos(
                                deltaAdjustedAzimuths) - 0.0) * distancePixelRatio).astype(np.uint16).astype(np.int) #0.0 was 5.0
                            # Zs = distances * sinElevationsRepeated

                            validXs = np.where(Xs < resolution)
                            Xs = Xs[validXs]
                            Ys = Ys[validXs]
                            validYs = np.where(Ys < resolution)
                            Xs = Xs[validYs].reshape((-1, 1))
                            Ys = Ys[validYs].reshape((-1, 1))

                            frame = np.zeros(shape=(resolution, resolution, 3), dtype=np.float)
                            frame[Ys, Xs] = np.array((1.0,1.0,1.0)) #(0.5, 0.5, 0.5)) #

                            if (timeStamp > self.startTime):
                                # frame = frame * np.array([1.0, 0.5, 1.0])
                                # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
                                # frame = cv2.dilate(frame, kernel, iterations=1)

                                self.frameOutQueue.put((frame,timeStamp))
                                # print(f'{(time.time() - lastFrameProcessTime):.7f}')
                                # lastFrameProcessTime = time.time()

                        lastAzimuth = azimuth

    def pixelCoordinatesToMeters(self,coordinates):
        coordinates = np.array(coordinates).astype(np.float)
        coordinates[:,0:2] = coordinates[:,0:2]/distancePixelRatio
        return coordinates

    def meterCoordinatesToPixels(self,coordinates):
        coordinates = np.array(coordinates)
        coordinates[:,0:2] = coordinates[:,0:2]*distancePixelRatio
        return coordinates

    def measurePcap(self):
        f = open(self.fileName, 'rb')
        pcap = dpkt.pcap.Reader(f)

        for _, buf in pcap:
            if (len(buf) == 1248):
                data = np.frombuffer(buf, dtype=np.uint8)
                self.firstPacketTime = struct.unpack_from('I', data[-6:-2])[0] / 1000000
                self.numberOfPackets += 1
                break
        for _, buf in pcap:
            if (len(buf) == 1248):
                data = np.frombuffer(buf, dtype=np.uint8)
                if((struct.unpack_from('I', data[-6:-2])[0] / 1000000) < self.endPacketTime and self.maxPacketTime == 0.0):
                    self.maxPacketTime = self.endPacketTime
                self.endPacketTime = self.maxPacketTime + (struct.unpack_from('I', data[-6:-2])[0] / 1000000)
                self.numberOfPackets += 1

''' This bkock of code visualizes the frame data using OpenCV library '''
class ImageViewer:

    def __init__(self, mouseCallback=None, trackBarCallback=None): #, playVideo=False
        cv2.namedWindow("output")
        # if(playVideo):
        #     cv2.namedWindow("video")
        if (trackBarCallback is None):
            cv2.createTrackbar('packetTrackbar','output',0,1000,self.passFunc2)
        else:
            cv2.createTrackbar('packetTrackbar','output',0,1000,trackBarCallback)
        if(mouseCallback is None):
            mouseCallback = self.passFunc
        self.callback = mouseCallback
        cv2.setMouseCallback("output", self.callback)

    def passFunc2(self,position):
        pass

    def passFunc(self,event, x, y, flags, param):
        pass

    def showImage(self, frame, dimensions=(0.0, 0.0), name="output"):
        if (dimensions == (0.0, 0.0)):
            cv2.imshow(name, frame)
        else:
            if(dimensions[0]/frame.shape[0] > dimensions[1]/frame.shape[1]):
                dimensions = (int(frame.shape[1]*(dimensions[0]/frame.shape[0])),int(dimensions[0]))
            else:
                dimensions = (int(dimensions[1]),int(frame.shape[0]*(dimensions[1]/frame.shape[1])))

            frame = cv2.resize(frame, dimensions)
            cv2.imshow(name, frame)

    def destroyWindows(self):
        cv2.destroyAllWindows()

def point_inside_polygon(x,y,poly):
    n = len(poly)
    inside =False
    p1x,p1y = poly[0]
    for i in range(n+1):
        p2x,p2y = poly[i % n]
        if y > min(p1y,p2y):
            if y <= max(p1y,p2y):
                if x <= max(p1x,p2x):
                    if p1y != p2y:
                        xinters = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x
                    if p1x == p2x or x <= xinters:
                        inside = not inside
        p1x,p1y = p2x,p2y
    return inside

''' This is where you can play with the data and run your code or even modify it '''
if __name__ == "__main__":

    '''Some Flags'''
    plot_object=True
    plot_loops=True
    plot_occupancy=True
    plot_tracking=True
    toggle_mode=['car','ped','bike','all']
    toggle_mode_index=4 # it starts with showing all modes
    fps=10.0

    ''' Exisitng file names '''
    file_name='DATA_20200323_154915'
    # file_name='DATA_20200323_160416'
    # file_name='DATA_20200323_161917'
    # file_name='DATA_20200323_163418'
    # file_name='DATA_20200323_164919'
    # file_name='DATA_20200323_170420'
    # file_name='DATA_20200323_171921'
    # file_name='DATA_20200323_173422'
    # file_name='DATA_20200323_174923'
    # file_name='DATA_20200323_180424'

    ''' Collisions Dictionary to store all objects '''
    road_object = {}

    ''' Incident Arrays to store all incidents to prevent from pushing same alerts'''
    incident_arrays = []

    PEDS_RADIUS = 8      # radius for the pedestrians
    CYCL_SIZE = 10       # size of the cyclists
    CAR_SIZE = 20        # size of the cars

    # List of addresses
    #   1. Latitude
    #   2. Longtitude
    #   3. Address
    address = [[49.886282, -119.476950, "Bernard Ave & Gordon Dr, Kelowna BC"],
               [49.886325, -119.482776, "Ethel St & Bernard Ave, Kelowna BC"],
               [49.885388, -119.488528, "Lawrence Ave & Richter St, Kelowna BC"],
               [49.886321, -119.493538, "Bernard Ave & Ellis St, Kelowna BC"],
               [49.887323, -119.496564, "Water St & Queensway, Kelowna BC"],
               [49.893563, -119.488507, "Richter St & Clement Ave, Kelowna BC"],
               [49.893563, -119.482746, "Clement Ave & Ethel St, Kelowna BC"],
               [49.893646, -119.477103, "Clement Ave & Gordon Dr, Kelowna BC"],
               [49.872774, -119.476033, "Guisachan Rd & Gordon Dr, Kelowna BC"],
               [49.865406, -119.482784, "Raymer Ave & Ethel St, Kelowna BC"]]
    address_size = len(address)
    address_index = 0

    pcapSniffer = PcapSniffer('./'+file_name+'.pcap') # read Pcap file NOTE: Modify path based on your computer

    P_car = pd.read_csv(file_name+'_frozenNotTensorRtBigGood.csv',skiprows=1) # Read CSV file with timestamp, Object ID, x and y cordination and mode, 1 for pedestrian, 2 for cars and 3 for cyclists
    P_ped = pd.read_csv(file_name+'_bigPed_pedestrians.csv') # Read CSV file with timestamp, Object ID, x and y cordination and mode, 1 for pedestrian, 2 for cars and 3 for cyclists
    loops = pd.read_csv("loops.csv") # Read the Virtual Loop locations

    loops_location=np.array(loops) # Convert Dataframe to numpy arrays
    P_car=np.array(P_car) # Convert Dataframe to numpy arrays
    P_ped=np.array(P_ped) # Convert Dataframe to numpy arrays
    P=np.concatenate((P_car, P_ped), axis=0) # Add data from ped and car in one array

    ''' Some initialization '''
    object_ID=1
    exit_flag=False
    frame_count=0
    frame_count_total=0
    scale=1
    font = cv2.FONT_HERSHEY_SIMPLEX
    out = cv2.VideoWriter('output.avi', cv2.VideoWriter_fourcc(*"XVID"), fps, (int(900*scale),int(900*scale)), True) # Record output as a video file
    color=[[(255,0,0)],[(0,255,0)],[(0,0,255)]] # Colors assigned for each object based on their mode

    imageViewer = ImageViewer()

    lastFrameProcessTime = time.time()

    P[:,2]=(P[:,2]/0.05) # Converts distance value from Excel to Pixel values. each pixel associated to 5cm in real world.
    P[:,3]=(P[:,3]/0.05) # Converts distance value from Excel to Pixel values. each pixel associated to 5cm in real world.
    P[:,0]=(P[:,0]*10) # time in seconds with flowting numbers (for example 101.1235 seconds) will be converted to 1011.235) based on 10 frame per second sample rate every 100 millisecond we will get a new frame
    P=P.astype(np.int) # Convert everything to integer to make it faster and easier

    while(not exit_flag):
        road_object = {}
        frame, timeStamp = pcapSniffer.frameOutQueue.get()
        timeStamp=int(timeStamp*10) # The same approach as line code 299
        frame=((frame/frame.max())*255).astype('uint8')

        frame_count_total=frame_count_total+1

        frame_count=frame_count+1
        if(frame_count>=10.0/fps):
            frame_count=0

            P_sub=P[P[:,0]==timeStamp,:] # Reads the data about the current frame returned from pcapSniffer

            if toggle_mode_index!=4:
                P_sub=P_sub[P_sub[:,4]==toggle_mode_index,:]

            ''' Now let's plot the info about all objects in the frame '''

#            unique_obj=np.unique(P_sub[:,1])
            for i in range(len(P_sub)):
                cx=P_sub[i,2]
                cy=P_sub[i,3]
                mode=P_sub[i,4]
                Oid=P_sub[i,1]

                if(mode == 2):  # Pedestrians
                    road_object[Oid] = [mode, collision.Circle(collision.Vector(cx, cy), PEDS_RADIUS)]
                    cv2.circle(frame, (cx, cy), PEDS_RADIUS, (255, 0, 0), 2)

                if(mode == 1):  # Cars
                    road_object[Oid] = [mode, collision.Poly.from_box(collision.Vector(cx, cy), CAR_SIZE * 2, CAR_SIZE * 2)]
                    cv2.rectangle(frame, (cx - CAR_SIZE, cy - CAR_SIZE), (cx + CAR_SIZE, cy + CAR_SIZE), (0, 0, 255), 2)

                if(mode == 3):  # Cyclists
                    road_object[Oid] = [mode, collision.Poly.from_box(collision.Vector(cx, cy), CYCL_SIZE * 2, CYCL_SIZE * 2)]
                    cv2.rectangle(frame, (cx - CYCL_SIZE, cy - CYCL_SIZE), (cx + CYCL_SIZE, cy + CYCL_SIZE), (0, 0, 255), 2)

                if(plot_object):
                    cv2.circle(frame,(cx,cy),3,color[mode-1][0],2)
                    cv2.putText(frame,'('+str(cx)+','+str(cy)+')',(cx+10,cy+10),font, 0.4, color[mode-1][0],1)
                    cv2.putText(frame,str(Oid),(cx+30,cy+30),font, 0.4, color[mode-1][0],1)

                '''Plot Tracking'''
                if(plot_tracking):
                    P_sub_obj=P[P[:,1]==Oid,:]
                    P_sub_obj=P_sub_obj[P_sub_obj[:,4]==mode,:]
                    P_sub_obj=P_sub_obj[P_sub_obj[:,0]<=timeStamp,:]
                    P_sub_obj_sorted = (np.array(sorted(P_sub_obj, key=lambda x : x[0], reverse=True))).astype(np.int)
                    if(len(P_sub_obj_sorted)>1):
                        for j in range(1,min(30,len(P_sub_obj_sorted))):
                            cv2.line(frame, (P_sub_obj_sorted[j-1,2],P_sub_obj_sorted[j-1,3]), (P_sub_obj_sorted[j,2],P_sub_obj_sorted[j,3]), color[mode-1][0], 2)

            '''Plot Virtual Loops'''
            if(plot_loops):
                loops=np.copy(loops_location)
                if toggle_mode_index!=4:
                    loops=loops[loops[:,3]==toggle_mode_index,:]

                for l in np.unique(loops[:,0]):

                    pt=loops[loops[:,0]==l,:]
                    pts = np.array([[pt[0,1],pt[0,2]],[pt[1,1],pt[1,2]],[pt[2,1],pt[2,2]],[pt[3,1],pt[3,2]]], np.int32)

                    cv2.polylines(frame,[pts],True,color[int(pt[0,3])-1][0])

                    '''Plot Occupancy'''
                    if(plot_occupancy):
                        for i in range(len(P_sub)):
                            cxo=P_sub[i,2]
                            cyo=P_sub[i,3]
                            mode=P_sub[i,4]
                            cx=(max(pt[:,1])-min(pt[:,1]))/2+min(pt[:,1])
                            cy=(max(pt[:,2])-min(pt[:,2]))/2+min(pt[:,2])

                            if(point_inside_polygon(cxo,cyo,list(pts))):
                                if int(pt[0,3])==mode:
                                    cv2.circle(frame,(int(cx),int(cy)),8,(0,0,255),15)
                    cx=(max(pt[:,1])-min(pt[:,1]))/2+min(pt[:,1])
                    cy=(max(pt[:,2])-min(pt[:,2]))/2+min(pt[:,2])
                    cv2.putText(frame,str(int(l)),(int(cx),int(cy)),font, 0.6, (255,255,255),2)

            object_list = road_object.keys()
            objects_combinations = list(itertools.combinations(object_list, 2))

            for i in objects_combinations:
                if((not (road_object[i[0]][0] == 2 and road_object[i[1]][0] == 2)) and collision.collide(road_object[i[0]][1] , road_object[i[1]][1])):
                    # Check if collisions are already reported or not
                    if(i[0] not in incident_arrays or i[1] not in incident_arrays):
                        incident_arrays.append(i[0])
                        incident_arrays.append(i[1])
                        cv2.putText(frame, 'Warning Collision Detected!', (road_object[i[0]][1].pos.x, road_object[i[0]][1].pos.y), font, 0.6, (0, 0, 255), 2)
                        print('Collision Detected!')
                        print(frame_count_total)
                        # Send alert to Web
                        send_alert(address[address_index], road_object[i[0]][0], road_object[i[1]][0])
                        address_index += 1
                        if(address_index >= address_size):
                            address_index = 0
                        break

            cv2.putText(frame,'Mode : '+toggle_mode[toggle_mode_index-1],(750,50),font, 0.6, (0,0,255),2)
            cv2.putText(frame,str(frame_count_total),(750,20),font, 0.6, (0,0,255),2)
            frame = cv2.resize(frame,None,fx=scale,fy=scale)
            imageViewer.showImage(frame)
            k = cv2.waitKey(100) & 0xFF
            if k == ord('m'): # Change the mode: car, ped,bike,all
                toggle_mode_index=toggle_mode_index+1
                if toggle_mode_index>4:
                    toggle_mode_index=1
            if k == 32: # Pause Playing with space key
                while(True):
                    k1 = cv2.waitKey(10)
                    if k1 == 32:
                        break
            if k == 27: # Exit with ESC key
                    exit_flag=True
                    break

            out.write(frame.astype('uint8'))

            lastFrameProcessTime = time.time()

    imageViewer.destroyWindows()
    out.release()


"""=== main Title =>

    Author : Saifeddine ALOUI
    Description :
        A code to test FaceAnalyzer: takes a still image and use the webcam to track a face and animate the image using the landmarks.
        A reference triangulation file is used. Dont forget to regenerate another one if you change the landmarks you use.
        The image file used in this code is under creative commons licence.
<================"""

import numpy as np
from pathlib import Path
import cv2
import matplotlib.pyplot as plt
import time
from PIL import Image
from FaceAnalyzer import FaceAnalyzer, Face
from pathlib import Path
import pickle

# Parameters
use_simplified_face = True  # If true then the simplified faster version of the face landmarks will be used instead of the full more accurate but slower version
use_seemless_cloning= False #If true, seemless cloning will be used to blend the mask onto the face
view_original = False # If true, the original image with triangles will be shown

# Select landmarks
if use_simplified_face:
    # Simplified facial features (fast)
    lm_indices = list(set(
                        Face.simplified_face_features+
                        Face.mouth_inner_indices+
                        Face.mouth_outer_indices+
                        Face.left_eyelids_indices+
                        Face.right_eyelids_indices+
                        Face.left_eye_contour_indices+
                        Face.right_eye_contour_indices)) # list(range(468)) #
else:
    # Full face (Slower)
    lm_indices = list(range(478))

# open reference triangulation file
file = Path(__file__).parent/"reference.pkl"
with open(str(file),"rb") as f:
    triangles = pickle.load(f)

# open an image and recover all faces inside it (here there is a single face)
fa_mask = FaceAnalyzer.from_image(str(Path(__file__).parent/"assets/pennywize.jpg"))
mask_face = fa_mask.faces[0]
min_ = mask_face.npLandmarks.min(axis=0)-np.array([150,50,0])
max_ = mask_face.npLandmarks.max(axis=0)+np.array([150,50,0])
peripherals = np.array([[min_[0],min_[1],0],[min_[0],max_[1],0],[max_[0],max_[1],0],[max_[0],min_[1],0]])
mask_face.npLandmarks=np.vstack([mask_face.npLandmarks,peripherals])
mask_face.triangles=triangles

cap = cv2.VideoCapture(0)
# cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
# cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

# Build a window
cv2.namedWindow('Face Mask', flags=cv2.WINDOW_NORMAL)
cv2.resizeWindow('Face Mask', (640,480))

# Plot
start_time = time.time_ns()
fa = FaceAnalyzer(max_nb_faces=3)
for face in fa.faces:
    face.triangles = mask_face.triangles

ref_inf =(mask_face.npLandmarks[:-4,...].max(axis=0)-mask_face.npLandmarks[:-4,...].min(axis=0))
while cap.isOpened():
    success, image = cap.read()
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    if view_original:
        image = fa_mask.image.copy()
        mask_face.draw_delaunay(image, landmark_indices=lm_indices)
    else:
        files = []
        fa.process(image)
        img = fa_mask.image.copy()

        if fa.found_faces:
            for i, face in enumerate(fa.faces):
                if face.ready:
                    n_lm= face.npLandmarks.shape[0]
                    extra_indices = [n_lm,n_lm+1,n_lm+2,n_lm+3]
                    mn = face.npLandmarks.mean(axis=0)
                    face.npLandmarks=face.npLandmarks-mn
                    scale=ref_inf/(face.npLandmarks.max(axis=0)-face.npLandmarks.min(axis=0))
                    face.npLandmarks*=scale
                    face.npLandmarks+=mask_face.npLandmarks[::-4,...].mean(axis=0)
                    face.npLandmarks=np.vstack([face.npLandmarks,peripherals])
                    #face.draw_landmarks(img,face.get_landmarks_pos(lm_indices+extra_indices))
                    #face.draw_delaunay(img,lm_indices+extra_indices)
                    image=img
                    mask_face.copyToFace(
                                            face,
                                            fa_mask.image,
                                            img,
                                            opacity = 1,
                                            landmark_indices=lm_indices+extra_indices,
                                            seemless_cloning=use_seemless_cloning,
                                            empty_fill_color=(0,0,0))
                break
    try:
        cv2.imshow('Face Mask', cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    except Exception as ex:
        print(ex)
    # Wait for key stroke
    wk = cv2.waitKey(5)
    if wk & 0xFF == 27: # If escape is pressed then return
        break
    if wk & 0xFF == 115: # If s is pressed then take a snapshot
        sc_dir = Path(__file__).parent/"screenshots"
        if not sc_dir.exists():
            sc_dir.mkdir(exist_ok=True, parents=True)
        i = 1
        file = sc_dir /f"sc_{i}.jpg"
        while file.exists():
            i+=1
            file = sc_dir /f"sc_{i}.jpg"
        cv2.imwrite(str(file),cv2.cvtColor(image,cv2.COLOR_BGR2RGB))
        print("Shot")
cap.release()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep  6 18:37:18 2021

@author: qiang
"""

import sys
import numpy as np
import cv2
from scipy.spatial.transform import Rotation
# import imutils
import itertools
from torchvision.utils import save_image
import torchvision.models as models
import torch
from torchvision import transforms
from torch.nn import functional as F
from rrc_example_package import rearrange_dice_env
import trifinger_simulation.tasks.rearrange_dice as task
from trifinger_object_tracking.py_lightblue_segmenter import segment_image
import torch.nn as nn
from rrc_example_package.scripts import batch_network
from torch.autograd import Variable
from copy import copy

import torch
import torch.nn.functional as F
import os
import numpy as np
import torch.nn.parallel
import torchvision.utils as vutils
from torchvision.utils import save_image
from rrc_example_package.scripts.utils import lineplot
from torchvision import datasets
import cv2
from rrc_example_package import rearrange_dice_env_GAN
from copy import copy
from rrc_example_package.scripts.vae_network import VAE
from rrc_example_package.scripts import ae_network
from torch.autograd import Variable
from torchvision import transforms
from random import randint
import time
import random
import os
from torch.nn import functional as F
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torchvision.utils as vutils
from torch.autograd import Variable
from rrc_example_package.scripts.utils import lineplot

from rrc_example_package.her.mpi_utils.simple_rms import RunningMeanStd

coord_norm = RunningMeanStd(shape=(75))

FACE_CORNERS = (
    (0, 1, 2, 3),
    (4, 5, 1, 0),
    (5, 6, 2, 1),
    (7, 6, 2, 3),
    (4, 7, 3, 0),
    (4, 5, 6, 7),
)

Encoder = torch.load('./Encoder_5000.pt')
Decoder = torch.load('./Decoder_5000.pt')
Dec = torch.load('./Dec_5000.pt')
Inc = torch.load('./Inc_5000.pt')
optimizerEncoder = optim.Adam(Encoder.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)
optimizerDec = optim.Adam(Dec.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)

def preprocess_observation_(observation, bit_depth):
  observation.div_(2 ** (8 - bit_depth)).floor_().div_(2 ** bit_depth).sub_(0.5)  # Quantise to given bit depth and centre
  observation.add_(torch.rand_like(observation).div_(2 ** bit_depth))  # Dequantise (to approx. match likelihood of PDF of continuous images vs. PMF of discrete images)

# Postprocess an observation for storage (from float32 numpy array [-0.5, 0.5] to uint8 numpy array [0, 255])
def postprocess_observation(observation, bit_depth):
  return np.clip(np.floor((observation + 0.5) * 2 ** bit_depth) * 2 ** (8 - bit_depth), 0, 2 ** 8 - 1).astype(np.uint8)

def _images_to_observation(images, bit_depth):
  # print(images)
  images = torch.tensor(cv2.resize(images, (128,128), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1), dtype=torch.float32)  # Resize and put channel first
  # images = cv2.resize(images, (128,128), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)  # Resize and put channel first
   # images = torch.tensor(cv2.resize(images, (64,64), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1), dtype=torch.float32)  # Resize and put channel first
  # preprocess_observation_(images, bit_depth)  # Quantise, centre and dequantise inplace
  return images.unsqueeze(dim=0)  # Add batch dimension
  # return images

gpu_id = 0
device = torch.device("cuda", gpu_id)

def generate_batch(env, batch_size):
    batch = np.ones((batch_size, 384))
    goals = np.ones((batch_size, 25 * 3))
    for i in range(batch_size):
        seg_mask = np.ones((3, 270, 270))
        g = task.sample_goal()
        goal = list(itertools.chain(*g))
		#goal = [g for i, g in enumerate(goal) if ((i+1) % 3) !=0]
        goals[i] = np.array(goal)
        for idx, c in enumerate(env.camera_params):
            seg_mask[idx,:,:] = generate_goal_mask(c, g)
        seg_mask = seg_mask.transpose(1, 2, 0)
        seg_mask =  cv2.resize(seg_mask, (128,128), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)
        seg_mask /= 255.0

        seg_mask = torch.from_numpy(seg_mask).unsqueeze(dim=0)
        seg_mask = Variable(seg_mask).type('torch.cuda.FloatTensor').to(device)

        # k1 = copy(seg_mask)
        seg_mask = Encoder(seg_mask)
        seg_mask = Dec(seg_mask)
        # seg_mask = Inc(seg_mask)
        # seg_mask = Decoder(seg_mask)
        # compare_x = torch.cat([k1, seg_mask])
        # save_image(compare_x.data.cpu(), './batch_image_%03d.png' % (i))

        batch[i] = seg_mask.cpu().detach().numpy()
# 		batch[i] = torch.tensor(cv2.resize(images, (128,128), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1), dtype=torch.float32)  # Resize and put channel first
# 		print(seg_mask.shape)
# 		a = seg_mask.transpose(1,2,0)
# 		cv2.imwrite('%d.jpg' % (i),a)
    return batch, goals

##########################################################################################################################
# def generate_batch(env, batch_size):
#  	batch = np.ones((batch_size, 3, 128, 128))
#  	goals = np.ones((batch_size, 25 * 3))
#  	#goals = np.ones((batch_size, 25 * 2))
#  	for i in range(batch_size):
# 		seg_mask = np.ones((3, 270, 270))
# 		g = task.sample_goal()
# 		goal = list(itertools.chain(*g))
# 		#goal = [g for i, g in enumerate(goal) if ((i+1) % 3) !=0]
# 		goals[i] = np.array(goal)
# 		for idx, c in enumerate(env.camera_params):
#  			seg_mask[idx,:,:] = generate_goal_mask(c, g)
# 		#segmentation_masks = np.array([segment_image(cv2.cvtColor(c.image, cv2.COLOR_RGB2BGR)) for c in obs.cameras])
# 		seg_mask = seg_mask.transpose(1, 2, 0)
# 		seg_mask =  cv2.resize(seg_mask, (128,128), interpolation=cv2.INTER_LINEAR).transpose(2, 0, 1)
# 		seg_mask /= 255.0
# 		batch[i] = seg_mask

#  	return batch, goals
##########################################################################################################################

def get_cell_corners_3d(pos):
    """Get 3d positions of the corners of the cell at the given position."""
    d = 0.022 / 2
    nppos = np.asarray(pos)

    # order of the corners is the same as in the cube model of the
    # trifinger_object_tracking package
    # people.tue.mpg.de/mpi-is-software/robotfingers/docs/trifinger_object_tracking/doc/cube_model.html
    return np.array(
        (
            nppos + (d, -d, d),
            nppos + (d, d, d),
            nppos + (-d, d, d),
            nppos + (-d, -d, d),
            nppos + (d, -d, -d),
            nppos + (d, d, -d),
            nppos + (-d, d, -d),
            nppos + (-d, -d, -d),
        )
    )

def generate_goal_mask(camera_parameters, goal):
    """Generate goal masks that can be used with :func:`evaluate_state`.
    A goal mask is a single-channel image where the areas at which dice are
    supposed to be placed are white and everything else is black.  So it
    corresponds more or less to a segmentation mask where all dice are at the
    goal positions.
    For rendering the mask, :data:`TARGET_WIDTH` is used for the die width to
    add some tolerance.
    Args:
        camera_parameters: List of camera parameters, one per camera.
        goal: The goal die positions.
    Returns:
        List of masks.  The number and order of masks corresponds to the input
        ``camera_parameters``.
    """
    #masks = []
    #for cam in camera_parameters:
    #mask = np.zeros((camera_parameters.image_height, camera_parameters.image_width), dtype=np.uint8)
    mask = np.zeros((270, 270), dtype=np.uint8)

    # get camera position and orientation separately
    tvec = camera_parameters.tf_world_to_camera[:3, 3]
    rmat = camera_parameters.tf_world_to_camera[:3, :3]
    rvec = Rotation.from_matrix(rmat).as_rotvec()

    for pos in goal:
        corners = get_cell_corners_3d(pos)

        # project corner points into the image
        projected_corners, _ = cv2.projectPoints(
            corners,
            rvec,
            tvec,
            camera_parameters.camera_matrix,
            camera_parameters.distortion_coefficients,
        )

        # draw faces in mask
        for face_corner_idx in FACE_CORNERS:
            points = np.array(
                [projected_corners[i] for i in face_corner_idx],
                dtype=np.int32,
            )
            mask = cv2.fillConvexPoly(mask, points, 255)

        #masks.append(mask)

    return mask

def loss_fn(recon_x, x, mu, logvar):
    # BCE = F.binary_cross_entropy(recon_x, x, size_average=False)
    # BCE = F.mse_loss(recon_x, x, size_average=False)
    BCE = F.mse_loss(recon_x, x, reduction='none').sum().mean()
    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD, BCE, KLD

##########################################################################################################################
# env =rearrange_dice_env.SimtoRealRearrangeDiceEnv(enable_cameras=True,visualization=False,step_size=1,image_size=270)
# env.reset()

# # def train_AE(env, batch_size):
# metrics = {'steps': [],'observation_loss': []}

# # Encoder = ae_network.Encoder().to(device)
# # Encoder.apply(ae_network.weights_init)
# # Decoder = ae_network.Decoder().to(device)
# # Decoder.apply(ae_network.weights_init)
# # Dec = ae_network.Dec().to(device)
# # Dec.apply(ae_network.weights_init)
# # Inc = ae_network.Inc().to(device)
# # Inc.apply(ae_network.weights_init)

# Encoder = torch.load('./Encoder_5000.pt')
# Decoder = torch.load('./Decoder_5000.pt')
# Dec = torch.load('./Dec_5000.pt')
# Inc = torch.load('./Inc_5000.pt')

# optimizerEncoder = optim.Adam(Encoder.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)
# optimizerDecoder = optim.Adam(Decoder.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)
# optimizerDec = optim.Adam(Dec.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)
# optimizerInc = optim.Adam(Inc.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)

# bs = 4
# try:
#     os.mkdir('FF3')
# except OSError:
#     pass

# results_dir = './FF3'

# epochs = 5000

# for epoch in range(epochs+1):

#     env.reset()
#     input_batch, goals = generate_batch(env, bs)
#     loss = 0
#     for index, mask in enumerate(input_batch):

#         mask = torch.from_numpy(mask).unsqueeze(dim=0)
#         image = Variable(mask).type('torch.cuda.FloatTensor').to(device)

#         generated = Encoder(image)
#         generated = Dec(generated)
#         generated = Inc(generated)
#         recon_x = Decoder(generated)

#         loss += F.mse_loss(recon_x, image, reduction='none').sum().mean()

#     Encoder.zero_grad()
#     Decoder.zero_grad()
#     Dec.zero_grad()
#     Inc.zero_grad()

#     loss.backward()

#     optimizerEncoder.step()
#     optimizerDecoder.step()
#     optimizerDec.step()
#     optimizerInc.step()

#     to_print = "Epoch[{}/{}] Loss: {:.3f}".format(epoch+1, epochs, loss.item()/bs)
#     print(to_print)

#     metrics['observation_loss'].append(loss.item()/bs)
#     metrics['steps'].append(epoch)
#     lineplot(metrics['steps'][-len(metrics['observation_loss']):], metrics['observation_loss'], 'observation_loss', results_dir)

#     if epoch % 10 == 0:

#         mask,coor = generate_batch(env, 1)
#         mask = mask[0]

#         mask = torch.from_numpy(mask)
#         mask = mask.unsqueeze(dim=0)
#         mask = Variable(mask).type('torch.cuda.FloatTensor').to(device)

#         generated = Encoder(mask)
#         generated = Dec(generated)
#         generated = Inc(generated)
#         recon_x = Decoder(generated)

#         compare_x = torch.cat([mask, recon_x])

#         save_image(compare_x.data.cpu(), './FF2/image_%03d.png' % (epoch))

#     if epoch % 20 ==0:
#         tempresults_dir = './FF2/%s%d' % ('epoch',epoch)
#         try:
#             os.mkdir(tempresults_dir)
#         except OSError:
#             pass
#         torch.save(Encoder,tempresults_dir+'/Encoder_%d.pt' % (epoch))
#         torch.save(Decoder,tempresults_dir+'/Decoder_%d.pt' % (epoch))
#         torch.save(Inc,tempresults_dir+'/Inc_%d.pt' % (epoch))
#         torch.save(Dec,tempresults_dir+'/Dec_%d.pt' % (epoch))

##########################################################################################################################
# bs = 16
# env =rearrange_dice_env.SimtoRealRearrangeDiceEnv(enable_cameras=True,visualization=False,step_size=1,image_size=270)
# env.reset()
# input_batch, goals = generate_batch(env, bs)
# print(input_batch.shape)
# print(goals.shape)

metrics = {'steps': [],'observation_loss': []}

Coor = batch_network.Coor().to(device)
Coor.apply(batch_network.weights_init)

optimizerCoor = optim.Adam(Coor.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)
# optimizerCoor = optim.SGD(Coor.parameters(), lr = 0.000001, momentum=0.6)

try:
    os.mkdir('Coor2')
except OSError:
    pass

results_dir = './Coor2'

epochs = 500000
bs = 32

env =rearrange_dice_env.SimtoRealRearrangeDiceEnv(enable_cameras=True,visualization=False,step_size=1,image_size=270)

coord_norm = RunningMeanStd(shape=(75))
coord_norm1 = RunningMeanStd(shape=(384))

for epoch in range(epochs+1):

    env.reset()
    input_batch, goals = generate_batch(env, bs)
    loss = 0
    for index, mask in enumerate(input_batch):

        mask = np.array(mask)
        coord_norm1.update(mask)
        mask = coord_norm1.normalize(mask)
        mask = torch.from_numpy(mask)
        image = Variable(mask).type('torch.cuda.FloatTensor').to(device)
        generated = Coor(image)

        goal = np.array(goals[index])
        coord_norm.update(goal)
        goal = coord_norm.normalize(goal)
        # goal = goal * 100
        goal = torch.from_numpy(goal)
        goal = Variable(goal).type('torch.cuda.FloatTensor').to(device)

        loss += F.mse_loss(generated, goal, reduction='none').sum().mean()

    Coor.zero_grad()

    loss.backward()

    optimizerCoor.step()

    to_print = "Epoch[{}/{}] Loss: {:.3f}".format(epoch+1, epochs, loss.item()/bs)
    print(to_print)

    metrics['observation_loss'].append(loss.item()/bs)
    metrics['steps'].append(epoch)
    lineplot(metrics['steps'][-len(metrics['observation_loss']):], metrics['observation_loss'], 'observation_loss', results_dir)

    if epoch % 30 == 0:
        mask,coor = generate_batch(env, 1)
        mask = mask[0]

        mask = np.array(mask)
        coord_norm1.update(mask)
        mask = coord_norm1.normalize(mask)
        mask = torch.from_numpy(mask)
        mask = mask.unsqueeze(dim=0)
        mask = Variable(mask).type('torch.cuda.FloatTensor').to(device)
        generated = Coor(image)

        coor = coor[0]
        # coor = coor * 100
        coor = torch.from_numpy(coord_norm.normalize(coor))
        coor = Variable(coor).type('torch.cuda.FloatTensor').to(device)

        print('coor',coor.data)
        print('generated',generated.data)

    if epoch % 500 ==0:
        tempresults_dir = './Coor2/%s%d' % ('epoch',epoch)
        try:
            os.mkdir(tempresults_dir)
        except OSError:
            pass

        torch.save(Coor,tempresults_dir+'/Coor2_%d.pt' % (epoch))

##########################################################################################################################
# import torch.optim as optim
# import os
# import torch
# import torch.nn.functional as F
# import os
# import numpy as np
# import torch.nn.parallel
# import torchvision.utils as vutils
# from torchvision.utils import save_image
# from rrc_example_package.scripts.utils import lineplot
# from torchvision import datasets
# import cv2
# from rrc_example_package import rearrange_dice_env_GAN
# from copy import copy
# from rrc_example_package.scripts.vae_network import VAE
# from rrc_example_package.scripts import ae_network
# from torch.autograd import Variable
# from torchvision import transforms
# from random import randint
# import time
# import random
# import os
# from torch.nn import functional as F
# import numpy as np
# import torch
# import torch.nn as nn
# import torch.nn.parallel
# import torch.optim as optim
# import torchvision.utils as vutils
# from torch.autograd import Variable
# from rrc_example_package.scripts.utils import lineplot
# from rrc_example_package.her.mpi_utils.simple_rms import RunningMeanStd

# metrics = {'steps': [],'observation_loss': []}

# Decoder = batch_network.Decoder().to(device)
# Decoder.apply(batch_network.weights_init)

# optimizerDecoder = optim.Adam(Decoder.parameters(), lr = 0.00015, betas = (0.5, 0.999), eps = 1e-8)

# try:
#     os.mkdir('DEC_V3')
# except OSError:
#     pass

# results_dir = './DEC_V3'

# epochs = 500000
# bs = 128

# env =rearrange_dice_env.SimtoRealRearrangeDiceEnv(enable_cameras=True,visualization=False,step_size=1,image_size=270)

# coord_norm = RunningMeanStd(shape=(75))

# for epoch in range(epochs+1):

#     env.reset()
#     input_batch, goals = generate_batch(env, bs)
#     loss = 0
#     for index, mask in enumerate(input_batch):
#         mask/=255.0
#         mask = torch.from_numpy(mask)
#         image = Variable(mask).type('torch.cuda.FloatTensor').to(device)

#         goal = np.array(goals[index])
#         coord_norm.update(goal)
#         goal = coord_norm.normalize(goal)
#         goal = torch.from_numpy(goal)
#         goal = Variable(goal).type('torch.cuda.FloatTensor').to(device)
#         generated = Decoder(goal)

#         loss += F.mse_loss(generated, image, reduction='none').sum().mean()

#     Decoder.zero_grad()

#     loss.backward()

#     optimizerDecoder.step()

#     to_print = "Epoch[{}/{}] Loss: {:.3f}".format(epoch+1, epochs, loss.item()/bs)
#     print(to_print)

#     metrics['observation_loss'].append(loss.item()/bs)
#     metrics['steps'].append(epoch)
#     lineplot(metrics['steps'][-len(metrics['observation_loss']):], metrics['observation_loss'], 'observation_loss', results_dir)

#     if epoch % 30 == 0:
#         mask,coor = generate_batch(env, 1)
#         mask = mask[0]
#         mask/=255.0

#         mask = torch.from_numpy(mask)
#         mask = mask.unsqueeze(dim=0)
#         mask = Variable(mask).type('torch.cuda.FloatTensor').to(device)

#         coor = coor[0]
#         coor = torch.from_numpy(coord_norm.normalize(coor))
#         coor = Variable(coor).type('torch.cuda.FloatTensor').to(device)

#         recon_x = Decoder(coor)

#         compare_x = torch.cat([mask, recon_x])

#         save_image(compare_x.data.cpu(), './DEC_V3/image_%03d.png' % (epoch))

#     if epoch % 2000 ==0:
#         tempresults_dir = './DEC_V3/%s%d' % ('epoch',epoch)
#         try:
#             os.mkdir(tempresults_dir)
#         except OSError:
#             pass

#         torch.save(Decoder,tempresults_dir+'/Decoder_%d.pt' % (epoch))

# optim = torch.optim.Adam(filter(lambda p: p.requires_grad, resnet.parameters()), lr=0.000001)
# min_cost = 10000
# while True:

#     input_batch, goals = generate_batch(env, 16)
#     input_batch = torch.from_numpy(input_batch).float()
#     goals = torch.from_numpy(goals).float()
#     loss = torch.nn.MSELoss()
#     if torch.cuda.is_available():
#         input_batch = input_batch.to('cuda')
#         goals = goals.to('cuda')
#         resnet.to('cuda')

#     """preprocess = transforms.Compose([
#     transforms.Resize(256),
#     transforms.CenterCrop(224),
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])"""

#     #input_batch = preprocess(input_batch)
#     input_batch = torch.nn.functional.normalize(input_batch)
#     out = resnet(input_batch)
#     cost = loss(out, goals)
#     cost.backward()
#     optim.step()
#     print("Loss: {}".format(cost))
#     if cost < min_cost:
#         min_cost = cost
#         torch.save(resnet.state_dict(), './best_model_resnet_2.pth')

# loss = nn.MSELoss()


#!/usr/bin/env python
# coding: utf-8

# In[ ]:

from tensorflow.keras.preprocessing.image import load_img, img_to_array
from moviepy.editor import VideoFileClip, ImageSequenceClip
from tensorflow.keras.regularizers import l2
from tensorflow.keras.layers import *
import matplotlib.pyplot as plt
from tqdm import tqdm
import tensorflow as tf
from glob import glob
import pandas as pd
import numpy as np
import json
import cv2
import os

# In[2]:

H, W = 608, 608
grid_size = [H//32, W//32]
nboxes = 9
classes = ['bike', 'bus', 'car', 'motor', 'person', 'rider', 'traffic light', 'traffic sign', 'train', 'truck']
class_map = {k:idx for idx,k in enumerate(classes)}
nclasses = len(class_map)
output_shape = grid_size + [nboxes*5 + nclasses]

# In[3]:

def compute_iou(boxes1, boxes2):
    boxes1_t = tf.stack([boxes1[..., 0] - boxes1[..., 2] / 2.0,
                         boxes1[..., 1] - boxes1[..., 3] / 2.0,
                         boxes1[..., 0] + boxes1[..., 2] / 2.0,
                         boxes1[..., 1] + boxes1[..., 3] / 2.0],
                        axis=-1)

    boxes2_t = tf.stack([boxes2[..., 0] - boxes2[..., 2] / 2.0,
                         boxes2[..., 1] - boxes2[..., 3] / 2.0,
                         boxes2[..., 0] + boxes2[..., 2] / 2.0,
                         boxes2[..., 1] + boxes2[..., 3] / 2.0],
                        axis=-1)

    # calculate the left up point & right down point
    lu = tf.maximum(boxes1_t[..., :2], boxes2_t[..., :2])
    rd = tf.minimum(boxes1_t[..., 2:], boxes2_t[..., 2:])

    # intersection
    intersection = tf.maximum(0.0, rd - lu)
    inter_square = intersection[..., 0] * intersection[..., 1]

    # calculate the boxs1 square and boxs2 square
    square1 = boxes1[..., 2] * boxes1[..., 3]
    square2 = boxes2[..., 2] * boxes2[..., 3]

    union_square = tf.maximum(square1 + square2 - inter_square, 1e-10)

    return tf.clip_by_value(inter_square / union_square, 0.0, 1.0)

offset = []
offset_tran = []
for i in range(grid_size[0]):
    row = []
    row_trans = []
    for j in range(grid_size[0]):
        row.append(j)
        row_trans.append(i)
    offset.append(row)
    offset_tran.append(row_trans)
offset = np.tile(np.array(offset)[None, :, :, None], reps=[1,1,1,nboxes])
offset_tran = np.tile(np.array(offset_tran)[None, :, :, None], reps=[1,1,1,nboxes])

offset = tf.constant(offset, dtype=tf.float32)
offset_tran = tf.constant(offset_tran, dtype=tf.float32)

def Yolo_Loss(y_true=None, y_pred=None, eval=False):
    pred_obj_conf = y_pred[:,:,:,:nboxes]
    pred_box_classes = y_pred[:,:,:,5*nboxes:]
    pred_box_offset_coord = y_pred[:,:,:, nboxes:5*nboxes]
    pred_box_offset_coord = tf.reshape(pred_box_offset_coord, shape=[-1, grid_size[0], grid_size[0], nboxes, 4])
    pred_box_normalized_coord = tf.stack([(pred_box_offset_coord[:,:,:,:,0] + offset)/grid_size[0],
                                         (pred_box_offset_coord[:,:,:,:,1] + offset_tran)/grid_size[0],
                                         tf.square(pred_box_offset_coord[:,:,:,:,2]),
                                         tf.square(pred_box_offset_coord[:,:,:,:,3])], axis=-1)
    if eval:
        return pred_obj_conf, pred_box_classes, pred_box_normalized_coord
    target_obj_conf = y_true[:,:,:,:1]
    target_box_classes = y_true[:,:,:,5:]
    target_box_coord = y_true[:,:,:,1:5]
    target_box_coord = tf.reshape(target_box_coord, shape=[-1, grid_size[0], grid_size[1], 1, 4])
    target_box_coord = tf.tile(target_box_coord, multiples=[1,1,1,nboxes,1])
    target_box_normalized_coord = target_box_coord / H
    target_box_offset_coord = tf.stack([target_box_normalized_coord[:,:,:,:,0]*grid_size[0] - offset,
                                        target_box_normalized_coord[:,:,:,:,1]*grid_size[0] - offset_tran,
                                        tf.sqrt(target_box_normalized_coord[:,:,:,:,2]),
                                        tf.sqrt(target_box_normalized_coord[:,:,:,:,3])], axis=-1)

    pred_ious = compute_iou(target_box_normalized_coord, pred_box_normalized_coord)
    predictor_mask = tf.reduce_max(pred_ious, axis=3, keepdims=True)
    predictor_mask = tf.cast(pred_ious>=predictor_mask, tf.float32) * target_obj_conf
    noobj_mask = tf.ones_like(predictor_mask) - predictor_mask

    # Computing the class loss
    class_loss = tf.reduce_mean(tf.reduce_sum(tf.square(target_obj_conf*(target_box_classes - pred_box_classes)), axis=[1, 2, 3]))

    # computing the confidence loss
    obj_loss = tf.reduce_mean(tf.reduce_sum(tf.square(predictor_mask*(pred_obj_conf - pred_ious)), axis=[1, 2, 3]))
    noobj_loss = tf.reduce_mean(tf.reduce_sum(tf.square(noobj_mask*(pred_obj_conf)), axis=[1, 2, 3]))

    # computing the localization loss
    predictor_mask = predictor_mask[:,:,:,:,None]
    loc_loss = tf.reduce_mean(tf.reduce_sum(tf.square(predictor_mask*(target_box_offset_coord - pred_box_offset_coord)), axis=[1, 2, 3]))

    loss = 10 * loc_loss + obj_loss + 0.1 * noobj_loss + 0.5 * class_loss
    return loss

# In[4]:

train_images = sorted(glob('/home/mia/backup/research/autonomous_driving/BDD/bdd100k/images/100k/train/*'))
train_labels = sorted(glob('/home/mia/backup/research/autonomous_driving/BDD/bdd100k/labels/100k/train/*'))

val_images = sorted(glob('/home/mia/backup/research/autonomous_driving/BDD/bdd100k/images/100k/val/*'))
val_labels = sorted(glob('/home/mia/backup/research/autonomous_driving/BDD/bdd100k/labels/100k/val/*'))

batch_size = 24
train_steps = len(train_images) // batch_size
val_steps = len(val_images) // batch_size

# In[5]:

def get_label(label_path, orig_h=720, orig_w=1280):
    label = np.zeros(shape=[*grid_size, 5+nclasses])
    with open(label_path, 'r') as f:
        temp = json.load(f)
    for obj in temp['frames'][0]['objects']:
        if 'box2d' in obj:
            x1 = obj['box2d']['x1'] * W / orig_w
            y1 = obj['box2d']['y1'] * H / orig_h
            x2 = obj['box2d']['x2'] * W / orig_w
            y2 = obj['box2d']['y2'] * H / orig_h
            x = (x2 + x1) / 2
            y = (y2 + y1) / 2
            w = x2 - x1
            h = y2 - y1
            category_id = class_map[obj['category']]
            class_vector = np.zeros((nclasses, ))
            class_vector[category_id] = 1
            label_vector = [1, x, y, w, h, *class_vector]
            grid_x = int(x / W * grid_size[1])
            grid_y = int(y / H * grid_size[0])
            try:
                label[grid_y, grid_x] = label_vector
            except:
                continue
    return label

# In[70]:

train_label_vectors = np.zeros(shape=[len(train_images), *grid_size, 5+nclasses])
for i, img in tqdm(enumerate(train_images)):
    fname = img.split('/')[-1].split('.')[0] + '.json'
    label_path = '/home/mia/backup/research/autonomous_driving/BDD/bdd100k/labels/100k/train/' + fname
    train_label_vectors[i] = get_label(label_path)

val_label_vectors = np.zeros(shape=[len(val_images), *grid_size, 5+nclasses])
for i, img in tqdm(enumerate(val_images)):
    fname = img.split('/')[-1].split('.')[0] + '.json'
    label_path = '/home/mia/backup/research/autonomous_driving/BDD/bdd100k/labels/100k/val/' + fname
    val_label_vectors[i] = get_label(label_path)

# In[5]:

def get_image(image_path):
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img)
    img = tf.cast(tf.image.resize(img, size=[H, W]), dtype=tf.float32)
    img = tf.image.random_brightness(img, max_delta=50.)
    img = tf.image.random_saturation(img, lower=0.5, upper=1.5)
    img = tf.image.random_hue(img, max_delta=0.2)
    img = tf.image.random_contrast(img, lower=0.5, upper=1.5)
    img = tf.clip_by_value(img, 0, 255)
    img /= 127.5
    img -= 1.
    return img

def load_data(image_path, labels):
    flip = tf.cast(tf.random.uniform(shape=[1,], minval=0, maxval=2, dtype=tf.int32), dtype=tf.bool).numpy()[0]
    return get_image(image_path), flip_labels(labels)

# In[5]:

def conv_block(x, n_filters, size, pool=False):
    x = Conv2D(filters=n_filters,
               kernel_size=size,
               padding='same',
               kernel_initializer='he_normal')(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(alpha=0.1)(x)
    if pool:
        x = MaxPool2D(pool_size=2)(x)
    return x

input_layer = Input(shape=(H, W, 3))
x = conv_block(input_layer, 32, 3, pool=True)
x = conv_block(x, 64, 3, pool=True)
x = conv_block(x, 128, 3)
x = conv_block(x, 64, 1)
x = conv_block(x, 128, 3, pool=True)
x = conv_block(x, 256, 3)
x = conv_block(x, 128, 1)
x = conv_block(x, 256, 3, pool=True)
x = conv_block(x, 512, 3)
x = conv_block(x, 256, 1)
x = conv_block(x, 512, 3)
x = conv_block(x, 256, 1)

skip = Lambda(lambda tensor : tf.nn.space_to_depth(tensor, block_size=2))(x)

x = conv_block(x, 512, 3, pool=True)
x = conv_block(x, 1024, 3)
x = conv_block(x, 512, 1)
x = conv_block(x, 1024, 3)
x = conv_block(x, 512, 1)
x = conv_block(x, 1024, 3)
x = conv_block(x, 1024, 3)
x = conv_block(x, 1024, 3)
x = concatenate([x, skip])
x = conv_block(x, 1024, 3)

output_layer = Conv2D(output_shape[-1], kernel_size=1)(x)
model = tf.keras.Model(inputs=input_layer, outputs=output_layer, name='Yolo')
model.load_weights('model/weights.h5')

# In[8]:

train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_label_vectors))
train_dataset = train_dataset.shuffle(buffer_size=256)
train_dataset = train_dataset.apply(tf.data.experimental.map_and_batch(map_func=load_data,
                                                   batch_size=batch_size,
                                                   num_parallel_calls=64,
                                                   drop_remainder=True))
train_dataset = train_dataset.repeat()
train_dataset = train_dataset.prefetch(-1)

val_dataset = tf.data.Dataset.from_tensor_slices((val_images, val_label_vectors))
val_dataset = val_dataset.shuffle(buffer_size=256)
val_dataset = val_dataset.apply(tf.data.experimental.map_and_batch(map_func=load_data,
                                                   batch_size=batch_size,
                                                   num_parallel_calls=64,
                                                   drop_remainder=True))
val_dataset = val_dataset.repeat()
val_dataset = val_dataset.prefetch(-1)

# In[7]:

# !rm -r model/*
# callbacks = [tf.keras.callbacks.ModelCheckpoint('model/weights.h5', save_best_only=True, save_weights_only=True)]
# model.fit(train_dataset,
#          steps_per_epoch=train_steps,
#          epochs=300,
#          validation_data=val_dataset,
#          validation_steps=val_steps,
#          callbacks=callbacks)

# In[75]:

label_map = {v:k for k,v in class_map.items()}

def read_img(image_path):
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img)
    img = tf.cast(tf.image.resize(img, size=[H, W]), dtype=tf.float32)
    return img.numpy()

def visualize(img, pred, iou=0.5, score_threshold=0.5, label_map={}, orig_h=720, orig_w=1280, resize=True, video=True):
    predict_object, predict_class, predict_normalized_box  = Yolo_Loss(pred, pred, eval=True)
    predict_object, predict_class, predict_normalized_box = predict_object[0], predict_class[0], predict_normalized_box[0].numpy()
    if not resize:
        predict_normalized_box = predict_normalized_box * H
    else:
        img = cv2.resize(img, (orig_w, orig_h))
        predict_normalized_box[...,0] = predict_normalized_box[...,0] * orig_w
        predict_normalized_box[...,1] = predict_normalized_box[...,1] * orig_h
        predict_normalized_box[...,2] = predict_normalized_box[...,2] * orig_w
        predict_normalized_box[...,3] = predict_normalized_box[...,3] * orig_h
        H, W = orig_h, orig_w
    predict_normalized_box = np.stack([predict_normalized_box[:,:,:,0]-0.5*predict_normalized_box[:,:,:,2],
                                       predict_normalized_box[:,:,:,1]-0.5*predict_normalized_box[:,:,:,3],
                                       predict_normalized_box[:,:,:,0]+0.5*predict_normalized_box[:,:,:,2],
                                       predict_normalized_box[:,:,:,1]+0.5*predict_normalized_box[:,:,:,3]], axis=-1)
    predict_object = np.reshape(predict_object, newshape=[-1, 1])
    predict_class = np.reshape(predict_class, newshape=[-1, 1, nclasses])
    predict_class = np.tile(predict_class, reps=[1, nboxes, 1])
    predict_class = np.reshape(predict_class, newshape=[-1, nclasses])
    predict_class_conf = np.max(predict_class, axis=-1)[:, None]
    predict_class_idx = np.argmax(predict_class, axis=-1)[:, None]
    scores = np.max(predict_class_conf * predict_object, axis=-1)
    predict_normalized_box = np.reshape(predict_normalized_box, newshape=[-1, 4])
    indices = tf.image.non_max_suppression(predict_normalized_box, scores=scores, iou_threshold=iou, score_threshold=score_threshold, max_output_size=100).numpy()

    boxes = np.int32(predict_normalized_box[indices])
    probs = scores[indices]
    category = np.reshape(predict_class_idx[indices], newshape=[-1])

    img = np.uint8(img)
    for idx, box in enumerate(boxes):
        color_map = [(161, 244, 171), (179, 165, 215), (0, 0, 255), (144, 158, 187), (255, 0, 150), (239, 100, 124), (188, 128, 221), (156, 171, 161), (131, 180, 204), (187, 208, 123)]
        bcolor = color_map[category[idx]]
#         label_w = 150
#         label_h = 45

#         cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), bcolor, 8)
        cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), bcolor, 2)

        cv2.rectangle(img, (box[0], box[1]), (box[0]+W//10, box[1]-H//32), (0, 255, 0), -1)
#         cv2.rectangle(img, (box[0], box[1]), (box[0]+label_w, box[1]-label_h), (0, 255, 0), -1)

#         cv2.putText(img, f'{label_map[category[idx]]}', (box[0]+5, box[1]-3), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (10,10,10), 4, lineType=cv2.LINE_8)
        cv2.putText(img, f'{label_map[category[idx]]}', (box[0]+5, box[1]-3), cv2.FONT_HERSHEY_SIMPLEX, .7, (10,10,10), 2, lineType=cv2.LINE_AA)

    if video:
        return img
    plt.figure(figsize=(16, 9))
    plt.imshow(img)
    plt.title('output')
    plt.axis('off')
    plt.savefig(f'outputs/{np.random.randint(0, 10000)}.png')
    plt.close()

def pipeline(frame):
    h,w = frame.shape[:2]
    disp = frame.copy()
    frame = np.float32(cv2.resize(frame, (W,H)))
    frame /= 255.
    pred = model.predict(frame[None, ...])
    output = visualize(disp, pred, iou=0.5, score_threshold=0.1, label_map=label_map, orig_h=h, orig_w=w)
    return output

# In[71]:

# color_map = {i:tuple(v) for i, v in enumerate(np.random.randint(100, 255, size=[10, 3]))}
# color_map = list(color_map.values())

# In[76]:

for i in tqdm(np.random.randint(0, 10000, 500)):
    img = val_images[i]
    img = read_img(img)
    disp = img.copy()
    img /= 255.
    pred = model.predict(img[None, ...])
    x = visualize(disp, pred, iou=0.5, score_threshold=0.4, label_map=label_map, video=False)

# In[62]:

clip = VideoFileClip('../../4k_driving/toronto.webm')
res = clip.fl_image(pipeline)
res.write_videofile('darknet_toronto.mp4', audio=False, threads=16)

# In[ ]:

clip = VideoFileClip('../../4k_driving/vancouver.webm')
res = clip.fl_image(pipeline)
res.write_videofile('darknet_vancouver.mp4', audio=False, threads=16)

# In[ ]:

clip = VideoFileClip('../../4k_driving/seattle.webm')
res = clip.fl_image(pipeline)
res.write_videofile('darknet_seattle.mp4', audio=False, threads=16)

# In[ ]:

clip = VideoFileClip('../../4k_driving/karol.mp4')
res = clip.fl_image(pipeline)
res.write_videofile('darknet_karol.mp4', audio=False, threads=16)

# In[ ]:


import cv2
import sys
import numpy as np
from model import EMR
import json

# prevents opencl usage and unnecessary logging messages
cv2.ocl.setUseOpenCL(False)

EMOTIONS = ['angry', 'happy', 'neutral', 'sad', 'scared']

data = {}
data['emotion'] = []

def format_image(image):
    """
    Function to format frame
    """
    if len(image.shape) > 2 and image.shape[2] == 3:
            # determine whether the image is color
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        # Image read from buffer
        image = cv2.imdecode(image, cv2.CV_LOAD_IMAGE_GRAYSCALE)

    cascade = cv2.CascadeClassifier(
        "./utility/emotionDetection/haarcascade_frontalface_default.xml")
    # faces = cascade.detectMultiScale(image, scaleFactor=1.3, minNeighbors=5)
    faces = cascade.detectMultiScale(image, scaleFactor=1.1, minNeighbors=3)

    if not len(faces) > 0:
        return [0]

    # initialize the first face as having maximum area, then find the one with max_area
    max_area_face = faces[0]
    for face in faces:
        if face[2] * face[3] > max_area_face[2] * max_area_face[3]:
            max_area_face = face
    face = max_area_face

    # extract ROI of face
    image = image[face[1]:(face[1] + face[2]), face[0]:(face[0] + face[3])]

    try:
        # resize the image so that it can be passed to the neural network
        image = cv2.resize(
            image, (48, 48), interpolation=cv2.INTER_CUBIC) / 255.
    except Exception:
        print("----->Problem during resize")
        return None

#     return image
    output = np.array(image)
    output = np.reshape(output, (48, 48, 1))

    return output

# Initialize object of EMR class
network = EMR()
network.build_network()

# name = input('Input videoname (abc.xyz) : ')
try:
    name = sys.argv[1]
    print(name)
except:
    print('please fill the videoname in input_video')
    exit()

cap = cv2.VideoCapture(name)

length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
frame_number = 0
fps = cap.get(cv2.CAP_PROP_FPS)
duration = length/fps

sum = [0, 0, 0, 0, 0]

font = cv2.FONT_HERSHEY_SIMPLEX
feelings_faces = []

frame_width = int(cap.get(3))
frame_height = int(cap.get(4))

# Output file
out = cv2.VideoWriter('./result/video/output.avi', cv2.VideoWriter_fourcc(
    'M', 'J', 'P', 'G'), 10, (frame_width, frame_height))

# append the list with the emoji images
for index, emotion in enumerate(EMOTIONS):
    feelings_faces.append(cv2.imread('./emojis/' + emotion + '.png', -1))

while 1:
    ret, frame = cap.read()
    frame_number += 1

    # Quit when the input video file ends
    if not ret:
        break

    facecasc = cv2.CascadeClassifier(
        './utility/emotionDetection/haarcascade_frontalface_default.xml')
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = facecasc.detectMultiScale(gray, 1.3, 5)
    new_frame = format_image(frame)

    # compute softmax probabilities
    if(len(new_frame) == 48):
        try:
            result = network.predict([new_frame])
        except:
            continue
    else:
        result = None

    # # compute softmax probabilities
    # if(len(new_frame) == 48):
    #     try:
    #         result = network.predict([format_image(frame)])
    #     except:
    #         continue
    # else:
    #     result = None

    if result is not None:
        # write the different emotions and have a bar to indicate probabilities for each class
        for index, emotion in enumerate(EMOTIONS):
            cv2.putText(frame, emotion, (10, index * 20 + 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            cv2.rectangle(frame, (130, index * 20 + 10), (130 +
                                                          int(result[0][index] * 100), (index + 1) * 20 + 4), (255, 0, 0), -1)

        # find the emotion with maximum probability and display it

        maxindex = np.argmax(result[0])
        if(result[0][maxindex] > float(sys.argv[2]) ):
            # print('hello')
            sum[maxindex] += 1
            # print(sum)

        if(frame_number % (int(fps/4)) == 0):
            maxindex2 = np.argmax(sum)
            sum = [0, 0, 0, 0, 0]
            print('time :{:.2f} emotion {}'.format(
                frame_number/fps, EMOTIONS[maxindex2]))

            # Add result to json file
            data['emotion'].append({
                'emotion': EMOTIONS[maxindex2],
                'time': str(frame_number/fps),
                'score': str(result[0][maxindex2])
            })

        font = cv2.FONT_HERSHEY_SIMPLEX
        cv2.putText(frame, EMOTIONS[maxindex], (10, 360),
                    font, 2, (255, 255, 255), 2, cv2.LINE_AA)
        face_image = feelings_faces[maxindex]

    if len(faces) > 0:
        # draw box around faces
        for face in faces:
            (x, y, w, h) = face
            frame = cv2.rectangle(
                frame, (x, y-30), (x+w, y+h+10), (255, 0, 0), 2)
            newimg = frame[y:y+h, x:x+w]
            newimg = cv2.resize(
                newimg, (48, 48), interpolation=cv2.INTER_CUBIC) / 255.
#             result = model.predict(newimg)
    out.write(frame)

    # cv2.imshow('Video', cv2.resize(frame, None, fx=2,
    #                                fy=2, interpolation=cv2.INTER_CUBIC))
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

with open('./result/JSON/data.json', 'w') as outfile:
    json.dump(data, outfile)