from nltk import CFG
from nltk import ChartParser # parse_cfg, ChartParser
from random import choice
import re
from enum import Enum, auto
from argparse import ArgumentParser
from os import listdir
from os.path import isfile, join
import os

this_dir = os.path.dirname(os.path.abspath(__file__))
name_segment_folder = join(this_dir, "../../name-segments/")

class EnumAutoName(Enum):
    # An enum where auto() will default to the enum name
    def _generate_next_value_(name, start, count, last_values):
        return name

    def __str__(self):
        return self.value

def StringToEnum(s):
    if s in [v.name for v in Name.NameBank]:
        return Name.NameBank[s]
    else:
        raise ValueError

def get_files_from_path(path, extension_filter=".txt"):
    files = []
    for f in listdir(path):
        full_f = join(path, f)
        if extension_filter in full_f:
            if isfile(full_f):
                files.append(f)
    return files

def get_available_namebanks_and_syllables():
    namebanks = get_available_namebanks()
    syllables = get_available_syllables()
    # return namebanks | syllables   # Nicer 3.9 Python syntax
    return {**namebanks, **syllables}

def get_available_syllables(where="syllables"):
    global name_segment_folder
    path = join(name_segment_folder,where)
    onlyfiles = get_files_from_path(path)
    available = {}
    for f in onlyfiles:
        f = f.replace(".txt","")
        f = f.replace("-female","")
        f = f.replace("-male","")
        for x in range(10):
            f = f.replace(f"-{x}","")
        f = f.capitalize()
        available[f] = auto()
    return available

def get_available_namebanks(where="forenames"):
    global name_segment_folder
    path = join(name_segment_folder,where)
    onlyfiles = get_files_from_path(path)
    available = {}
    for f in onlyfiles:
        f = f.replace(".txt","")
        f = f.replace("-female","")
        f = f.replace("-male","")
        f = f.capitalize()
        available[f] = auto()
    return available

def get_available_origins(where="nouns"):
    global name_segment_folder
    path = join(name_segment_folder,where)
    onlyfiles = get_files_from_path(path)
    available = {}
    for f in onlyfiles:
        f = f.replace(".txt","")
        f = f.replace("-female","")
        f = f.replace("-male","")
        f = f.capitalize()
        available[f] = auto()
    return available

class Name:
    class NameOrder(EnumAutoName):
        Eastern = auto()
        Forename_Only = auto()
        Surname_Only = auto()
        Western = "Western"

    namebank_values = get_available_namebanks_and_syllables()
    NameBank = EnumAutoName('NameBank', namebank_values)

    origin_values = get_available_origins()
    Origin = EnumAutoName('Origin', origin_values)

    class NameType(EnumAutoName):
        Forename = auto()
        Surname = auto()

    class Origin(EnumAutoName):
        Aquatic = auto()
        Desert = auto()
        Mountain = auto()
        Tundra = auto()
        Urban = auto()
        Forest = auto()
        Air = auto()

    def __init__(self):
        self.gender_male = False
        self.gender_female = False
        self.gender_neutral = False

        self.has_position = False
        self.order = Name.NameOrder.Western

class FileFetcher():
    def __init__(self):
        pass

    def get_gender_endings(self, config, always_neutral=False):
        e = []
        if config.gender_male:
            e.append("male")
        if config.gender_female:
            e.append("female")
        if config.gender_neutral or always_neutral:
            e.append("")
        if len(e) == 0:
            print("No Gender Selection. Defaulting to gender neutral")
            config.gender_neutral = True
            e.append("")
        return e

    def get_position_files(self, config):
        ges = self.get_gender_endings(config)
        pt = []
        for g in ges:
            g = f"-{g}" if g != "" else g
            pt.append(f'prefixes/positions{g}.txt')

        return pt

    def SyllableLength(self, namebank):
        global name_segment_folder
        path = join(name_segment_folder,"syllables")
        onlyfiles = get_files_from_path(path)
        unique_syllables = {}
        for f in onlyfiles:
            if namebank in f:
                f = f.replace("-female","")
                f = f.replace("-male","")
                unique_syllables[f] = True
        return len(unique_syllables.items())

class Grammar:
    def __init__(self, config):
        self.config = config
        self.obj = {}
        self.root = "S"
        self.ff = FileFetcher()

    def initialize(self):
        self.obj[self.root]= ["PRE", "CORE", "POST"]

        self.basic_tokens()

    def basic_tokens(self):
        self.obj["SPC"] = ["' '"]
        self.obj["OF"] = ["'o'", "'f'"]

    def define_position(self, config, optional=False):

        # Prefix
        positions = self.ff.get_position_files(config)
        positions = [f"['{p}']" for p in positions]
        self.obj["PRE"] = ["TITLE", "SPC"]
        if optional:
            self.obj["PRE"].append(None)

        self.obj["TITLE"] = positions

        # Postfix
        origin = config.origin.name.lower()

        self.obj["POST"] = ["SPC", "OF", "SPC", "WHERE"]
        if optional:
            self.obj["POST"].append(None)
        # TODO: Allow multiple origins
        self.obj["WHERE"] = [f"['postfixes/{origin}.txt']",]

    def setNameOrder(self, order):
        if order == Name.NameOrder.Western:
            self.obj["CORE"] = ["FORENAME", "SPC", "SURNAME"]
        elif order == Name.NameOrder.Eastern:
            self.obj["CORE"] = ["SURNAME", "SPC", "FORENAME"]
        elif order == Name.NameOrder.Forename_Only:
            self.obj["CORE"] = ["FORENAME"]
        elif order == Name.NameOrder.Surname_Only:
            self.obj["CORE"] = ["FORENAME"]
        else:
            print("Unimplemented Name Order: ", order, ". Defaulting to Western")
            self.setNameOrder(Name.NameOrder.Western)

    def getNamesFromSyllables(self, config, name_type):
        ges = self.ff.get_gender_endings(config)
        namebank = config.namebank.name.lower()
        name_type = name_type.name.upper()

        global name_segment_folder

        # TODO: Check compatibile with namebanks
        syls = self.ff.SyllableLength(namebank)
        self.obj[name_type] = []
        for x in range(syls):
            self.obj[name_type].append(f"SYLLABLE{x}")

        for x in range(syls):
            pt = []
            for g in ges:
                g = f"-{g}" if g != "" else g
                f = f'syllables/{namebank}{g}-{x}.txt'
                if os.path.exists(join(name_segment_folder, f)):
                    pt.append(f)
                else:
                    print(f"Warn/Err: No syllable file found: {f}. May produce bad name.")
            self.obj[f"SYLLABLE{x}"] = [pt]

    def getNamesFromBank(self, config, name_type):
        ges = self.ff.get_gender_endings(config)
        namebank = config.namebank.name.lower()
        name_type = name_type.name.upper()

        pt = []
        for g in ges:
            g = f"-{g}" if g != "" else g
            # TODO: s shouldnt be there.
            pt.append(f'{name_type.lower()}s/{namebank}{g}.txt')
        self.obj[name_type] = [pt]

    def constructName(self, config, name_type):
        origin = config.origin.name.lower()
        name_type = name_type.name.upper()

        self.obj[name_type] = ["ADJ", "NOUN"]

        self.buildAdjBank(config)
        self.buildNounBank(config)

    def buildAdjBank(self, config):
        origin = config.origin.name.lower()

        pt = []
        # TODO: Dodginess/Alignment. John Bloodsword seems more evil than John Goldheart
        pt.append(f"['adjectives/{origin}.txt']")

        self.obj["ADJ"] = pt

    def buildNounBank(self, config):
        origin = config.origin.name.lower()

        pt = []
        # TODO: Dodginess/Alignment. John Poisonblood seems more evil than John Goldheart
        pt.append(f"['nouns/{origin}.txt']")
        self.obj["NOUN"] = pt

    def write(self, dir="", filename="custom.grammar"):
        # TODO: order carefully
        s = ""
        for key, value in self.obj.items():
            s += f"{key} -> "
            for i, v in enumerate(value):
                sep = " "
                if v is None:
                    v = " | "
                s += f"{v}{sep}"
            s += "\n"

        self.string_repr = s

        filename = os.path.join(dir, filename)
        print("---------------------")
        print(filename)
        print(filename)
        print(filename)
        print(filename)
        print("---------------------")

        f = open(filename, "w")
        f.write(s)
        f.close()
        filename = os.path.abspath(filename)
        return filename

    def __str__(self):
        if hasattr(self, "string_repr"):
            return self.string_repr
        else:
            return "Not Finalized"

def define_grammar(config, where=""):

    grammar = Grammar(config)
    grammar.initialize()
    if config.has_position:
        grammar.define_position(config)
    grammar.setNameOrder(config.order)

    # Prefer Forenames to be syllable generated
    if grammar.ff.SyllableLength(config.namebank.name.lower()) > 0:
        grammar.getNamesFromSyllables(config, Name.NameType.Forename)
    else:
        grammar.getNamesFromBank(config, Name.NameType.Forename)

    # TODO: Use namebank for Surnames
    grammar.constructName(config, Name.NameType.Surname)

    return grammar.write(where)

def resolve_grammar(G):
    def file_contents(s):
        global name_segment_folder
        filename = join(name_segment_folder,str(s.group(1)))

        try:
            terms = open(filename).readlines()
            s = ""
            for i, t in enumerate(terms):
                t = t.replace("\n","")
				# Allow Commenting
                if "#" not in t:
                    seperator = "|" if i > 0 else ""
                    s += f"{seperator} '{t}' "
        except FileNotFoundError:
            print("Warn/Err: File doesn't exist:", filename, ". May produce bad names.")
            s = ""
        return s

    G = re.sub(r"\[\'([a-zA-Z\-\.\/0-9]*)\'\]", file_contents, G)
    return G

def generate_name(G, ):
    grammar = CFG.fromstring(G)

    parser = ChartParser(grammar)

    gr = parser.grammar()
    tokens = produce(gr, gr.start())
    name = ''.join(tokens)
    return name.title()

def produce(grammar, symbol):
    words = []
    productions = grammar.productions(lhs = symbol)
    production = choice(productions)
    for sym in production.rhs():
        if isinstance(sym, str):
            words.append(sym)
        else:
            words.extend(produce(grammar, sym))
    return words

def generate(args, where=""):
    config = Name()
    config.has_position = True
    config.origin = args.origin
    config.namebank = args.namebank
    config.order = args.order
    config.gender_male = args.gender_male
    config.gender_female = args.gender_female
    config.gender_neutral = args.gender_neutral
    grammar_file = define_grammar(config, where)

    G = resolve_grammar(open(grammar_file).read())
    name = generate_name(G)

    if args.verbose:
        print("Your Character:", name)
    return name

def parse_args():
    ap = ArgumentParser(description="Generate a character name")

    # Gender
    ap.add_argument('--gender-male', action="store_true")
    ap.add_argument('--gender-female', action="store_true")
    ap.add_argument('--gender-neutral', action="store_true")

    # Origins
    ap.add_argument('--order', type=Name.NameOrder, choices=list(Name.NameOrder), nargs="?", default=Name.NameOrder.Western)
    ap.add_argument('--origin', type=Name.Origin, choices=list(Name.Origin), nargs="?", default=Name.Origin.Mountain)
    ap.add_argument('--namebank', type=Name.NameBank, choices=Name.NameBank, nargs="?", default=Name.NameBank.Dwarf)
    args = ap.parse_args()
    args.verbose = True
    return args

if __name__ == "__main__":
    a = parse_args()
    generate(a)


import numpy as np
import cv2
import sys
import argparse

# Creating the parser
ap = argparse.ArgumentParser()
ap.add_argument("-m", "--model", required=True,
	help="path to the model used to make the prediction and generate the class activation maps")
# Parsing the arguments
args = vars(ap.parse_args())

from utilities.gapModels import MobileNetGAP
from utilities.classifier import Classifier
from utilities.helpers import *

print("⏳" + BLUE + " Loading model ... " + END)

model = MobileNetGAP(path=args["model"])
clf = Classifier(model, name='mobilenet')

print("💾" + BLUE + " Model loaded." + END)

def addContours(input_img, output_img, draw_bounding_box=True, draw_contours=False, threshold=100):
    """
	>>> Work In Progress <<<
    Detects the bounding boxes and/or contours in the input image and adds them to the output image
    Returns the modified output_img
	>>> Work In Progress <<<
    """
	# Convert image to gray
    gray = cv2.cvtColor(input_img, cv2.COLOR_BGR2GRAY)

    # Threshold the image
    _, threshed_img = cv2.threshold(gray, threshold, 255, cv2.THRESH_BINARY)

    # Get the external contours
    _, contours, _ = cv2.findContours(threshed_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if draw_contours:
        # Draw the contours
        cv2.drawContours(output_img , contours, -1, (0, 255, 0), 5)

    if draw_bounding_box:
        # Draw the bounding boxes
        for c in contours:
            # Get the bounding rectangle
            x, y, w, h = cv2.boundingRect(c)
            # Draw it
            cv2.rectangle(output_img, (x, y), (x + w, y + h), (255, 0, 0), 3)

    return output_img

def show_detection(img, prediction):
	height, width, _ = img.shape
	middle = (int(height//2), int(width//2))
	if prediction>0.6: cv2.rectangle(img,(10,10),(width-10, height-10),(255, 255, 255),thickness=40)
	cv2.rectangle(img,(0,0),(width, 40),(56, 38, 50),thickness=-1)
	cv2.rectangle(img,(0,0),(int(width*prediction), 40),(118, 230, 0),thickness=-1)

	return img

cap = cv2.VideoCapture(0)

while 1:
	ret, img = cap.read()

	# Get the cam and prediction made by the model
	cam, prediction = clf.cam(img, class_number=1)

	# Detect the contours and or bounding boxes in the cam
	# img = addContours(input_img=cam, output_img=img, draw_bounding_box=True, draw_contours=False, threshold=100)

	# Add the cam to the original image
	img = cv2.addWeighted(cam, 0.5, img, 0.8, 0)

	# Indicators of the probability of presence of a human
	img = show_detection(img, prediction[1])

	cv2.imshow('img',img)

	if cv2.waitKey(30) & 0xFF == ord('q'):
		break

cap.release()
cv2.destroyAllWindows()


# Written by P. Xydi, Feb 2022

######################################
# Import libraries
######################################

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

import matplotlib.cm as cm
color_1 = cm.get_cmap("Set2")(2) # set blue
color_2 = cm.get_cmap("Set2")(1) # set orange
from sklearn.metrics import ConfusionMatrixDisplay

######################################
def sentence_distribution(dataset, label = 'training', to_plot = False):
    '''
    Plots the distribution of sentences in a given dataset

    INPUTS:
    - dataset:  list, list of samples
    - label:    str, used in the title of the output plot

    OUTPUT:
    - histogram of the distribution of sentences in dataset
    - nbr_sents : list, number of sentences per sample in dataset

    '''
######################################

    # Create empty list
    nbr_sents = []

    for i in range(len(dataset)):
        nbr_sents.append(len(dataset[i]))

    if to_plot:

        # Plot the sentence distibution

        # Barplot and font specifications
        barplot_specs = {"color": color_1, "alpha": 0.7, "edgecolor": "grey"}
        label_specs   = {"fontsize": 12}
        title_specs   = {"fontsize": 14, "fontweight": "bold", "y": 1.03}

        plt.figure(figsize=(8,4))

        plt.hist(nbr_sents, bins = 20, **barplot_specs)
        plt.xlabel('Nbr of sentences per sample', **label_specs)
        plt.ylabel('Nbr of samples',**label_specs)
        plt.title('Distribution of sentences in {} set'.format(label),**title_specs)
        plt.show()

    return np.sum(nbr_sents)

######################################
def plot_token_distribution(dataset, label):
    '''
    Plots the distribution of tokens in the sentences of a
    given dataset.

    INPUTS:
    - dataset:  list, list of samples
    - label:    str, used in the title of the output plot

    OUTPUT:
    - histogram of the distribution of tokens in dataset
    - nbr_tokens: list, number of tokens per sentence in dataset

    '''
######################################

    # Create empty list
    nbr_tokens = []

    # Count tokens in sentences and append to list
    for i in range(len(dataset)):
        for j in range(len(dataset[i])):
            nbr_tokens.append(len(dataset[i][j]))

    # Plot the sentence distibution

    # Barplot and font specifications
    barplot_specs = {"color": "mediumpurple", "alpha": 0.7, "edgecolor": "grey"}
    label_specs   = {"fontsize": 12}
    title_specs   = {"fontsize": 14, "fontweight": "bold", "y": 1.03}

    plt.figure(figsize=(8,4))

    plt.hist(nbr_tokens, bins = 20, **barplot_specs)
    plt.xlabel('Nbr of tokens per sentence', **label_specs)
    plt.ylabel('Nbr of sentences',**label_specs)
    plt.title('Distribution of tokens in {} set'.format(label),**title_specs)
    plt.show()

    return nbr_tokens

######################################
def target_sample_distribution(labels):
    """
    Plots the distribution of samples in target variable (categorical)

    Input:
    - labels : list, list of target values

    """
######################################

    w = pd.value_counts(labels)

    # Barplot and font specifications
    barplot_specs = {"color": color_2, "alpha": 0.7, "edgecolor": "grey"}
    label_specs   = {"fontsize": 12}
    title_specs   = {"fontsize": 14, "fontweight": "bold", "y": 1.02}

    plt.figure(figsize=(8,4.5))
    sns.barplot(x=w.index,y=w.values, **barplot_specs);
    plt.ylabel('Counts',**label_specs);
    plt.xticks(rotation=45)
    plt.yscale('log')
    plt.title('Sample distribution in target variable',**title_specs);

######################################
def plot_loss_accuracy_curves(history):
######################################

    title_specs = {"fontsize": 16}
    label_specs = {"fontsize": 14}

    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(12, 5))

    # Plot loss values
    ax1.set_title('Validation loss: {:.4f}'.format(history.history['val_loss'][-1]))
    ax1.plot(history.history['loss'], color =color_1, label='training set')
    ax1.plot(history.history['val_loss'], color =color_2, label='validation set')
    ax1.set_xlabel('Epochs',**label_specs)
    ax1.set_ylabel('Loss',**label_specs)
    ax1.set_ylim([0,None])
    ax1.legend()

    # plot accuracy values
    ax2.set_title('Validation accuracy: {:.2f}%'.format(history.history['val_accuracy'][-1]*100))
    ax2.plot(history.history['accuracy'], color =color_1, label='training set')
    ax2.plot(history.history['val_accuracy'],  color =color_2, label='validation set')
    ax2.set_xlabel('Epochs',**label_specs)
    ax2.set_ylabel('Accuracy',**label_specs)
    ax2.set_ylim([None,1])
    ax2.legend()

    plt.tight_layout()

######################################
def plot_confusion_matrix(y, y_pred, labels, suptitle):
######################################

    # Create two subplots
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

    title_specs   = {"fontsize": 14, "fontweight": "bold", "y": 1.03}
    plt.suptitle(suptitle,**title_specs)

    # Plots the standard confusion matrix
    ax1.set_title("Confusion Matrix (counts)", y= 1.02)

    ConfusionMatrixDisplay.from_predictions(y,
                                            y_pred,
                                            display_labels=labels,
                                            cmap=plt.cm.Blues,
                                            values_format='d',
                                            ax=ax1)

    ax1.set_xticklabels(labels = labels, rotation=90)

    # Plots the normalized confusion matrix
    ax2.set_title("Confusion Matrix (ratios)", y= 1.02)

    ConfusionMatrixDisplay.from_predictions(y,
                                            y_pred,
                                            normalize="true",
                                            display_labels=labels,
                                            cmap=plt.cm.Blues,
                                            values_format='.1g',
                                            ax=ax2)

    ax2.set_xticklabels(labels = labels, rotation=90)

    plt.tight_layout()


from models.pointnet import PointNetDenseCls
import torch
import torch.nn as nn
import torch.nn.functional as F
import hydra
import os
from datasets import kpnet
import logging
from itertools import combinations
import numpy as np
from tqdm import tqdm

def pdist(vectors):
    distance_matrix = -2 * vectors.mm(torch.t(vectors)) + vectors.pow(2).sum(dim=1).view(1, -1) + vectors.pow(2).sum(
        dim=1).view(-1, 1)
    return distance_matrix

class AverageMeter(object):
    """Computes and stores the average and current value
       Imported from https://github.com/pytorch/examples/blob/master/imagenet/main.py#L247-L262
    """
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

class PairSelector:
    """
    Implementation should return indices of positive pairs and negative pairs that will be passed to compute
    Contrastive Loss
    return positive_pairs, negative_pairs
    """

    def __init__(self):
        pass

    def get_pairs(self, embeddings, labels):
        raise NotImplementedError

# reference: https://github.com/adambielski/siamese-triplet
class HardNegativePairSelector(PairSelector):
    """
    Creates all possible positive pairs. For negative pairs, pairs with smallest distance are taken into consideration,
    matching the number of positive pairs.
    """

    def __init__(self, cpu=True):
        super(HardNegativePairSelector, self).__init__()
        self.cpu = cpu

    def get_pairs(self, embeddings, labels):
        if self.cpu:
            embeddings = embeddings.cpu()
        distance_matrix = pdist(embeddings)

        labels = labels.cpu().data.numpy()
        all_pairs = np.array(list(combinations(range(len(labels)), 2)))
        all_pairs = torch.LongTensor(all_pairs)
        positive_pairs = all_pairs[(
            labels[all_pairs[:, 0]] == labels[all_pairs[:, 1]]).nonzero()]
        negative_pairs = all_pairs[(
            labels[all_pairs[:, 0]] != labels[all_pairs[:, 1]]).nonzero()]

        negative_distances = distance_matrix[negative_pairs[:,
                                                            0], negative_pairs[:, 1]]
        negative_distances = negative_distances.cpu().data.numpy()
        top_negatives = np.argpartition(negative_distances, len(positive_pairs))[
            :len(positive_pairs)]
        top_negative_pairs = negative_pairs[torch.LongTensor(top_negatives)]

        return positive_pairs, top_negative_pairs

# reference: https://github.com/adambielski/siamese-triplet
class OnlineContrastiveLoss(nn.Module):
    """
    Online Contrastive loss
    Takes a batch of embeddings and corresponding labels.
    Pairs are generated using pair_selector object that take embeddings and targets and return indices of positive
    and negative pairs
    """

    def __init__(self, margin, pair_selector, mean_distance=None):
        super(OnlineContrastiveLoss, self).__init__()
        self.margin = margin
        self.pair_selector = pair_selector
        if mean_distance is not None:
            self.mean_distance = mean_distance[0].cuda()
        else:
            self.mean_distance = None

    def forward(self, embeddings, target):
        positive_pairs, negative_pairs = self.pair_selector.get_pairs(
            embeddings, target)
        if embeddings.is_cuda:
            positive_pairs = positive_pairs.cuda()
            negative_pairs = negative_pairs.cuda()
        positive_loss = (embeddings[positive_pairs[:, 0]] -
                         embeddings[positive_pairs[:, 1]]).pow(2).sum(1)

        labels_1 = tuple(target[negative_pairs[:, 0]].tolist())
        labels_2 = tuple(target[negative_pairs[:, 1]].tolist())
        label_pair = (labels_1, labels_2)

        if self.mean_distance is not None:
            negative_loss = F.relu(
                self.mean_distance[label_pair] - ((embeddings[negative_pairs[:, 0]] - embeddings[negative_pairs[:, 1]]).pow(2).sum(
                    1) + 1e-6).sqrt()).pow(2)
        else:
            negative_loss = F.relu(
                self.margin - ((embeddings[negative_pairs[:, 0]] - embeddings[negative_pairs[:, 1]]).pow(2).sum(
                    1) + 1e-6).sqrt()).pow(2)
        loss = torch.cat([positive_loss, negative_loss], dim=0)
        return loss.mean()

@hydra.main(config_path='config', config_name='config')
def main(cfg):
    logger = logging.getLogger(__name__)

    train_dataset = kpnet.KeypointDataset(cfg)
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, drop_last=True)

    model = PointNetDenseCls(feature_transform=True, cfg=cfg).cuda()

    logger.info('Start training on 3D embeddings')

    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=1e-3
    )
    criterion = OnlineContrastiveLoss(1., HardNegativePairSelector())

    meter = AverageMeter()
    for epoch in range(cfg.max_epoch + 1):
        train_iter = tqdm(train_dataloader)

        # Training
        meter.reset()
        model.train()
        for i, (pc, kp_idxs) in enumerate(train_iter):
            pc, kp_idxs = pc.cuda(), kp_idxs.cuda()
            outputs = model(pc.transpose(1, 2))

            embeddings = []
            labels = []
            for i in range(cfg.batch_size):
                embedding_model = outputs[i]
                keypoints = kp_idxs[i]
                for idx in range(len(keypoints)):
                    kp_idx = keypoints[idx]
                    if kp_idx < 0:
                        continue

                    embedding_kp = embedding_model[kp_idx]
                    embeddings.append(embedding_kp)
                    labels.append(idx)

            embeddings = torch.stack(embeddings)
            labels = torch.tensor(labels).cuda()

            loss = criterion(embeddings, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_iter.set_postfix(loss=loss.item())
            meter.update(loss.item())

        logger.info(
                f'Epoch: {epoch}, Average Train loss: {meter.avg}'
            )

        torch.save(model.state_dict(), f'epoch{epoch}.pth')

if __name__ == '__main__':
    main()


import sys

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits import mplot3d
import math

class Interpolator:
    def __init__(self):
        DEFAULT_STEP = 1  # 10 ** 5

        self.c = 0.9  # Smoothing factor
        self.e = 0.1  # sys.float_info.epsilon  # Really small number

        self.actions = []
        self.qualities = []
        self.knots_count = 0
        self.step = DEFAULT_STEP

    def distance(self, chosen_action, i, q_max):
        return np.linalg.norm(np.subtract(chosen_action, self.actions[i])) ** 2 + self.c * (
                q_max - self.qualities[i]) + self.e

    def wsum(self, chosen_action):
        output = 0
        q_max = max(self.qualities)
        for i in range(self.knots_count):
            output += self.qualities[i] / self.distance(chosen_action, i, q_max)
        return output

    def norm(self, chosen_action):
        output = 0
        q_max = max(self.qualities)
        for i in range(self.knots_count):
            output += 1 / self.distance(chosen_action, i, q_max)
        return output

    def get_quality(self, action):
        value = self.wsum(action) / self.norm(action)
        if math.isnan(value):
            return 0
        else:
            return value

    def update_function_2(self, action, quality, update_action=True):
        q = np.array(self.qualities)

        knot_count = len(q)

        optimal_action = action
        action = np.array(self.actions)
        Q_new = quality

        num = 0
        den = 0
        deriv_q = []
        deriv_u0 = []
        deriv_u1 = []

        for it in range(0, knot_count):
            weight = np.linalg.norm(optimal_action - action[it]) + self.c * (q.max() - q[it] + self.e)
            den = den + (1.0 / weight)
            num = num + (q[it] / weight)
            deriv_q.append((den * (weight + q[it] * self.c) - num * self.c) / pow((weight * den), 2))
            deriv_u0.append(((num - den * q[it]) * 2 * (action[it][0] - optimal_action[0])) / (pow(weight * den, 2)))
            deriv_u1.append(((num - den * q[it]) * 2 * (action[it][1] - optimal_action[1])) / (pow(weight * den, 2)))

        Q_dash = num / den
        error = Q_new - Q_dash

        for it in range(0, knot_count):
            q[it] = q[it] + error * deriv_q[it]
            action[it][0] = action[it][0] + error * deriv_u0[it]
            action[it][1] = action[it][1] + error * deriv_u1[it]

        if update_action:
            self.actions = action
        self.qualities = q

    def update_function(self, action, quality, update_action=False):
        knot_count = len(self.qualities)
        # print("qualities:", self.qualities)
        if type(self.qualities) == np.ndarray:
            self.qualities = self.qualities.tolist()
        if type(self.qualities[0]) == list:
            self.qualities = [e[0] for e in self.qualities]
        max_list = self.qualities + [float(quality)]

        q_max = max(max_list)
        for it in range(0, knot_count):
            self.qualities[it] += self.e * \
                                  (quality - self.qualities[it]) \
                                  / self.distance(action, it, q_max) ** 2

    def set_u(self, actions):
        self.actions = actions
        self.knots_count = len(self.actions)

    def set_q(self, qualities):
        self.qualities = qualities

    def set_step(self, step):
        self.step = step

    def get_u(self):
        return self.actions

    def get_q(self):
        return self.qualities

if __name__ == "__main__":
    from output_visualizer import OutputVisualizer
    import cv2

    u = []
    interpolator = Interpolator()
    for i in np.arange(-1, 1.1, 0.5):
        for j in np.arange(-1, 1.1, 0.5):
            u.append(np.array([i, j]))
    q = [0.04448929, 0.5086165, 0.76275706, -0.2851543, 0.39455223,
         -0.19585085, -0.52812827, 0.25080782, 0.4987614, 0.26595366,
         -0.3598364, 0.41622806, 0.10484912, -0.11532316, -0.11455766,
         -0.14297369, -0.04747943, 0.19820265, 0.5723205, 0.13500524,
         -0.24156858, 0.15854892, 0.22840545, 0.35542938, -0.5061423]

    visualizer = OutputVisualizer()
    visualizer.render(np.append(u, [[e] for e in q], axis=1))
    cv2.waitKey(3000)

    interpolator.set_q(q)
    interpolator.set_u(u)

    # for _ in range(5):
    #    interpolator.update_function_2(np.array([0, 0]), 2)  # , update_action=False)
    # interpolator.update_function(np.array([-1, 0]), 2)#, update_action=False)
    interpolator.update_function(np.array([-0.5, 1.0]), -0.6402964293956757)  # , update_action=False)

    q = interpolator.get_q()
    u = interpolator.get_u()
    visualizer.render(np.append(u, [[e] for e in q], axis=1))
    cv2.waitKey(3000)

    # print(interpolator.get_quality(np.array([0.75, 0])))
    '''
    fig = plt.figure()
    ax = plt.axes()  # projection="3d")

    X = []
    Y = []
    Z = []
    for throttle in np.arange(-1, 1.1, 0.1):
        for steering in np.arange(-1, 1.1, 0.1):
            X.append(throttle)
            Y.append(steering)
            Z.append(interpolator.get_quality(np.array([throttle, steering])))
    '''
    # ax.plot_trisurf(np.array(X), np.array(Y), np.array(Z), cmap=cm.bwr)

    # throttles = [a[0] for a in u]
    # steerings = [a[1] for a in u]
    # ax.plot_trisurf(np.array(throttles), np.array(steerings), np.array(q))

    # interpolator.update_function(np.array([1, 0]), 20)
    # interpolator.update_function(np.array([1, 0]), 20)

    '''
    X = []
    Y = []
    Z = []
    for throttle in np.arange(-1, 1.1, 0.1):
        for steering in np.arange(-1, 1.1, 0.1):
            X.append(throttle)
            Y.append(steering)
            Z.append(interpolator.get_quality(np.array([throttle, steering])))
    ax.plot_trisurf(np.array(X), np.array(Y), np.array(Z), cmap=cm.bwr)
    '''

    '''
    u = interpolator.get_u()
    q = np.reshape(interpolator.get_q(), (-1, 5))
    throttles = np.reshape([a[0] for a in u], (-1, 5))
    steerings = np.reshape([a[1] for a in u], (-1, 5))
    ax.contourf(np.array(throttles), np.array(steerings), np.array(q), cmap=cm.bwr)

    plt.show()
    '''


from Actors.Actor import Actor
from Actors.Maze import Maze
from Actors.Direction import Direction
from Util.Timer import Timer
import pyrr
import math
import numpy

class Pacman(Actor):

    # assume position is 2d vector (x cell, z cell)
    def __init__(self, position : list, direction : Direction, speed : float, id : str):

        super().__init__(position, direction, speed, id)

        self.animation = self.__PacmanAnimation()

        self.updateTimer = Timer()
        self.updateTimer.restart()

        self.isMoving = False

    def update(self, maze : Maze):

        deltaTime = self.updateTimer.getElapsedTime()
        distance = deltaTime * self.speed
        self.updateTimer.restart()

        # if self.turnSignal == None:
        #     pass

        if self.canApplyTurnSignal(maze):
            self.applyTurnSignal()

        currentCellPos = Maze.worldCoordsToCellCoords(self.position)
        if not maze.isVacantSpotInSpecifiedDirection(currentCellPos, self.currectDirection) and maze.isTheMiddleOfTheCell(self.position, self.currectDirection):
            self.isMoving = False
            return

        self.position += self.frontVector * pyrr.Vector3([distance, distance, distance])
        self.isMoving = True

    def notify(self, message : str):
        tokens = message.split('/')

        if tokens[0] == "turn":
            direction = tokens[1]

            if direction[0] == 'a':
                directionValue = Direction.LEFT
            elif direction[0] == 'w':
                directionValue = Direction.UP
            elif direction[0] == 'd':
                directionValue = Direction.RIGHT
            elif direction[0] == 's':
                directionValue = Direction.DOWN

            self.setTurnSignal(directionValue)

    def getLowerThenUpperJawRotations(self):
        return self.animation.getLowerThenUpperJawRotations()

    class __PacmanAnimation:

        def __init__(self):
            self.animationPeriod = 300
            self.amplitude = 60.0

            self.openMouse = True

            self.currentRotationLowerJaw = 0.0
            self.currentRotationUpperJaw = 0.0

            self.timer = Timer()
            self.timer.restart()

        def getLowerThenUpperJawRotations(self):

            deltaTime = self.timer.getElapsedTime()

            delta_x_degrees = (deltaTime * self.amplitude) / self.animationPeriod

            if self.currentRotationLowerJaw < 10.0:
                self.openMouse = True
            elif self.currentRotationLowerJaw > self.amplitude:
                self.openMouse = False

            if self.openMouse == True:
                self.currentRotationLowerJaw += delta_x_degrees
            else:
                self.currentRotationLowerJaw -= delta_x_degrees

            self.currentRotationUpperJaw = -self.currentRotationLowerJaw

            self.timer.restart()

            return self.currentRotationLowerJaw, self.currentRotationUpperJaw


import os
import pickle
import types
import shutil
import time
import multiprocessing as mtp

packageSample={"mode":"join", "from":"id123", "to":["id234", "id789"], "time":"20181012122123","type":"unknown", "tag":["dog", "white"], "dataSet":"image object"}

class cellComm(object):
    def __init__(self):
        self.package="package.pkl"
        self.pklTemp="temp.pkl"
        self.commScript="comm.py"
        self.cellScript="cell.py"
        self.database="cell.db"
        self.cwdir=os.path.dirname(__file__)
        self.pwdir=os.path.dirname(self.cwdir)

    #check if any file received to awake cell
    def listen(self):
        if os.path.exists(self.package):
            pfile=open(self.package, "rb")
            data=pickle.load(pfile)
            pfile.close()
            if type(data)==dict:
                self.awakeCell()
            os.remove(self.package)

    #return data passed on
    def recv(self):
        if os.path.exists(self.package):
            pfile=open(self.package, "rb")
            self.data=pickle.load(pfile)
            pfile.close()
            if type(self.data)==dict:
                return self.data
            else:
                os.remove(self.package)
                return False

    #send() function accepts packaged dictionary object
    def send(self, packageObject):
        self.packageObject=packageObject
        if type(self.packageObject)==dict:
            temp=open(self.pklTemp, "wb")
            pickle.dump(self.packageObject, temp)
            temp.close()
            for item in self.packageObject["to"]:
                shutil.copy(self.pklTemp, (item+r"/"+self.package))
                self.awakeComm(item+r"/"+self.commScript)
            os.remove(self.pklTemp)

    #awake "comm.py" of target cell to activate communication
    def awakeComm(self, desObject):
        p=mtp.Process(target=execfile(desObject))
        p.daemon=False
        p.start()

    #awake "cell.py" of target cell to activate it's body
    def awakeCell(self):
        p=mtp.Process(target=execfile(self.cellScript))
        p.daemon=False
        p.start()

    def run(self):
        self.listen()