import os
import cv2
from glob import glob
import numpy as np
from utils.dataprocessor import *
from utils.HDRutils import *
from torch.utils.data import Dataset

class KalantariDataset(Dataset):
    def __init__(self, configs,
                 input_name= 'input_*_aligned.tif',
                 ref_name = 'ref_*_aligned.tif',
                 input_exp_name = 'input_exp.txt',
                 ref_exp_name = 'ref_exp.txt',
                 ref_hdr_name = 'ref_hdr_aligned.hdr'):
        super().__init__()
        print('====> Start preparing training data.')
        filepath = os.path.join(configs.data_path, 'train')
        self.scene_dirs = [scene_dir for scene_dir in os.listdir(filepath)
                            if os.path.isdir(os.path.join(filepath, scene_dir))]
        self.scene_dirs = sorted(self.scene_dirs)
        self.num_scenes = len(self.scene_dirs)

        self.patch_size = configs.patch_size
        self.patch_stride = configs.patch_stride
        self.num_shots = configs.num_shots
        self.patch_path = configs.patch_dir
        self.count = len(os.listdir(self.patch_path))

        if self.count == 0:
            for i, scene_dir in enumerate(self.scene_dirs):
                print(i)
                cur_scene_dir = os.path.join(filepath, scene_dir)
                in_LDR_paths = sorted(glob(os.path.join(cur_scene_dir, input_name)))
                tmp_img = get_image(in_LDR_paths[0])
                h, w, c = tmp_img.shape
                in_LDRs = np.zeros((h, w, c * self.num_shots))
                for j, in_LDR_path in enumerate(in_LDR_paths):
                    in_LDRs[:, :, j * c:(j + 1) * c] = get_image(in_LDR_path)
                in_LDRs = in_LDRs.astype(np.float32)
                in_exps_path = os.path.join(cur_scene_dir, input_exp_name)
                in_exps = np.array(open(in_exps_path).read().split('\n')[:self.num_shots]).astype(np.float32)
                ref_HDR = get_image(os.path.join(cur_scene_dir, ref_hdr_name))
                ref_LDR_paths = sorted(glob(os.path.join(cur_scene_dir, ref_name)))
                ref_LDRs = np.zeros((h, w, c * self.num_shots))
                for j, ref_LDR_path in enumerate(ref_LDR_paths):
                    ref_LDRs[:, :, j * c:(j + 1) * c] = get_image(ref_LDR_path)
                ref_exps_path = os.path.join(cur_scene_dir, ref_exp_name)
                ref_exps = np.array(open(ref_exps_path).read().split('\n')[:self.num_shots]).astype(np.float32)

                # Cut the images into several patches
                # Store patches into files to save memory.
                for _h in range(0, h - self.patch_size[0] + 1, self.patch_stride):
                    for _w in range(0, w - self.patch_size[1] + 1, self.patch_stride):
                        store_patch(_h, _h + self.patch_size[0], _w, _w + self.patch_size[1], in_LDRs, in_exps,
                                    ref_HDR, ref_LDRs, ref_exps, self.patch_path, self.count)
                        self.count += 1
                if h % self.patch_size[0]:
                    for _w in range(0, w - self.patch_size[1] + 1, self.patch_stride):
                        store_patch(h - self.patch_size[0], h, _w, _w + self.patch_size[1], in_LDRs, in_exps,
                                    ref_HDR, ref_LDRs, ref_exps, self.patch_path, self.count)
                        self.count += 1
                if w % self.patch_size[1]:
                    for _h in range(0, h - self.patch_size[0] + 1, self.patch_stride):
                        store_patch(_h, _h + self.patch_size[0], w - self.patch_size[1], w, in_LDRs, in_exps,
                                    ref_HDR, ref_LDRs, ref_exps, self.patch_path, self.count)
                        self.count += 1
                if h % self.patch_size[0] and w % self.patch_size[1]:
                    store_patch(h - self.patch_size[0], h, w - self.patch_size[1], w, in_LDRs, in_exps,
                                ref_HDR, ref_LDRs, ref_exps, self.patch_path, self.count)
                    self.count += 1
        print('====> Finish preparing training data!')

    def __len__(self):
        return self.count

    def __getitem__(self, index):
        distortions = np.random.uniform(0.0, 1.0, 2)

        data = get_patch_from_file(self.patch_path, index)
        in_LDR = data['in_LDR']
        ref_LDR = data['ref_LDR']
        ref_HDR = data['ref_HDR']
        in_exp = data['in_exp']
        ref_exp = data['ref_exp']

        # Horizontal flip
        if distortions[0] < 0.5:
            in_LDR = np.flip(in_LDR, axis=1)
            ref_LDR = np.flip(ref_LDR, axis=1)
            ref_HDR = np.flip(ref_HDR, axis=1)

        # Rotation
        k = int(distortions[1] * 4 + 0.5)
        in_LDR = np.rot90(in_LDR, k)
        ref_LDR = np.rot90(ref_LDR, k)
        ref_HDR = np.rot90(ref_HDR, k)
        in_exp = 2 ** in_exp
        ref_exp = 2 ** ref_exp

        in_HDR = LDR2HDR_batch(in_LDR, in_exp)

        # In pytorch, channels is in axis=1, so ijk -> kij
        in_LDR = np.einsum("ijk->kij", in_LDR)
        ref_LDR = np.einsum("ijk->kij", ref_LDR)
        in_HDR = np.einsum("ijk->kij", in_HDR)
        ref_HDR = np.einsum("ijk->kij", ref_HDR)
        return in_LDR.copy().astype(np.float32), ref_LDR.copy().astype(np.float32), \
               in_HDR.copy().astype(np.float32), ref_HDR.copy().astype(np.float32), \
               in_exp.copy().astype(np.float32), ref_exp.copy().astype(np.float32)

class KalantariTestDataset(Dataset):
    def __init__(self, configs,
                 input_name= 'input_*_aligned.tif',
                 input_exp_name = 'input_exp.txt',
                 ref_hdr_name = 'ref_hdr_aligned.hdr'):
        super().__init__()
        print('====> Start preparing testing data.')
        self.filepath = os.path.join(configs.data_path, 'test')
        self.scene_dirs = [scene_dir for scene_dir in os.listdir(self.filepath)
                            if os.path.isdir(os.path.join(self.filepath, scene_dir))]
        self.scene_dirs = sorted(self.scene_dirs)
        self.num_scenes = len(self.scene_dirs)
        self.patch_size = configs.patch_size
        self.patch_stride = configs.patch_stride
        self.num_shots = configs.num_shots
        self.sample_path = configs.sample_dir
        self.input_name = input_name
        self.input_exp_name = input_exp_name
        self.ref_hdr_name = ref_hdr_name
        print('====> Finish preparing testing data!')

    def __len__(self):
        return self.num_scenes

    def __getitem__(self, index):
        scene_dir = self.scene_dirs[index]
        scene_path = os.path.join(self.filepath, scene_dir)
        sample_path = os.path.join(self.sample_path, scene_dir)
        LDR_path = os.path.join(scene_path, self.input_name)
        exp_path = os.path.join(scene_path, self.input_exp_name)
        ref_HDR_path = os.path.join(scene_path, self.ref_hdr_name)
        in_LDRs, in_HDRs, in_exps, ref_HDRs = get_input(LDR_path, exp_path, ref_HDR_path)
        in_LDRs = np.einsum("ijk->kij", in_LDRs)
        in_HDRs = np.einsum("ijk->kij", in_HDRs)
        ref_HDRs = np.einsum("ijk->kij", ref_HDRs)
        return sample_path, \
               in_LDRs.copy().astype(np.float32), \
               in_HDRs.copy().astype(np.float32), \
               in_exps.copy().astype(np.float32), \
               ref_HDRs.copy().astype(np.float32)


import numpy as np
import cv2
import face_recognition
import datetime

name = ''

video_capture = cv2.VideoCapture(0)

aryan_image = face_recognition.load_image_file('known/aryan.jpg')

aryan_face_encoding = face_recognition.face_encodings(aryan_image)[0]

known_face_encodings = [aryan_face_encoding]
known_face_names = ['Aryan']

def recognize():
    global name

    ret, frame = video_capture.read()

    # Converting the frame from OpenCV's BGR format to the RGB format
    rgb_frame = frame[:, :, ::-1]

    # Finding the face locations and encodings in each frame
    face_locations = face_recognition.face_locations(rgb_frame)
    face_encodings = face_recognition.face_encodings(rgb_frame, face_locations)

    # Now to loop through each face in this frame
    for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):

        # Checking if the face is a match for known faces
        matches = face_recognition.compare_faces(known_face_encodings, face_encoding)

        name = 'Unknown'

        # Use the known face with the smallest vector distance to the new face
        face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
        best_match_index = np.argmin(face_distances)

        if matches[best_match_index]:
            name = known_face_names[best_match_index]
            if name == 'Aryan':
                log_book = open("log_book.txt", "w")
                log_book.write(datetime.datetime.now().strftime('%d-%m-%Y || %H:%M:%S __%A__\n'))
                log_book.close()
                return 0

        # Draw a box around the face
        cv2.rectangle(frame, (left, top), (right, bottom), (255, 0, 0), 2)

        # Draw a label with the name below the face
        cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (255, 0, 0), cv2.FILLED)
        font = cv2.FONT_HERSHEY_DUPLEX
        cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)

    # Display the image
    cv2.imshow('Video', frame)

    # Hit 'q' on the keyboard to quit
    if cv2.waitKey(1) & 0xFF == ord('q'):
        pass

if __name__ == '__main__':
    while True:
        try:
            recognize()
            if recognize() == 0:
                video_capture.release()
                cv2.destroyAllWindows()
        except TypeError:
            break


# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""
RoBERTa: A Robustly Optimized BERT Pretraining Approach.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import sys

import logging
import numpy as np

from fairseq import utils
from fairseq.models import (
    FairseqDecoder,
    FairseqLanguageModel,
    register_model,
    register_model_architecture,
)
from fairseq.modules import (
    LayerNorm,
    TransformerSentenceEncoder,
    TransformerMultiEncoder
)
from fairseq.modules.transformer_sentence_encoder import init_bert_params
#from fairseq.models.roberta import RobertaModel
from fairseq import checkpoint_utils
import pdb
from fairseq.models.roberta import RobertaModel

logger = logging.getLogger(__name__)
#from .hub_interface import RobertaHubInterface

@register_model('semiemo')
class SemiemoModel(FairseqLanguageModel):

    def __init__(self, args, encoder):
        super().__init__(encoder)
        self.args = args

        # We follow BERT's random weight initialization
        self.apply(init_bert_params)

        self.classification_heads = nn.ModuleDict()

        ############################## Adding the pretrained SSL models to extract features###############

        if self.args.a_only or self.args.all_in:

            self.roberta_vqwav2vec = RobertaModel.from_pretrained('/hpc/gsir059/INTERSPEECH/MOSI-SEMI/trained_ssl/wav2vec/vq-wav2vec-Kmeans-Roberta',checkpoint_file='bert_kmeans.pt')

            if self.args.frozen_ssl:
                for param in self.roberta_vqwav2vec.parameters():
                    param.requires_grad = False

        if self.args.t_only or self.args.all_in:
            roberta = torch.hub.load('pytorch/fairseq', 'roberta.large')

            ########################### Freezing pretrained SSL paramtere###################################
            self.model_text2vec=roberta
            if self.args.frozen_ssl:

                for param in self.model_text2vec.parameters():
                    param.requires_grad = False

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        parser.add_argument('--encoder-layers', type=int, metavar='L',
                            help='num encoder layers in plain trans')
        parser.add_argument('--encoder-layers-cross', type=int, metavar='L',
                            help='num encoder layers in cross modal trans')
        parser.add_argument('--encoder-embed-dim', type=int, metavar='H',
                            help='encoder embedding dimension')
        parser.add_argument('--encoder-embed-dim-concat', type=int, metavar='H',
                            help='encoder embedding dimension concatenated')
        parser.add_argument('--encoder-embed-dim-a', type=int, metavar='H',
                            help='encoder embedding dimension cross modal audio')
        parser.add_argument('--encoder-embed-dim-v', type=int, metavar='H',
                            help='encoder embedding dimension cross modal audio')
        parser.add_argument('--encoder-ffn-embed-dim', type=int, metavar='F',
                            help='encoder embedding dimension for FFN')
        parser.add_argument('--encoder-attention-heads', type=int, metavar='A',
                            help='num encoder attention heads')
        parser.add_argument('--activation-fn',
                            choices=utils.get_available_activation_fns(),
                            help='activation function to use')
        parser.add_argument('--pooler-activation-fn',
                            choices=utils.get_available_activation_fns(),
                            help='activation function to use for pooler layer')
        parser.add_argument('--encoder-normalize-before', action='store_true',
                            help='apply layernorm before each encoder block')
        parser.add_argument('--dropout', type=float, metavar='D',
                            help='dropout probability')
        parser.add_argument('--attention-dropout', type=float, metavar='D',
                            help='dropout probability for attention weights')
        parser.add_argument('--activation-dropout', type=float, metavar='D',
                            help='dropout probability after activation in FFN')
        parser.add_argument('--pooler-dropout', type=float, metavar='D',
                            help='dropout probability in the masked_lm pooler layers')
        parser.add_argument('--max-positions-t', type=int,
                            help='number of positional embeddings to learn')
        parser.add_argument('--load-checkpoint-heads', action='store_true',
                            help='(re-)register and load heads when loading checkpoints')

        parser.add_argument('--max-positions-v', type=int,
                            help='number of positional embeddings to learn in video stream')
        parser.add_argument('--max-positions-a', type=int,
                            help='number of positional embeddings to learn in audio stream')

        parser.add_argument('--t-only', action='store_true', default=False,
                            help='do you need only text')

        parser.add_argument('--v-only', action='store_true', default=False,
                                    help='do you need only video')

        parser.add_argument('--a-only', action='store_true', default=False,
                                    help='do you need only audio')

        parser.add_argument('--all-in', action='store_true', default=False,
                                    help='do you need all the embeddings')

        parser.add_argument('--stack-up', action='store_true', default=False,
                                    help='do you need to add an architecture on top of SSL layers')

        parser.add_argument('--frozen-ssl', action='store_true', default=False,
                                    help='do you need to keep the ssl frozen')

    def set_num_updates(self, num_updates):

        self.curr_step = num_updates

    @classmethod
    def build_model(cls, args, task):

        """Build a new model instance."""
        # make sure all arguments are present
        #base_architecture(args) #Supplying the relevence arguments to create the baseline transformer architecture

        encoder = SemiemoEncoder(args, task.source_dictionary)

        return cls(args, encoder)

    def forward(self, src, features_only=False, return_all_hiddens=False, classification_head_name=None, **kwargs):

        final_loss_input = {
            'sup_logits': None, #we can take sup targets as usual
            'sup_targets': None,
            'ori_uda_logits': None,
            'aug_uda_logits': None,
            'simCLR_soft_logits': None,
            'simCLR_soft_labels': None,
            'mixup_logits':None,
            'mixup_targets':None,
            'sup_split':None,

        }

        sup_src=src['sup']
        #unsup_src=src['unsup']
        unsup_src=None

        if sup_src is None or len(sup_src) == 0:
            sup_src_tokens=None

        else:
            sup_src_tokens=sup_src['net_input']
            sup_src_target=sup_src['target']
            sup_src_split=sup_src['split']
            final_loss_input.update(sup_split=sup_src_split)

        if unsup_src is None or len(unsup_src) == 0:
            unsup_src_tokens=None

        else:
            unsup_src_tokens=unsup_src['net_input']

        #unsup_src_tokens=None

        if sup_src_tokens is not None:

            #batch_size_sup=sup_src_target.shape[0]

            aug=False # because we are not augmenting anything

            if classification_head_name is not None:
                features_only = True

            data_dict_sup={} #Create a dictionary with features
            data_dict_sup['raw_data']=sup_src_tokens

            if self.args.t_only or self.args.all_in:
                tokens_text_sup=sup_src_tokens['text']

                #Text SSL feature extraction  # [2, 100, 1024] B X T X D
                roberta_feature_sup=self.model_text2vec.extract_features(tokens_text_sup)
                data_dict_sup['Text']=roberta_feature_sup

            if self.args.a_only or self.args.all_in:

                tokens_audio_sup=sup_src_tokens['audio']

                roberta_vqwav2vec_feature_sup=self.roberta_vqwav2vec.extract_features(tokens_audio_sup)

                data_dict_sup['Audio']=roberta_vqwav2vec_feature_sup

                #Audio SSL feature extraction [2, 512, 310] B X D X T

            #This will output the main models whole features as well as token features
            x, extr_sup = self.decoder(data_dict_sup,features_only, return_all_hiddens,aug, **kwargs) #here the decoder means the encoder (to have the interface fixed)

            if classification_head_name is not None:

                x_sup,_ = self.classification_heads[classification_head_name](extr_sup,Final_rep=False) #Here we send the extracs ("First token")

            final_loss_input.update(sup_logits=x_sup,sup_targets=sup_src_target)

        # if unsup_src_tokens is not None:

        #     aug=True

        #     data_dict_unsup={} #Create a dictionary with features
        #     data_dict_unsup['raw_data']=unsup_src_tokens

        #     if self.args.t_only or self.args.all_in:
        #         tokens_text_unsup=unsup_src_tokens['text']

        #         #Text SSL feature extraction  # [2, 100, 1024] B X T X D
        #         roberta_feature_unsup=self.model_text2vec.extract_features(tokens_text_unsup)
        #         data_dict_unsup['Text']=roberta_feature_unsup

        #     if self.args.a_only or self.args.all_in:

        #         tokens_audio_unsup=unsup_src_tokens['audio']

        #         roberta_vqwav2vec_feature_unsup=self.roberta_vqwav2vec.extract_features(tokens_audio_unsup)

        #         data_dict_unsup['Audio']=roberta_vqwav2vec_feature_unsup

        #         #Audio SSL feature extraction [2, 512, 310] B X D X T

        #     if classification_head_name is not None:
        #         features_only = True

        #     # ori_u_score, aug_u_score= self.UDA_unsup(data_dict_unsup,classification_head_name,return_all_hiddens)

        #     # final_loss_input.update(ori_uda_logits=ori_u_score,aug_uda_logits=aug_u_score)

        # if (unsup_src_tokens is not None) and (sup_src_tokens is not None):

        #     mixup_logits,mixed_target=self.rep_mix_up_BA(data_dict_sup, data_dict_unsup,sup_src_target,classification_head_name,return_all_hiddens)
        #     final_loss_input.update(mixup_logits=mixup_logits,mixup_targets=mixed_target)

        #     #print("no mixmatch")

        #     #use cross entropy
        #     # contrastive_logits,contrastive_labels= self.SimCLR_unsup(data_dict_sup, data_dict_unsup,sup_src_target,classification_head_name,return_all_hiddens)
        #     # final_loss_input.update(simCLR_soft_logits=contrastive_logits,simCLR_soft_labels=contrastive_labels)

        return final_loss_input

    def register_classification_head(self, name, num_classes=None, inner_dim=None, **kwargs):
        """Register a classification head."""

        if name in self.classification_heads:
            prev_num_classes = self.classification_heads[name].out_proj.out_features
            prev_inner_dim = self.classification_heads[name].dense.out_features
            if num_classes != prev_num_classes or inner_dim != prev_inner_dim:
                print(
                    'WARNING: re-registering head "{}" with num_classes {} (prev: {}) '
                    'and inner_dim {} (prev: {})'.format(
                        name, num_classes, prev_num_classes, inner_dim, prev_inner_dim
                    )
                )

        self.classification_heads[name] = SemiemoClassificationHead(
            self.args.encoder_embed_dim_concat,#self.args.encoder_embed_dim,
            #inner_dim or self.args.encoder_embed_dim,
            inner_dim or self.args.encoder_embed_dim_concat,
            num_classes,
            self.args.pooler_activation_fn,
            self.args.pooler_dropout,
            self.args
        )

    def rep_mix_up(self, sup_sample, unsup_sample, targets_x, classification_head_name, return_all_hiddens, **kwargs):

            targets_x=torch.nn.functional.one_hot(targets_x.long(), 4).float()

            #################################################################
            ##only for iemocap
            #targets_x= targets_x.view(-1)
            #targets_x=torch.nn.functional.one_hot(targets_x.long(), 2).float()

            ###################################################################

            Temparature=0.5

            if classification_head_name is not None:
                features_only = True

            if self.args.stack_up:

                with torch.no_grad(): #get the average over the predictions for two augmented versions
                    # compute guessed labels of unlabel samples
                    features_only=True
                    _,rep_u =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #randomly we get two aug
                    _,rep_u2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)

                    outputs_u1,_ = self.classification_heads[classification_head_name](rep_u,Final_rep=False) #from the usual setting
                    outputs_u2,_ = self.classification_heads[classification_head_name](rep_u2,Final_rep=False) #change the text  [cls] position littlebit

                    ############################################################
                    # outputs_u1=outputs_u1.view(-1, 2) #only for iemocap
                    # outputs_u2=outputs_u2.view(-1, 2)
                    ##################################################################################

                    average_u_score_ori = (torch.softmax(outputs_u1, dim=1)+torch.softmax(outputs_u2, dim=1)) / 2
                    pt = average_u_score_ori**(1/Temparature)
                    targets_u=pt / pt.sum(dim=1, keepdim=True)
                    targets_u = targets_u.detach()

                _,rep_unsup =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup = self.classification_heads[classification_head_name](rep_unsup,Final_rep=False)

                _,rep_unsup2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup2 = self.classification_heads[classification_head_name](rep_unsup2,Final_rep=False)

                _,rep_sup =self.decoder(sup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #if we need we can put the  augmentation true
                _,final_sup = self.classification_heads[classification_head_name](rep_sup,Final_rep=False)

                all_inputs = torch.cat([final_sup,final_unsup,final_unsup2], dim=0)  #make a mixed batch
                all_targets = torch.cat([targets_x, targets_u,targets_u], dim=0) #all the examples are together

                # ############################################################################################
                # all_targets=all_targets.view(all_inputs.shape[0],-1) #only for binary emocap thing

            alpha=0.75
            Temparature=0.5

            l = np.random.beta(alpha, alpha)
            l = max(l, 1-l)
            idx = torch.randperm(all_inputs.size(0))

            input_a, input_b = all_inputs, all_inputs[idx]
            target_a, target_b = all_targets, all_targets[idx]

            mixed_input = l * input_a + (1 - l) * input_b  #mixing the inputs with
            mixed_target = l * target_a + (1 - l) * target_b
            mixup_logits,_=self.classification_heads[classification_head_name](mixed_input,Final_rep=True)

            return mixup_logits,mixed_target

    def rep_mix_up_BA(self, sup_sample, unsup_sample, targets_x, classification_head_name, return_all_hiddens, **kwargs):

            targets_x=torch.nn.functional.one_hot(targets_x.long(), 4).float()
            ################################################################
            #only for iemocap
            targets_x= targets_x.view(-1)
            targets_x=torch.nn.functional.one_hot(targets_x.long(), 2).float()
            ##################################################################

            Temparature=0.5

            if classification_head_name is not None:
                features_only = True

            if self.args.stack_up:

                with torch.no_grad(): #get the average over the predictions for two augmented versions
                    # compute guessed labels of unlabel samples
                    features_only=True
                    _,rep_u =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #randomly we get two aug
                    _,rep_u2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)

                    outputs_u1,_ = self.classification_heads[classification_head_name](rep_u,Final_rep=False) #from the usual setting
                    outputs_u2,_ = self.classification_heads[classification_head_name](rep_u2,Final_rep=False) #change the text  [cls] position littlebit

                    ###########################################################
                    outputs_u1=outputs_u1.view(-1, 2) #only for iemocap
                    outputs_u2=outputs_u2.view(-1, 2)
                    #################################################################################

                    average_u_score_ori = (torch.softmax(outputs_u1, dim=1)+torch.softmax(outputs_u2, dim=1)) / 2
                    pt = average_u_score_ori**(1/Temparature)
                    targets_u=pt / pt.sum(dim=1, keepdim=True)
                    targets_u = targets_u.detach()

                _,rep_unsup =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup = self.classification_heads[classification_head_name](rep_unsup,Final_rep=False)

                _,rep_unsup2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup2 = self.classification_heads[classification_head_name](rep_unsup2,Final_rep=False)

                _,rep_sup =self.decoder(sup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #if we need we can put the  augmentation true
                _,final_sup = self.classification_heads[classification_head_name](rep_sup,Final_rep=False)

                all_inputs = torch.cat([final_sup,final_unsup,final_unsup2], dim=0)  #make a mixed batch
                all_targets = torch.cat([targets_x, targets_u,targets_u], dim=0) #all the examples are together

                # ############################################################################################
                all_targets=all_targets.view(all_inputs.shape[0],-1) #only for binary emocap thing

            alpha=0.75
            Temparature=0.5

            l = np.random.beta(alpha, alpha)
            l = max(l, 1-l)
            idx = torch.randperm(all_inputs.size(0))

            input_a, input_b = all_inputs, all_inputs[idx]
            target_a, target_b = all_targets, all_targets[idx]

            mixed_input = l * input_a + (1 - l) * input_b  #mixing the inputs with
            mixed_target = l * target_a + (1 - l) * target_b
            mixup_logits,_=self.classification_heads[classification_head_name](mixed_input,Final_rep=True)

            return mixup_logits,mixed_target

    def rep_mix_up_no_stack(self, sup_sample, unsup_sample, targets_x, classification_head_name, return_all_hiddens, **kwargs):

            targets_x=torch.nn.functional.one_hot(targets_x.long(), 4).float()

            #################################################################
            ##only for iemocap
            #targets_x= targets_x.view(-1)
            #targets_x=torch.nn.functional.one_hot(targets_x.long(), 2).float()

            ###################################################################

            Temparature=0.5

            if classification_head_name is not None:
                features_only = True

            if self.args.stack_up:

                with torch.no_grad(): #get the average over the predictions for two augmented versions
                    # compute guessed labels of unlabel samples
                    features_only=True
                    _,rep_u =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #randomly we get two aug
                    _,rep_u2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=False, **kwargs)

                    outputs_u1,_ = self.classification_heads[classification_head_name](rep_u,Final_rep=False) #from the usual setting
                    outputs_u2,_ = self.classification_heads[classification_head_name](rep_u2,Final_rep=False) #change the text  [cls] position littlebit

                    ############################################################
                    # outputs_u1=outputs_u1.view(-1, 2) #only for iemocap
                    # outputs_u2=outputs_u2.view(-1, 2)
                    ##################################################################################

                    average_u_score_ori = (torch.softmax(outputs_u1, dim=1)+torch.softmax(outputs_u2, dim=1)) / 2
                    pt = average_u_score_ori**(1/Temparature)
                    targets_u=pt / pt.sum(dim=1, keepdim=True)
                    #targets_u = targets_u.detach()

                _,rep_unsup =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup = self.classification_heads[classification_head_name](rep_unsup,Final_rep=False)

                _,rep_unsup2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
                _,final_unsup2 = self.classification_heads[classification_head_name](rep_unsup2,Final_rep=False)

                _,rep_sup =self.decoder(sup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #if we need we can put the  augmentation true
                _,final_sup = self.classification_heads[classification_head_name](rep_sup,Final_rep=False)

                all_inputs = torch.cat([final_sup,final_unsup,final_unsup2], dim=0)  #make a mixed batch
                all_targets = torch.cat([targets_x, targets_u,targets_u], dim=0) #all the examples are together

                # ############################################################################################
                # all_targets=all_targets.view(all_inputs.shape[0],-1) #only for binary emocap thing

            alpha=0.75
            Temparature=0.5

            l = np.random.beta(alpha, alpha)
            l = max(l, 1-l)
            idx = torch.randperm(all_inputs.size(0))

            input_a, input_b = all_inputs, all_inputs[idx]
            target_a, target_b = all_targets, all_targets[idx]

            mixed_input = l * input_a + (1 - l) * input_b  #mixing the inputs with
            mixed_target = l * target_a + (1 - l) * target_b
            mixup_logits,_=self.classification_heads[classification_head_name](mixed_input,Final_rep=True)

            return mixup_logits,mixed_target

    def manifold_mix_up(self, sup_sample, unsup_sample, targets_x, classification_head_name, return_all_hiddens, **kwargs):

            targets_x=torch.nn.functional.one_hot(targets_x.long(), 4).float()
            targets_x_hat=torch.nn.functional.one_hot(unsup_sample['targets'].long(), 4).float()

            #################################################################
            ##only for iemocap
            #targets_x= targets_x.view(-1)
            #targets_x=torch.nn.functional.one_hot(targets_x.long(), 2).float()

            ###################################################################

            Temparature=0.5

            # if classification_head_name is not None:
            #     features_only = True

            # if self.args.stack_up:

            #     with torch.no_grad(): #get the average over the predictions for two augmented versions
            #         # compute guessed labels of unlabel samples
            #         features_only=True
            #         _,rep_u =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #randomly we get two aug
            #         _,rep_u2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)

            #         outputs_u1,_ = self.classification_heads[classification_head_name](rep_u,Final_rep=False) #from the usual setting
            #         outputs_u2,_ = self.classification_heads[classification_head_name](rep_u2,Final_rep=False) #change the text  [cls] position littlebit

            #         ############################################################
            #         # outputs_u1=outputs_u1.view(-1, 2) #only for iemocap
            #         # outputs_u2=outputs_u2.view(-1, 2)
            #         ##################################################################################

            #         average_u_score_ori = (torch.softmax(outputs_u1, dim=1)+torch.softmax(outputs_u2, dim=1)) / 2
            #         pt = average_u_score_ori**(1/Temparature)
            #         targets_u=pt / pt.sum(dim=1, keepdim=True)
            #         targets_u = targets_u.detach()

            _,rep_unsup =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
            _,final_unsup = self.classification_heads[classification_head_name](rep_unsup,Final_rep=False)

            _,rep_unsup2 =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)
            _,final_unsup2 = self.classification_heads[classification_head_name](rep_unsup2,Final_rep=False)

            _,rep_sup =self.decoder(sup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #if we need we can put the  augmentation true
            _,final_sup = self.classification_heads[classification_head_name](rep_sup,Final_rep=False)

            all_inputs = torch.cat([final_sup,final_unsup,final_unsup2], dim=0)  #make a mixed batch
            all_targets = torch.cat([targets_x, targets_u,targets_u], dim=0) #all the examples are together

            alpha=0.75
            Temparature=0.5

            l = np.random.beta(alpha, alpha)
            l = max(l, 1-l)
            idx = torch.randperm(all_inputs.size(0))

            input_a, input_b = all_inputs, all_inputs[idx]
            target_a, target_b = all_targets, all_targets[idx]

            mixed_input = l * input_a + (1 - l) * input_b  #mixing the inputs with
            mixed_target = l * target_a + (1 - l) * target_b
            mixup_logits,_=self.classification_heads[classification_head_name](mixed_input,Final_rep=True)

            return mixup_logits,mixed_target

    def UDA_unsup(self, unsup_sample,classification_head_name,return_all_hiddens, **kwargs):

        if classification_head_name is not None:
            features_only = True

        with torch.no_grad(): #get the average over the predictions for two augmented versions
            # compute guessed labels of unlabel samples
            _,rep_u_ori =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=False, **kwargs) #randomly we get two aug
            outputs_u1_ori,_ = self.classification_heads[classification_head_name](rep_u_ori,Final_rep=False) #from the usual setting

        _,rep_u =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #randomly we get two aug
        outputs_u1,_ = self.classification_heads[classification_head_name](rep_u,Final_rep=False) #from the usual setting

        ############################

        return outputs_u1,outputs_u1_ori

    def SimCLR_unsup(self, sup_sample, unsup_sample,targets_x,classification_head_name,return_all_hiddens, **kwargs):

        self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
        self.mask_samples_from_same_repr = self._get_correlated_mask(targets_x.shape[0]*2).type(torch.bool)

        if classification_head_name is not None:
            features_only = True

        _,rep_u_dict =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #randomly we get two aug
        _,rep_u2_dict =self.decoder(unsup_sample,features_only, return_all_hiddens,aug=True, **kwargs)

        _,final_rep_u = self.classification_heads[classification_head_name](rep_u_dict,Final_rep=False)
        _,final_rep_u2=self.classification_heads[classification_head_name](rep_u2_dict,Final_rep=False)

        _,rep_sup_dict =self.decoder(sup_sample,features_only, return_all_hiddens,aug=True, **kwargs) #newly added
        _,rep_sup2_dict =self.decoder(sup_sample,features_only, return_all_hiddens,aug=True, **kwargs)#newly added

        _,final_rep_sup = self.classification_heads[classification_head_name](rep_sup_dict,Final_rep=False)#newly added
        _,final_rep_sup2=self.classification_heads[classification_head_name](rep_sup2_dict,Final_rep=False)#newly added

        final_rep=torch.cat([final_rep_sup,final_rep_u], dim=0)
        final_rep2=torch.cat([final_rep_sup2,final_rep_u2], dim=0)

        # normalize projection feature vectors
        # repis = F.normalize(final_rep_u, dim=1)  #first augmetation copy of the batch
        # repjs = F.normalize(final_rep_u2, dim=1) #second augmentation copy of the batch
        repis = F.normalize(final_rep, dim=1)  #first augmetation copy of the batch
        repjs = F.normalize(final_rep2, dim=1) #second augmentation copy of the batch
        representations = torch.cat([repis, repjs], dim=0)

        similarity_matrix = self._cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0)) #squeeze is really important

        batch_size=targets_x.shape[0]*2
        temperature=0.5

        # filter out the scores from the positive samples
        l_pos = torch.diag(similarity_matrix, batch_size)
        r_pos = torch.diag(similarity_matrix, -batch_size)
        positives = torch.cat([l_pos, r_pos]).view(2 * batch_size, 1)

        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * batch_size, -1)

        logits = torch.cat((positives, negatives), dim=1)

        logits = logits/temperature
        labels = torch.zeros(2 * batch_size).long() #every time it is the correct location

        #use cross entropy loss as normal - reduction sum
        #make sure that the batch size from this is twwise

        return logits,labels

    def _get_correlated_mask(self,batch_size): #for the SimCLR method
        diag = np.eye(2 * batch_size)
        l1 = np.eye((2 * batch_size), 2 * batch_size, k=-batch_size)
        l2 = np.eye((2 * batch_size), 2 * batch_size, k=batch_size)
        mask = torch.from_numpy((diag + l1 + l2))
        mask = (1 - mask).type(torch.bool)
        return mask  #.to(self.device)

    @property
    def supported_targets(self):
        return {'self'}

    def upgrade_state_dict_named(self, state_dict, name):
        prefix = name + '.' if name != '' else ''
        current_head_names = [] if not hasattr(self, 'classification_heads') else \
            self.classification_heads.keys()

        # Handle new classification heads present in the state dict.
        keys_to_delete = []
        for k in state_dict.keys():
            if not k.startswith(prefix + 'classification_heads.'):
                continue

            #print(k)

            head_name = k[len(prefix + 'classification_heads.'):].split('.')[0]
            num_classes = state_dict[prefix + 'classification_heads.' + head_name + '.out_proj.weight'].size(0)
            #inner_dim = state_dict[prefix + 'classification_heads.' + head_name + '.dense.weight'].size(0)
            #Because we do not have a dense layer
            inner_dim =state_dict[prefix + 'classification_heads.' + head_name + '.out_proj.weight'].size(1)

            if getattr(self.args, 'load_checkpoint_heads', False):
                if head_name not in current_head_names:
                    self.register_classification_head(head_name, num_classes, inner_dim)
            else:
                if head_name not in current_head_names:
                    print(
                        'WARNING: deleting classification head ({}) from checkpoint '
                        'not present in current model: {}'.format(head_name, k)
                    )
                    keys_to_delete.append(k)
                elif (

                    num_classes != self.classification_heads[head_name].out_proj.out_features
                    # or inner_dim != self.classification_heads[head_name].dense.out_features
                    or inner_dim != self.classification_heads[head_name].out_proj.in_features #because we do not have a dense layer

                ):
                    print(
                        'WARNING: deleting classification head ({}) from checkpoint '
                        'with different dimensions than current model: {}'.format(head_name, k)
                    )
                    keys_to_delete.append(k)

        for k in keys_to_delete:
            del state_dict[k]

        # Copy any newly-added classification heads into the state dict
        # with their current weights.
        if hasattr(self, 'classification_heads'):
            cur_state = self.classification_heads.state_dict()
            for k, v in cur_state.items():
                if prefix + 'classification_heads.' + k not in state_dict:
                    print('Overwriting', prefix + 'classification_heads.' + k)
                    state_dict[prefix + 'classification_heads.' + k] = v

def upgrade_state_dict_named(self, state_dict, name):
    super().upgrade_state_dict_named(state_dict, name)

    prefix = name + '.' if name != '' else ''
    current_head_names = [] if not hasattr(self, 'classification_heads') else \
        self.classification_heads.keys()

    # Handle new classification heads present in the state dict.
    keys_to_delete = []
    for k in state_dict.keys():
        if not k.startswith(prefix + 'classification_heads.'):
            continue

        head_name = k[len(prefix + 'classification_heads.'):].split('.')[0]
        num_classes = state_dict[prefix + 'classification_heads.' + head_name + '.out_proj.weight'].size(0)
        #inner_dim = state_dict[prefix + 'classification_heads.' + head_name + '.dense.weight'].size(0)
        inner_dim =state_dict[prefix + 'classification_heads.' + head_name + '.out_proj.weight'].size(1)

        if getattr(self.args, 'load_checkpoint_heads', False):
            if head_name not in current_head_names:
                self.register_classification_head(head_name, num_classes, inner_dim)
        else:
            if head_name not in current_head_names:
                logger.warning(
                    'deleting classification head ({}) from checkpoint '
                    'not present in current model: {}'.format(head_name, k)
                )
                keys_to_delete.append(k)
            elif (
                num_classes != self.classification_heads[head_name].out_proj.out_features
                or inner_dim != self.classification_heads[head_name].dense.out_features
            ):
                logger.warning(
                    'deleting classification head ({}) from checkpoint '
                    'with different dimensions than current model: {}'.format(head_name, k)
                )
                keys_to_delete.append(k)
    for k in keys_to_delete:
        del state_dict[k]

    # Copy any newly-added classification heads into the state dict
    # with their current weights.
    if hasattr(self, 'classification_heads'):
        cur_state = self.classification_heads.state_dict()
        for k, v in cur_state.items():
            if prefix + 'classification_heads.' + k not in state_dict:
                logger.info('Overwriting ' + prefix + 'classification_heads.' + k)
                state_dict[prefix + 'classification_heads.' + k] = v

class SemiemoClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""

    def __init__(self, input_dim, inner_dim, num_classes, activation_fn, pooler_dropout,args):
        super().__init__()

        #inner_dim needs to change

        #self.dense = nn.Linear(input_dim, inner_dim)
        #self.activation_fn = utils.get_activation_fn(activation_fn) #Activation function is Tanh
        self.dropout = nn.Dropout(p=pooler_dropout)
        self.out_proj = nn.Linear(inner_dim, num_classes)

        self.args=args

    def forward(self, features,Final_rep, **kwargs):

        if not Final_rep: #we do not get the final representation drectly

            if self.args.t_only & (not self.args.stack_up):
                Final=features['j_text']

            if self.args.a_only & (not self.args.stack_up) :
                Final=features['j_aud']

            if (self.args.all_in) or (self.args.a_only and self.args.t_only):

                if self.args.stack_up:

                    #exit("haha check the stack up section")

                    #exit("###############################################################")

                    T_A=features['t2a_r']
                    A_T=features['a2t_r']
                    Final=torch.cat((T_A,A_T),dim=1)

                else:
                    T_A=features['j_text']
                    A_T=features['j_aud']
                    Final=torch.cat((T_A,A_T),dim=1)

        else: #we get the final represetation directly mix match case
            Final=features

        x =Final# features[:, 0, :]  # take <s> token (equiv. to [CLS])
        x = self.dropout(x)
        x = self.out_proj(x)

        return x,Final # we also output the final representation

class SemiemoEncoder(FairseqDecoder):
    """RoBERTa encoder.

    Implements the :class:`~fairseq.models.FairseqDecoder` interface required
    by :class:`~fairseq.models.FairseqLanguageModel`.
    """

    def __init__(self, args, dictionary):

        super().__init__(dictionary)

        self.args = args

        self.sentence_encoder = TransformerMultiEncoder(     #This encodes all the
            padding_idx=1,#dictionary.pad(), #Check!
            #vocab_size=0,#len(dictionary),#Check!
            #num_encoder_layers=args.encoder_layers,
            num_encoder_layers_cross=args.encoder_layers_cross,
            #embedding_dim=args.encoder_embed_dim,
            embedding_dim_text=args.encoder_embed_dim_t,
            embedding_dim_audio=args.encoder_embed_dim_a,
            #embedding_dim_video=args.encoder_embed_dim_v,
            ffn_embedding_dim=args.encoder_ffn_embed_dim,
            num_attention_heads=args.encoder_attention_heads,
            dropout=args.dropout,
            attention_dropout=args.attention_dropout,
            activation_dropout=args.activation_dropout,
            #max_seq_len_text=args.max_positions_t,#This is the text
            #max_seq_len_audio=args.max_positions_a,
            #max_seq_len_video=args.max_positions_v,
            #num_segments=0,
            #encoder_normalize_before=True,
            apply_bert_init=True,
            activation_fn=args.activation_fn,
            is_only_text=args.t_only,
            is_only_audio=args.a_only,
            #is_only_video=args.v_only,
            is_all_in=args.all_in,
            is_stack_up=args.stack_up

        )

    def forward(self, src_tokens, features_only=False, return_all_hiddens=False,aug=False, **unused):
        """
        Args:
            src_tokens (LongTensor): input tokens of shape `(batch, src_len)`
            features_only (bool, optional): skip LM head and just return
                features. If True, the output will be of shape
                `(batch, src_len, embed_dim)`.
            return_all_hiddens (bool, optional): also return all of the
                intermediate hidden states (default: False).

        Returns:
            tuple:
                - the  LM output of shape `(batch, src_len, vocab)`
                - a dictionary of additional data, where 'inner_states'
                  is a list of hidden states.
        """

        x, extra = self.extract_features(src_tokens, return_all_hiddens,aug) #Get the output after cross attentions and self attention mechanisms

        #Next track the loss function
        if not features_only:
            x = self.output_layer(x)
        return x, extra

    def extract_features(self, src_tokens, return_all_hiddens=False,aug=False, **unused):  #This

        inner_states, seq_rep = self.sentence_encoder(  #use the foewardfunction in TransformerMultiEncoder object
            src_tokens, last_state_only=not return_all_hiddens,is_aug=aug
        )

        return inner_states,seq_rep

    def max_positions(self):
        """Maximum length supported by the model."""
        return None #self.args.max_positions
        #return sys.maxsize

@register_model_architecture('semiemo', 'semiemo')
def semiemo_architecture(args):
    args.encoder_layers = getattr(args, 'encoder_layers', 24)
    args.encoder_layers_cross = getattr(args, 'encoder_layers_cross', 1)
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 1024)
    args.encoder_embed_dim_concat = getattr(args, 'encoder_embed_dim_concat',1792 )  #2048(1024 + 512 + 256)
    #args.encoder_embed_dim_concat = getattr(args, 'encoder_embed_dim_concat',1280 )  #2048(1024 + 512 + 256) #audio only
    args.encoder_embed_dim_t = getattr(args, 'encoder_embed_dim_t', 1024)
    args.encoder_embed_dim_a = getattr(args, 'encoder_embed_dim_a', 768)
    args.encoder_embed_dim_v = getattr(args, 'encoder_embed_dim_v', 256)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 1024) #previously this was 1024
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 16)

    args.activation_fn = getattr(args, 'activation_fn', 'gelu')
    args.pooler_activation_fn = getattr(args, 'pooler_activation_fn', 'tanh')

    args.dropout = getattr(args, 'dropout', 0.1)
    args.attention_dropout = getattr(args, 'attention_dropout', 0.1)
    args.activation_dropout = getattr(args, 'activation_dropout', 0.0)
    args.pooler_dropout = getattr(args, 'pooler_dropout', 0.0)


import numpy as np
from numpy import asarray
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import Rectangle
import os
from cv2 import *
import sys
from tensorflow import keras
from keras.preprocessing.image import img_to_array, load_img
from tensorflow.keras.models import Sequential, save_model, load_model
from mtcnn.mtcnn import MTCNN
from keras_vggface.utils import preprocess_input
from keras_vggface.vggface import VGGFace
from scipy.spatial.distance import cosine

"""
python filename.py (groupimage/single photo) (face)

Try to make sure that the faces themselves in the images are atleast 224x224

This script uses mctnn and vggface, both are MIT licensed.
https://github.com/ipazc/mtcnn
https://github.com/rcmalli/keras-vggface
"""

os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

firstImage = sys.argv[1]
secondImage = sys.argv[2]

#Input is face pixel data, return facial features
def get_model_scores(faces):
    samples = asarray(faces, 'float32')

    # prepare the data for the model
    samples = preprocess_input(samples, version=2)

    # create a vggface model object
    model = VGGFace(model='resnet50',
      include_top=False,
      input_shape=(224, 224, 3),
      pooling='avg')

    # perform prediction
    return model.predict(samples)

# Input is the filename, return is the coordinates of al the faces in the photo
def getFaces(filename):
    faces = []
    image = Image.open(filename)
    image = image.convert('RGB')
    pixels = asarray(image)
    detector = MTCNN()
    results = detector.detect_faces(pixels)

    #Go through every image append and return the processed data
    for result in results:
        x1, y1, width, height = result['box']
        # bug fix
        x1, y1 = abs(x1), abs(y1)
        x2, y2 = x1 + width, y1 + height
        x1, y1, width, height = result['box']
        face = pixels[y1:y2, x1:x2]
        point_array.append(result['box'])
        image = Image.fromarray(face)
        image = image.resize((224, 224))
        face_array = asarray(image)
        faces.append(face_array)
    return faces

point_array = []
test = getFaces(FirstImage)
face = getFaces(secondImage)
score_face = get_model_scores(face)

#Compare the returned features with each face and determine whether or not it's the same face.
counter = 0
for singleFace in test:
    test = np.expand_dims(singleFace, axis=0)
    score_test = get_model_scores(test)
    if cosine(score_test, score_face) <= 0.4:
        print("Faces Matched " + str(counter))
        start_point = point_array[counter]
        print (start_point)

        #Create a red box around the face and save the image.
        x1, y1, width, height = start_point
        x1, y1 = abs(x1), abs(y1)
        im = np.array(Image.open(firstImage), dtype=np.uint8)
        fig,ax = plt.subplots(1)
        ax.imshow(im)
        rect = patches.Rectangle((x1,y1),width,height,linewidth=1,edgecolor='r',facecolor='none')
        ax.add_patch(rect)
        plt.savefig('output.png', dpi=300, bbox_inches='tight')

        counter += 1
    else:
        print("Faces don't match " + str(counter))
        counter += 1


# coding:utf-8

# 配合main3.py使用的模型，将SRL预测和Sent预测分离开来

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

use_cuda = torch.cuda.is_available()

class Model4(nn.Module):
    def __init__(self, args, embed=None):
        super(Model4, self).__init__()
        self.model_name = 'Model4'
        self.args = args
        self.V = args.embed_num
        self.D = args.embed_dim
        self.H = args.hidden_size
        self.P = args.pos_dim

        # word embedding层
        self.embed = nn.Embedding(self.V, self.D, padding_idx=0)
        if embed is not None:
            self.embed.weight.data.copy_(embed)

        self.word_RNN = nn.GRU(
            input_size=self.D,
            hidden_size=self.H,
            batch_first=True,
            bidirectional=True
        )

        self.sent_RNN = nn.GRU(
            input_size=2 * self.H,
            hidden_size=self.H,
            batch_first=True,
            bidirectional=True
        )

        self.doc_RNN = nn.GRU(
            input_size=2 * self.H,
            hidden_size=self.H,
            batch_first=True,
            bidirectional=True
        )

        self.event_RNN = nn.GRU(
            input_size=self.D,
            hidden_size=self.H,
            batch_first=True,
            bidirectional=True
        )

        # position embedding，将sent、doc相对位置映射成一个位置向量
        self.doc_pos_embed = nn.Embedding(self.args.pos_doc_size, self.P)
        self.sent_pos_embed = nn.Embedding(self.args.pos_sent_size, self.P)

        # event预测层，考虑SRL内容，SRL与所在sent、doc、blog的相关性，SRL所在位置，SRL的TF值
        self.event_content = nn.Linear(2 * self.H, 1, bias=False)
        self.event_salience = nn.Bilinear(2 * self.H, 2 * self.H, 1, bias=False)
        self.event_tf = nn.Parameter(torch.FloatTensor(1).uniform_(-0.1, 0.1))
        self.event_bias = nn.Parameter(torch.FloatTensor(1).uniform_(-0.1, 0.1))

        # sent预测层，考虑sent内容，sent与所在doc及blog相关性，sent所在doc的位置，sent在doc中的位置，sent的SRL信息
        self.sent_content = nn.Linear(2 * self.H, 1, bias=False)
        self.sent_salience = nn.Bilinear(2 * self.H, 4 * self.H, 1, bias=False)
        self.sent_doc_pos = nn.Linear(self.P, 1, bias=False)
        self.sent_pos = nn.Linear(self.P, 1, bias=False)
        self.event_rel = nn.Bilinear(2 * self.H, 2 * self.H, 1)  # 计算句子和event相关性
        self.event_para = nn.Parameter(torch.FloatTensor(1).uniform_(-0.1, 0.1))
        self.event_zero = nn.Parameter(torch.FloatTensor(1).uniform_(-0.1, 0.1))
        self.sent_bias = nn.Parameter(torch.FloatTensor(1).uniform_(-0.1, 0.1))

    def max_pool1d(self, x, seq_lens):
        out = []
        for index, t in enumerate(x):
            if seq_lens[index] == 0:
                if use_cuda:
                    out.append(torch.zeros(1, 2 * self.H, 1).cuda())
                else:
                    out.append(torch.zeros(1, 2 * self.H, 1))
            else:
                t = t[:seq_lens[index], :]
                t = torch.t(t).unsqueeze(0)
                out.append(F.max_pool1d(t, t.size(2)))

        out = torch.cat(out).squeeze(2)
        return out

    def avg_pool1d(self, x, seq_lens):
        out = []
        for index, t in enumerate(x):
            if seq_lens[index] == 0:
                if use_cuda:
                    out.append(torch.zeros(1, 2 * self.H, 1).cuda())
                else:
                    out.append(torch.zeros(1, 2 * self.H, 1))
            else:
                t = t[:seq_lens[index], :]
                t = torch.t(t).unsqueeze(0)
                out.append(F.avg_pool1d(t, t.size(2)))

        out = torch.cat(out).squeeze(2)
        return out

    def forward(self, x, doc_lens, events, event_doc_lens, event_tfs, train_srl):

        event_lens = [4] * events.size(0)
        if self.args.embed_frozen:
            events = self.embed(events).detach()
        else:
            events = self.embed(events)
        events, hn = self.event_RNN(events)
        event_vec = self.max_pool1d(events, event_lens)
        t = event_vec.unsqueeze(0)
        event_sum = self.avg_pool1d(t, [t.size(1)]).squeeze(0)

        if train_srl:
            event_probs = []
            for i, event in enumerate(event_vec):
                event_content = self.event_content(event)
                event_salience = self.event_salience(event, event_sum)
                event_tf = self.event_tf * event_tfs[i]
                event_pre = event_content + event_tf + self.event_bias + event_salience
                event_probs.append(event_pre)
            event_probs = torch.cat(event_probs).squeeze()
            return event_probs

        # word => sent
        sent_lens = torch.sum(torch.sign(x), dim=1).data
        if self.args.embed_frozen:
            x = self.embed(x).detach()  # total_sent_num * word_num * D
        else:
            x = self.embed(x)
        x, hn = self.word_RNN(x)  # total_sent_num * word_num * (2*H)
        sent_vec = self.max_pool1d(x, sent_lens)  # total_sent_num * (2*H)

        # sent => doc
        docs = self.seq_split(sent_vec, doc_lens)
        doc_vec = []
        for i, doc in enumerate(docs):
            tmp_h, hn = self.sent_RNN(doc.unsqueeze(0))
            doc_vec.append(self.max_pool1d(tmp_h, [doc_lens[i]]).squeeze(0))
        doc_vec = torch.cat(doc_vec).view(len(doc_lens), -1)  # total_doc_num * (2*H)

        # doc => blog
        x = doc_vec.unsqueeze(0)  # 1 * total_doc_num * (2*H)
        x, hn = self.doc_RNN(x)  # 1 * total_doc_num * (2*H)
        blog_vec = self.max_pool1d(x, [x.size(1)]).squeeze(0)  # (2*H)

        doc_num = float(len(doc_lens))
        event_scores = [.0] * len(event_tfs)
        idxs = np.argsort(event_tfs.detach().cpu().numpy()).tolist()
        idxs.reverse()
        num_one = int(len(event_tfs) * 0.15)
        for i in idxs[:num_one]:
            event_scores[i] = 1.0
        event_scores = torch.FloatTensor(event_scores)
        if use_cuda:
            event_scores = event_scores.cuda()
        # 直接使用event score
        # event_scores = event_probs.detach()

        # sent predictor
        sent_probs = []
        sent_idx = 0
        event_start = 0
        for i in range(0, len(doc_lens)):
            context = torch.cat((blog_vec, doc_vec[i]))
            cur_event_vec = event_vec[event_start: event_start + event_doc_lens[i]]
            cur_event_prob = event_scores[event_start: event_start + event_doc_lens[i]]
            event_start += event_doc_lens[i]
            for j in range(0, doc_lens[i]):
                # 计算当前句子和所在doc的所有event的相似度
                if len(cur_event_vec) == 0:
                    event_rel = self.event_zero
                else:
                    event_sim = self.event_rel(
                        sent_vec[sent_idx].repeat(cur_event_vec.size(0)).view(cur_event_vec.size(0), -1),
                        cur_event_vec).squeeze(1)
                    event_rel = self.event_para * torch.dot(event_sim, cur_event_prob)
                cur_sent = sent_vec[sent_idx]
                sent_content = self.sent_content(cur_sent)
                sent_salience = self.sent_salience(cur_sent, context)
                sent_doc_index = torch.LongTensor([[int(i * self.args.pos_doc_size / doc_num)]])
                sent_index = torch.LongTensor([[int(j * self.args.pos_sent_size / doc_lens[i])]])
                if use_cuda:
                    sent_doc_index = sent_doc_index.cuda()
                    sent_index = sent_index.cuda()
                sent_doc_pos = self.sent_doc_pos(self.doc_pos_embed(sent_doc_index).squeeze(0))
                sent_pos = self.sent_pos(self.sent_pos_embed(sent_index).squeeze(0))
                sent_pre = sent_content + sent_salience + sent_doc_pos + sent_pos + event_rel + self.sent_bias
                sent_probs.append(sent_pre)
                sent_idx += 1
        return torch.cat(sent_probs).squeeze()

    @staticmethod
    def seq_split(vecs, seq_lens):
        rst = []
        start = 0
        for seq_len in seq_lens:
            rst.append(vecs[start: start + seq_len])
            start += seq_len
        assert start == len(vecs)
        return rst

    @staticmethod
    def seq_accumulate(seq):
        rst = [0] * len(seq)
        for i in range(0, len(seq)):
            if i == 0:
                rst[i] = seq[i]
            else:
                rst[i] = seq[i] + rst[i - 1]
        return rst

    def save(self, dir):
        checkpoint = {'model': self.state_dict(), 'args': self.args}
        torch.save(checkpoint, dir)


"""
game.py

Create a game instance which acts as a container for all the other elements (snake, apple).
"""
from random import randrange

import numpy as np

from environment.snake import Snake
from utils.direction import DOWN, LEFT, RIGHT, turn_left, turn_right
from utils.exceptions import PositionException
from utils.pos import Pos

class Game:
    __slots__ = {
        'width', 'height', 'dim', 'pixels', 'score', 'steps',
        'snake',
        'apple', 'apple_separate',
        'board'
    }

    def __init__(self,
                 width=11,
                 height=11,
                 pixels=20,
                 apple_separate=True,
                 ):
        """
        Game object used as an environment to contain the snake entity.

        :param width: Width of the playing field
        :param height: Height of the playing field
        :param pixels: Number of pixels for each tile of the game
        :param apple_separate: Give the apple its own board
        """
        self.width = width if width % 2 == 1 else width + 1  # Always odd
        self.height = height if height % 2 == 1 else height + 1  # Always odd
        self.dim = (self.width, self.height)
        self.pixels = pixels
        self.score = 0
        self.steps = 0

        # Initialise the snake
        self.snake = Snake(game=self)

        # Initialise the apple
        self.apple = None
        self.apple_separate: bool = apple_separate
        self.set_apple_pos()

        # Initialise the board
        self.board = self.create_board()

    # ----------------------------------------------------> MAIN <---------------------------------------------------- #

    def step(self, a, uses_tuple: bool = False):
        """
        Update the game with the corresponding action.

        :param a: Action to perform
        :param uses_tuple: Actions are in tuple representation (action, randomised)
        :return: Boolean indicating if game is still running
        """
        random_a = False
        if uses_tuple:
            random_a = a[1]
            a = a[0]
        assert a in [0, 1, 2]  # Straight, left, right
        if a == 1: self.snake.direction = turn_left(self.snake.direction)
        if a == 2: self.snake.direction = turn_right(self.snake.direction)

        # Update snake position
        try:
            if self.snake.step(apple=self.apple):
                self.set_apple_pos()
                self.score += .5  # Reward for eating apple

            # Update the board
            self.update_board()
            return True
        except PositionException:
            self.score -= .2 if random_a else 1  # Punish for hitting into a wall, less punishment if random action
            return False

    def reset(self):
        """Reset the game environment."""
        self.score = 0
        self.steps = 0
        self.clear_board()
        self.snake = Snake(game=self)
        self.set_apple_pos()

    # ----------------------------------------------------> BOARD <--------------------------------------------------- #

    def create_board(self):
        """Create the initial board, which consists out of two layers: (snake+wall) and apple."""
        layers = 2 if self.apple_separate else 1
        board = np.zeros((self.width, self.height, layers))

        # Add walls
        board[0, :, 0] = np.ones((self.width,)) * -1
        board[-1, :, 0] = np.ones((self.width,)) * -1
        board[:, 0, 0] = np.ones((self.height,)) * -1
        board[:, -1, 0] = np.ones((self.height,)) * -1

        # Add snake and apple
        for p in self.snake.body: board[p.x, p.y, 0] = -1
        board[self.apple.x, self.apple.y, 1 if self.apple_separate else 0] = 1
        return board

    def clear_board(self):
        """Clear the snake positions from the board."""
        self.board[1:-1, 1:-1, :] *= 0

    def update_board(self):
        """Update the position of the snake and apple."""
        self.clear_board()
        for p in self.snake.body: self.board[p.x, p.y, 0] = -1
        self.board[self.apple.x, self.apple.y, 1 if self.apple_separate else 0] = 1

    def show_board(self):
        """Print out the board."""
        for row in reversed(range(self.width)):
            for col in range(self.height):
                if self.board[col, row, 0] == -1:
                    print(" # ", end="")
                elif self.board[col, row, 0] == 1:
                    print(" o ", end="")
                else:
                    print("   ", end="")
            print()
        print("---" * self.width)

    def get_board_relative(self):
        """Transform the board position to first person viewing."""
        # Get board dimensions
        width = len(self.board)
        height = len(self.board[0])

        # Center board
        delta_width = int(width / 2) - self.snake.body[0][0]
        delta_height = int(height / 2) - self.snake.body[0][1]
        board = np.roll(self.board, delta_width, axis=0)
        board = np.roll(board, delta_height, axis=1)

        # Rotate board
        if self.snake.direction == RIGHT: board = np.rot90(board, 1)
        if self.snake.direction == DOWN: board = np.rot90(board, 2)
        if self.snake.direction == LEFT: board = np.rot90(board, 3)
        return board  # shape = (height, width, depth)

    # ----------------------------------------------------> APPLE <--------------------------------------------------- #

    def set_apple_pos(self):
        """Get random free location in the game."""
        self.apple = None
        while not self.apple or self.apple in self.snake.body:
            self.apple = Pos(x=randrange(1, self.width - 1),
                             y=randrange(1, self.height - 1))


# -*- coding: utf-8 -*-
"""
Created on Mon Apr 27 14:15:06 2020

@author: Shreyansh Satvik
"""
from keras.models import load_model
from glob import glob
import argparse
import os
import matplotlib.pyplot as plt
from matplotlib.pyplot import imshow
import scipy.io
import scipy.misc
import numpy as np
import pandas as pd
import PIL
import tensorflow as tf
from keras.layers import Input, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D
from keras.layers import AveragePooling2D, MaxPooling2D, Dropout, GlobalMaxPooling2D, GlobalAveragePooling2D
from keras import backend as K
from keras.layers import Input, Lambda, Conv2D
from keras.models import load_model, Model
from keras.models import model_from_json
from yolo_utils import read_classes, read_anchors, generate_colors, preprocess_image, draw_boxes, scale_boxes
from yad2k.models.keras_yolo import  yolo_boxes_to_corners, preprocess_true_boxes, yolo_loss, yolo_body,yolo_head
import h5py
from keras.models import Sequential
from keras.applications.vgg16 import VGG16
from PIL import Image
from keras.models import Sequential
from utils import load_weights, Box, yolo_net_out_to_car_boxes, draw_box
import cv2
import os

from configparser import ConfigParser

def yolo_filter_boxes(box_confidence, boxes, box_class_probs, threshold = .6):

    # Step 1: Compute box scores
    ### START CODE HERE ### (≈ 1 line)
    box_scores = box_confidence*box_class_probs
    ### END CODE HERE ###

    # Step 2: Find the box_classes using the max box_scores, keep track of the corresponding score
    ### START CODE HERE ### (≈ 2 lines)
    box_classes = K.argmax(box_scores,axis=-1)
    box_class_scores = K.max(box_scores,axis=-1)
    ### END CODE HERE ###

    # Step 3: Create a filtering mask based on "box_class_scores" by using "threshold". The mask should have the
    # same dimension as box_class_scores, and be True for the boxes you want to keep (with probability >= threshold)
    ### START CODE HERE ### (≈ 1 line)
    filtering_mask = box_class_scores >= threshold
    ### END CODE HERE ###

    # Step 4: Apply the mask to box_class_scores, boxes and box_classes
    ### START CODE HERE ### (≈ 3 lines)
    scores = tf.boolean_mask(box_class_scores,filtering_mask)
    boxes = tf.boolean_mask(boxes,filtering_mask)
    classes = tf.boolean_mask(box_classes,filtering_mask)
    ### END CODE HERE ###

    return scores, boxes, classes

def yolo_non_max_suppression(scores, boxes, classes, max_boxes = 10, iou_threshold = 0.5):

    max_boxes_tensor = K.variable(max_boxes, dtype='int32')     # tensor to be used in tf.image.non_max_suppression()
    K.get_session().run(tf.variables_initializer([max_boxes_tensor])) # initialize variable max_boxes_tensor

    # Use tf.image.non_max_suppression() to get the list of indices corresponding to boxes you keep
    ### START CODE HERE ### (≈ 1 line)
    nms_indices = tf.image.non_max_suppression(boxes,scores,max_boxes_tensor,iou_threshold,name=None)
    ### END CODE HERE ###

    # Use K.gather() to select only nms_indices from scores, boxes and classes
    ### START CODE HERE ### (≈ 3 lines)
    scores = K.gather(scores,nms_indices)
    boxes = K.gather(boxes,nms_indices)
    classes = K.gather(classes,nms_indices)
    ### END CODE HERE ###

    return scores, boxes, classes

def yolo_eval(yolo_outputs, image_shape , max_boxes=10, score_threshold=.6, iou_threshold=.5):

    # Retrieve outputs of the YOLO model (≈1 line)
    box_confidence, box_xy, box_wh, box_class_probs = yolo_outputs

    # Convert boxes to be ready for filtering functions (convert boxes box_xy and box_wh to corner coordinates)
    boxes = yolo_boxes_to_corners(box_xy, box_wh)

    # Use one of the functions you've implemented to perform Score-filtering with a threshold of score_threshold (≈1 line)
    scores, boxes, classes = yolo_filter_boxes(box_confidence, boxes, box_class_probs, score_threshold)

    # Scale boxes back to original image shape.
    boxes = scale_boxes(boxes, image_shape)

    # Use one of the functions you've implemented to perform Non-max suppression with
    # maximum number of boxes set to max_boxes and a threshold of iou_threshold (≈1 line)
    scores, boxes, classes = yolo_non_max_suppression(scores, boxes, classes, max_boxes , iou_threshold)

    ### END CODE HERE ###

    return scores, boxes, classes

"""
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(128, 128, 3)))
model.add(BatchNormalization(axis=3,name='bn0'))
model.add(Conv2D(16, (3, 3),strides=(1,1), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(BatchNormalization(axis=3,name='bn1'))

model.add(Conv2D(128, (2, 2), activation='relu'))
model.add(BatchNormalization(axis=3,name='bn2'))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(BatchNormalization(axis=3,name='bn3'))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

"""

####################################
yolo_model = load_model("yolonew.h5")
####################################################
sess = K.get_session()
class_names = read_classes("coco_classes(2).txt")
anchors = read_anchors("yolo_anchors.txt")

image_shape = (1280., 720.)
#new_model=yolo_model.layers.pop()
yolo_outputs = yolo_head(yolo_model.output, anchors, len(class_names))

scores, boxes, classes = yolo_eval(yolo_outputs, image_shape)

def predict(sess, image_file):
    """
    Runs the graph stored in "sess" to predict boxes for "image_file". Prints and plots the predictions.

    Arguments:
    sess -- your tensorflow/Keras session containing the YOLO graph
    image_file -- name of an image stored in the "images" folder.

    Returns:
    out_scores -- tensor of shape (None, ), scores of the predicted boxes
    out_boxes -- tensor of shape (None, 4), coordinates of the predicted boxes
    out_classes -- tensor of shape (None, ), class index of the predicted boxes

    Note: "None" actually represents the number of predicted boxes, it varies between 0 and max_boxes.
    """

    # Preprocess your image
    image, image_data = preprocess_image("images/" + image_file, model_image_size = (608, 608))

    # Run the session with the correct tensors and choose the correct placeholders in the feed_dict.
    # You'll need to use feed_dict={yolo_model.input: ... , K.learning_phase(): 0})
    ### START CODE HERE ### (≈ 1 line)
    out_scores, out_boxes, out_classes = sess.run(fetches=[scores,boxes,classes],feed_dict={yolo_model.input:image_data,
                  K.learning_phase():0})
    ### END CODE HERE ###

    # Print predictions info
    print('Found {} boxes for {}'.format(len(out_boxes), image_file))
    # Generate colors for drawing bounding boxes.
    colors = generate_colors(class_names)
    # Draw bounding boxes on the image file
    draw_boxes(image, out_scores, out_boxes, out_classes, class_names, colors)
    # Save the predicted bounding box on the image
    image.save(os.path.join("out", image_file), quality=90)
    # Display the results in the notebook
    output_image = scipy.misc.imread(os.path.join("out", image_file))
    imshow(output_image)

    return out_scores, out_boxes, out_classes

    out_scores, out_boxes, out_classes = predict(sess, "person2.jpg")


from typing import List
import tensorflow as tf
import numpy as np
import lunzi as lz
from lunzi import nn, rl

from boots import *
from boots.envs.batched_env import BatchedEnv
from boots.partial_envs import make_env, FLAGS as EnvFLAGS
from boots.normalizer import Normalizers
from boots.envs.virtual_env import VirtualEnv
from boots.utils import get_tf_config

class FLAGS(lz.BaseFLAGS):
    model = DeterministicModel.FLAGS
    env = EnvFLAGS
    SAC = SoftActorCritic.FLAGS
    plan = CEMPlanner.FLAGS

    n_test_samples = 10000
    n_train_samples = 1
    n_dev_samples = 1
    max_buf_size = 1000000
    n_dev_size = 2048
    n_model_ensemble = 1
    use_normalizer = True
    saving_schedule = 'exp'

    n_expl_samples = 10000

    class batch(lz.BaseFLAGS):
        n_model_grad_iters = 1
        n_policy_grad_iters = 1

        n_real = 64
        n_starts = 64
        horizon = 3

    n_iters = 1000000
    n_eval_iters = 10000
    n_save_iters = 10000

    policy_hidden_sizes = [32, 32]
    qfn_hidden_sizes = [256, 256]

    gamma = 0.99

    normalizer_decay = 1.0

    @classmethod
    def finalizer(cls):
        cls.model.n_dev_iters = 1000
        cls.SAC.batch_size = cls.batch.n_real + cls.batch.n_starts * cls.batch.horizon

@FLAGS.inject
def make_real_runner(*, n_envs, env):
    return rl.Runner(BatchedEnv([make_env] * n_envs), max_steps=env['max_steps'])

@FLAGS.inject
def evaluate(step, settings, tag, *, n_test_samples, _log, _writer):
    for runner, policy, name in settings:
        runner.reset()
        _, ep_infos = runner.run(policy, n_test_samples)
        returns = np.array([ep_info['return'] for ep_info in ep_infos])
        mean, std = np.mean(returns), np.std(returns)
        _log.info(f'# {step}, Tag = {tag}, Reward on {name} ({len(returns)} episodes): {mean:.6f} ± {std:.6f}')
        _writer.add_scalar(f'{tag}/{name}/reward/mean', mean, global_step=step)
        _writer.add_scalar(f'{tag}/{name}/reward/std', std, global_step=step)
        _writer.add_scalar(f'{tag}/{name}/reward/n', len(returns), global_step=step)

@lz.main(FLAGS)
@FLAGS.inject
def train_mb_sac(*, SAC, _log, _fs):
    tf.Session(config=get_tf_config()).__enter__()

    env = make_env()
    dim_state = int(np.prod(env.observation_space.shape))
    dim_action = int(np.prod(env.action_space.shape))
    env.verify()

    normalizers = Normalizers(dim_action=dim_action, dim_state=dim_state, decay=FLAGS.normalizer_decay)
    net_normalizer = normalizers.state if FLAGS.use_normalizer else None

    dtype = rl.gen_dtype(env, 'state action next_state reward done timeout')
    datasets = {
        'train': lz.ExtendableDataset(dtype, FLAGS.max_buf_size),
        'dev': lz.ExtendableDataset(dtype, FLAGS.n_dev_size * 2),
    }

    policy = TanhGaussianMLPPolicy(dim_state, dim_action, FLAGS.policy_hidden_sizes, normalizer=net_normalizer)
    det_policy = DetTanhPolicy(policy)
    models = []
    model_trainers = []
    for i in range(FLAGS.n_model_ensemble):
        model = DeterministicModel(dim_state, dim_action, normalizers)
        model_trainer = DeterministicModelTrainer(model, dim_state, dim_action, normalizers, datasets=datasets)
        models.append(model)
        model_trainers.append(model_trainer)
    ensemble = EnsembleModel(models, dim_state)

    virt_env = VirtualEnv(ensemble, make_env(), FLAGS.batch.n_starts)
    virt_runner = rl.Runner(virt_env, FLAGS.env.max_steps)
    # if FLAGS.batch.horizon > 0:
    #     virt_runner = VirtualRunner(dim_state, ensemble, det_policy, FLAGS.batch.horizon, dtype).build()
    qfns = nn.ModuleList([
        MLPQFunction(dim_state, dim_action, FLAGS.qfn_hidden_sizes, normalizer=net_normalizer),
        MLPQFunction(dim_state, dim_action, FLAGS.qfn_hidden_sizes, normalizer=net_normalizer),
    ])
    min_qfn = MinQFunction(dim_state, dim_action, qfns).build()

    algo = SoftActorCritic(qfns, policy, dim_state, dim_action)

    tf.get_default_session().run(tf.global_variables_initializer())

    planner = CEMPlanner(dim_state, dim_action, model, policy, min_qfn).build()
    behavior_policy = policy

    runners = {
        'collect': make_real_runner(n_envs=1),
        'dev': make_real_runner(n_envs=1),
        'eval/policy': make_real_runner(n_envs=4),
        'eval/det-policy': make_real_runner(n_envs=4),
        'eval/boots': make_real_runner(n_envs=1),
        'expand': virt_runner,
    }
    settings = [
        (runners['eval/policy'], policy, 'policy'),
        (runners['eval/det-policy'], det_policy, 'det-policy'),
        (runners['eval/boots'], planner, 'boots'),
    ]

    saver = nn.ModuleDict({'policy': policy, 'model': ensemble, 'qfns': qfns})
    _log.info(saver)

    _log.warning("Collect random trajectory...")
    datasets['train'].extend(runners['collect'].run(UniformPolicy(dim_action), FLAGS.n_expl_samples)[0])
    datasets['dev'].extend(runners['dev'].run(UniformPolicy(dim_action), FLAGS.n_expl_samples)[0])
    normalizers.update(datasets['train'][:FLAGS['n_expl_samples']])

    _log.warning("Pretraining model...")
    for model_trainer in model_trainers:
        model_trainer.pretrain()

    _log.warning("Start training...")
    for T in range(FLAGS.n_expl_samples, FLAGS.n_iters):
        if T % FLAGS.n_eval_iters == 0:
            evaluate(T, settings, 'episode')

        # collect data
        recent_train_set, ep_infos = runners['collect'].run(behavior_policy, FLAGS.n_train_samples)
        datasets['train'].extend(recent_train_set)
        datasets['dev'].extend(runners['dev'].run(policy, FLAGS.n_dev_samples)[0])

        normalizers.update(recent_train_set)

        n_policy_iters = FLAGS.batch.n_policy_grad_iters
        n_model_iters = (FLAGS.batch.n_model_grad_iters + n_policy_iters - 1) // n_policy_iters
        for _ in range(n_policy_iters):
            for _ in range(n_model_iters):
                for model_trainer in model_trainers:
                    model_trainer.step()

            real_samples = datasets['train'].sample(FLAGS.batch.n_real)
            if FLAGS.batch.horizon > 0:
                fake_starts = datasets['train'].sample(FLAGS.batch.n_starts)['state']
                runners['expand'].set_state(fake_starts)
                rollout_policy = policy
                fake_samples = runners['expand'].run(rollout_policy, FLAGS.batch.n_starts * FLAGS.batch.horizon)[0]
                samples = np.concatenate([real_samples, fake_samples]).astype(fake_samples.dtype).view(lz.Dataset)
            else:
                samples = real_samples

            algo.train(samples)

        if T % FLAGS.n_save_iters == 0:
            t = T // FLAGS.n_save_iters
            _fs.save('$LOGDIR/final.npy', saver.state_dict())
            if FLAGS.saving_schedule == 'exp' and t & (t - 1) == 0:
                _fs.save(f'$LOGDIR/stage-{T}.npy', saver.state_dict())
            elif FLAGS.saving_schedule == 'every':
                _fs.save(f'$LOGDIR/stage-{T}.npy', saver.state_dict())

if __name__ == '__main__':
    # with tf.device('/device:GPU:0'):
    train_mb_sac()


import logging

import numpy as np
import ray
from ray.rllib.utils.sgd import standardized
from ray.rllib.agents import with_common_config
from ray.rllib.agents.mbmpo.mbmpo_torch_policy import MBMPOTorchPolicy
from ray.rllib.agents.trainer_template import build_trainer
from ray.rllib.evaluation.metrics import get_learner_stats
from ray.rllib.execution.common import STEPS_SAMPLED_COUNTER, \
    STEPS_TRAINED_COUNTER, LEARNER_INFO, _get_shared_metrics
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.execution.metric_ops import CollectMetrics
from ray.util.iter import from_actors
from ray.rllib.agents.mbmpo.model_ensemble import DynamicsEnsembleCustomModel
from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
from ray.rllib.utils.torch_ops import convert_to_torch_tensor
from ray.rllib.evaluation.metrics import collect_episodes
from ray.rllib.agents.mbmpo.model_vector_env import custom_model_vector_env
from ray.rllib.evaluation.metrics import collect_metrics
from ray.rllib.agents.mbmpo.utils import calculate_gae_advantages, \
    MBMPOExploration

logger = logging.getLogger(__name__)

# yapf: disable
# __sphinx_doc_begin__
DEFAULT_CONFIG = with_common_config({
    # If true, use the Generalized Advantage Estimator (GAE)
    # with a value function, see https://arxiv.org/pdf/1506.02438.pdf.
    "use_gae": True,
    # GAE(lambda) parameter
    "lambda": 1.0,
    # Initial coefficient for KL divergence
    "kl_coeff": 0.0005,
    # Size of batches collected from each worker
    "rollout_fragment_length": 200,
    # Stepsize of SGD
    "lr": 1e-3,
    # Share layers for value function
    "vf_share_layers": False,
    # Coefficient of the value function loss
    "vf_loss_coeff": 0.5,
    # Coefficient of the entropy regularizer
    "entropy_coeff": 0.0,
    # PPO clip parameter
    "clip_param": 0.5,
    # Clip param for the value function. Note that this is sensitive to the
    # scale of the rewards. If your expected V is large, increase this.
    "vf_clip_param": 10.0,
    # If specified, clip the global norm of gradients by this amount
    "grad_clip": None,
    # Target value for KL divergence
    "kl_target": 0.01,
    # Whether to rollout "complete_episodes" or "truncate_episodes"
    "batch_mode": "complete_episodes",
    # Which observation filter to apply to the observation
    "observation_filter": "NoFilter",
    # Number of Inner adaptation steps for the MAML algorithm
    "inner_adaptation_steps": 1,
    # Number of MAML steps per meta-update iteration (PPO steps)
    "maml_optimizer_steps": 8,
    # Inner Adaptation Step size
    "inner_lr": 1e-3,
    # Horizon of Environment (200 in MB-MPO paper)
    "horizon": 200,
    # Dynamics Ensemble Hyperparameters
    "dynamics_model": {
        "custom_model": DynamicsEnsembleCustomModel,
        # Number of Transition-Dynamics Models for Ensemble
        "ensemble_size": 5,
        # Hidden Layers for Model Ensemble
        "fcnet_hiddens": [512, 512, 512],
        # Model Learning Rate
        "lr": 1e-3,
        # Max number of training epochs per MBMPO iter
        "train_epochs": 500,
        # Model Batch Size
        "batch_size": 500,
        # Training/Validation Split
        "valid_split_ratio": 0.2,
        # Normalize Data (obs, action, and deltas)
        "normalize_data": True,
    },
    "exploration_config": {
        "type": MBMPOExploration,
    },
    # Workers sample from dynamics models
    "custom_vector_env": custom_model_vector_env,
    # How many iterations through MAML per MBMPO iteration
    "num_maml_steps": 10,
})
# __sphinx_doc_end__
# yapf: enable

# Select Metric Keys for MAML Stats Tracing
METRICS_KEYS = [
    "episode_reward_mean", "episode_reward_min", "episode_reward_max"
]

class MetaUpdate:
    def __init__(self, workers, num_steps, maml_steps, metric_gen):
        """Computes the MetaUpdate step in MAML, adapted for MBMPO
        for multiple MAML Iterations

        Args:
            workers (WorkerSet): Set of Workers
            num_steps (int): Number of meta-update steps per MAML Iteration
            maml_steps (int): MAML Iterations per MBMPO Iteration
            metric_gen (Iterator): Generates metrics dictionary

        Returns:
            metrics (dict): MBMPO metrics for logging
        """
        self.workers = workers
        self.num_steps = num_steps
        self.step_counter = 0
        self.maml_optimizer_steps = maml_steps
        self.metric_gen = metric_gen
        self.metrics = {}

    def __call__(self, data_tuple):
        """Args:
            data_tuple (tuple): 1st element is samples collected from MAML
            Inner adaptation steps and 2nd element is accumulated metrics
        """
        # Metaupdate Step
        print("Meta-Update Step")
        samples = data_tuple[0]
        adapt_metrics_dict = data_tuple[1]
        self.postprocess_metrics(
            adapt_metrics_dict, prefix="MAMLIter{}".format(self.step_counter))

        # MAML Meta-update
        for i in range(self.maml_optimizer_steps):
            fetches = self.workers.local_worker().learn_on_batch(samples)
        fetches = get_learner_stats(fetches)

        # Update KLS
        def update(pi, pi_id):
            assert "inner_kl" not in fetches, (
                "inner_kl should be nested under policy id key", fetches)
            if pi_id in fetches:
                assert "inner_kl" in fetches[pi_id], (fetches, pi_id)
                pi.update_kls(fetches[pi_id]["inner_kl"])
            else:
                logger.warning("No data for {}, not updating kl".format(pi_id))

        self.workers.local_worker().foreach_trainable_policy(update)

        # Modify Reporting Metrics
        metrics = _get_shared_metrics()
        metrics.info[LEARNER_INFO] = fetches
        metrics.counters[STEPS_TRAINED_COUNTER] += samples.count

        if self.step_counter == self.num_steps - 1:
            td_metric = self.workers.local_worker().foreach_policy(
                fit_dynamics)[0]

            # Sync workers with meta policy
            self.workers.sync_weights()

            # Sync TD Models with workers
            sync_ensemble(self.workers)
            sync_stats(self.workers)

            metrics.counters[STEPS_SAMPLED_COUNTER] = td_metric[
                STEPS_SAMPLED_COUNTER]

            # Modify to CollectMetrics
            res = self.metric_gen.__call__(None)
            res.update(self.metrics)
            self.step_counter = 0
            print("MB-MPO Iteration Completed")
            return [res]
        else:
            print("MAML Iteration {} Completed".format(self.step_counter))
            self.step_counter += 1

            # Sync workers with meta policy
            print("Syncing Weights with Workers")
            self.workers.sync_weights()
            return []

    def postprocess_metrics(self, metrics, prefix=""):
        """Appends prefix to current metrics

        Args:
            metrics (dict): Dictionary of current metrics
            prefix (str): Prefix string to be appended
        """
        for key in metrics.keys():
            self.metrics[prefix + "_" + key] = metrics[key]

def post_process_metrics(prefix, workers, metrics):
    """Update Current Dataset Metrics and filter out specific keys

    Args:
        prefix (str): Prefix string to be appended
        workers (WorkerSet): Set of workers
        metrics (dict): Current metrics dictionary
    """
    res = collect_metrics(remote_workers=workers.remote_workers())
    for key in METRICS_KEYS:
        metrics[prefix + "_" + key] = res[key]
    return metrics

def inner_adaptation(workers, samples):
    # Each worker performs one gradient descent
    for i, e in enumerate(workers.remote_workers()):
        e.learn_on_batch.remote(samples[i])

def fit_dynamics(policy, pid):
    return policy.dynamics_model.fit()

def sync_ensemble(workers):
    """Syncs dynamics ensemble weights from main to workers

    Args:
        workers (WorkerSet): Set of workers, including main
    """

    def get_ensemble_weights(worker):
        policy_map = worker.policy_map
        policies = policy_map.keys()

        def policy_ensemble_weights(policy):
            model = policy.dynamics_model
            return {
                k: v.cpu().detach().numpy()
                for k, v in model.state_dict().items()
            }

        return {
            pid: policy_ensemble_weights(policy)
            for pid, policy in policy_map.items() if pid in policies
        }

    def set_ensemble_weights(policy, pid, weights):
        weights = weights[pid]
        weights = convert_to_torch_tensor(weights, device=policy.device)
        model = policy.dynamics_model
        model.load_state_dict(weights)

    if workers.remote_workers():
        weights = ray.put(get_ensemble_weights(workers.local_worker()))
        set_func = ray.put(set_ensemble_weights)
        for e in workers.remote_workers():
            e.foreach_policy.remote(set_func, weights=weights)

def sync_stats(workers):
    def get_normalizations(worker):
        policy = worker.policy_map[DEFAULT_POLICY_ID]
        return policy.dynamics_model.normalizations

    def set_normalizations(policy, pid, normalizations):
        policy.dynamics_model.set_norms(normalizations)

    if workers.remote_workers():
        normalization_dict = ray.put(
            get_normalizations(workers.local_worker()))
        set_func = ray.put(set_normalizations)
        for e in workers.remote_workers():
            e.foreach_policy.remote(
                set_func, normalizations=normalization_dict)

def post_process_samples(samples, config):
    # Instead of using NN for value function, we use regression
    split_lst = []
    for sample in samples:
        indexes = np.asarray(sample["dones"]).nonzero()[0]
        indexes = indexes + 1

        reward_list = np.split(sample["rewards"], indexes)[:-1]
        observation_list = np.split(sample["obs"], indexes)[:-1]

        paths = []
        for i in range(0, len(reward_list)):
            paths.append({
                "rewards": reward_list[i],
                "observations": observation_list[i]
            })

        paths = calculate_gae_advantages(paths, config["gamma"],
                                         config["lambda"])

        advantages = np.concatenate([path["advantages"] for path in paths])
        sample["advantages"] = standardized(advantages)
        split_lst.append(sample.count)
    return samples, split_lst

# Similar to MAML Execution Plan
def execution_plan(workers, config):
    # Train TD Models
    workers.local_worker().foreach_policy(fit_dynamics)

    # Sync workers policy with workers
    workers.sync_weights()

    # Sync TD Models and normalization stats with workers
    sync_ensemble(workers)
    sync_stats(workers)

    # Dropping metrics from the first iteration
    episodes, to_be_collected = collect_episodes(
        workers.local_worker(),
        workers.remote_workers(), [],
        timeout_seconds=9999)

    # Metrics Collector
    metric_collect = CollectMetrics(
        workers,
        min_history=0,
        timeout_seconds=config["collect_metrics_timeout"])

    inner_steps = config["inner_adaptation_steps"]

    def inner_adaptation_steps(itr):
        buf = []
        split = []
        metrics = {}
        for samples in itr:
            print("Collecting Samples, Inner Adaptation {}".format(len(split)))
            # Processing Samples (Standardize Advantages)
            samples, split_lst = post_process_samples(samples, config)

            buf.extend(samples)
            split.append(split_lst)

            adapt_iter = len(split) - 1
            prefix = "DynaTrajInner_" + str(adapt_iter)
            metrics = post_process_metrics(prefix, workers, metrics)

            if len(split) > inner_steps:
                out = SampleBatch.concat_samples(buf)
                out["split"] = np.array(split)
                buf = []
                split = []

                yield out, metrics
                metrics = {}
            else:
                inner_adaptation(workers, samples)

    # Iterator for Inner Adaptation Data gathering (from pre->post adaptation)
    rollouts = from_actors(workers.remote_workers())
    rollouts = rollouts.batch_across_shards()
    rollouts = rollouts.transform(inner_adaptation_steps)

    # Metaupdate Step with outer combine loop for multiple MAML iterations
    train_op = rollouts.combine(
        MetaUpdate(workers, config["num_maml_steps"],
                   config["maml_optimizer_steps"], metric_collect))
    return train_op

def get_policy_class(config):
    return MBMPOTorchPolicy

def validate_config(config):
    config["framework"] = "torch"
    if config["framework"] != "torch":
        raise ValueError("MB-MPO not supported in Tensorflow yet!")
    if config["inner_adaptation_steps"] <= 0:
        raise ValueError("Inner Adaptation Steps must be >=1.")
    if config["maml_optimizer_steps"] <= 0:
        raise ValueError("PPO steps for meta-update needs to be >=0")
    if config["entropy_coeff"] < 0:
        raise ValueError("entropy_coeff must be >=0")
    if config["batch_mode"] != "complete_episodes":
        raise ValueError("truncate_episodes not supported")
    if config["num_workers"] <= 0:
        raise ValueError("Must have at least 1 worker/task.")

MBMPOTrainer = build_trainer(
    name="MBMPO",
    default_config=DEFAULT_CONFIG,
    default_policy=MBMPOTorchPolicy,
    get_policy_class=get_policy_class,
    execution_plan=execution_plan,
    validate_config=validate_config)


import sys
import numpy as np
from .Scope import Scope
from .Processor import Processor

# Designed to be run from the command line to turn a grammar into a .obj file
if __name__ == '__main__':
    if len(sys.argv) != 12:
         print("bad args")
    else:
        gp = Processor()
        coords = [float(coord) for coord in sys.argv[1:7]]
        scope = Scope.freshScope(np.array(coords[:3]),np.array(coords[3:6]))
        gp.grammarToManyObjs(sys.argv[7], sys.argv[8], scope, int(sys.argv[9]), int(sys.argv[10]), sys.argv[11])
        print("Completed building generation")