# imports
import pandas as pd
from sklearn.model_selection import train_test_split
import random
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
from tqdm import tqdm, trange
import os
import torch.nn.functional as F
import random

SEED = 48
random.seed(48)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# define functions

# Generate a song
def text_generation(prompt_list):
    '''
    :param prompt_list: list
    :return: single song or list of songs; n returned = number of prompts given
    '''
    songs = []
    if len(prompt_list) > 1:
        for i in range(len(prompt_list)):
            prompt = prompt_list[i]
            song = generate(model.to('cpu'), tokenizer, prompt)
            songs.append(song)
        return songs
    else:
        song = generate(model.to('cpu'), tokenizer, prompt_list[0])
        return song

def generate(
        model,
        tokenizer,
        prompt,
        entry_length=250,  # maximum number of words
        top_p=0.8,
        temperature=1.,
):
    model.eval()
    generated_num = 0
    generated_list = []

    filter_value = -float("Inf")

    with torch.no_grad():

        entry_finished = False
        generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0)

        for i in range(entry_length):
            outputs = model(generated, labels=generated)
            loss, logits = outputs[:2]
            logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)

            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
                                                ..., :-1
                                                ].clone()
            sorted_indices_to_remove[..., 0] = 0

            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            logits[:, indices_to_remove] = filter_value

            next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
            generated = torch.cat((generated, next_token), dim=1)

            if next_token in tokenizer.encode("<|endoftext|>"):
                entry_finished = True

            if entry_finished:
                generated_num = generated_num + 1

                output_list = list(generated.squeeze().numpy())
                output_text = tokenizer.decode(output_list)
                generated_list.append(output_text)
                break

        if not entry_finished:
            output_list = list(generated.squeeze().numpy())
            output_text = f"{tokenizer.decode(output_list)}<|endoftext|>"
            generated_list.append(output_text)

    return generated_list

#---------LOAD MODEL--------------------
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
# new_tokens = ['<newline>', '<verse>', '<chorus>', '<prechorus>', '<bridge>', '<outro>', '<intro>', '<refrain>', '<hook>', '<postchorus>', '<other>']
new_tokens = ['<newline>', '<SONGBREAK>']
tokenizer.add_special_tokens({'additional_special_tokens': new_tokens}) # add tokens for verses
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.resize_token_embeddings(len(tokenizer)) # resize embeddings for added special tokens
model.load_state_dict(torch.load('model_2.pt'))
model = model.to(device)

songs = text_generation(['here i am thinking about you', 'hey what do you know', 'cold dark days are ahead', "i'm not thinking about"])
print(songs)


#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
Documentation: .py
  Classes and functions:
    N.A.
  Description:
    This class handles the menu which appears when the player hits the pause
    button inside a running game.
"""

__author__ = "Fireclaw the Fox"
__license__ = """
Simplified BSD (BSD 2-Clause) License.
See License.txt or http://opensource.org/licenses/BSD-2-Clause for more info
"""

from menus.Menu import Menu
from direct.gui.DirectGui import DirectButton
from panda3d.core import TextNode

class MenuPause(Menu):

    def __init__(self):
        Menu.__init__(self)
        self.frameMain.setTransparency(True)
        self.frameMain.frameColor = (0.25, 0, 1, 0.25)

        screenResMultiplier = 1.33

        self.btnReturn = DirectButton(
            # size of the button
            scale = (0.25, 0.25, 0.25),
            # some temp text
            text = "ABC",
            # size of the text
            text_scale = (0.5*screenResMultiplier, 0.5, 0.5),
            # set the alignment to center
            text_align = TextNode.ACenter,
            # put the text on the left side of the button
            text_pos = (0, -0.15),
            # set the text color to white
            text_fg = (1,1,1,1),
            # set the font of the text
            text_font = self.defaultFont,
            # set the buttons images
            geom = (self.defaultBtnMaps.find("**/button_ready"),
                    self.defaultBtnMaps.find("**/button_click"),
                    self.defaultBtnMaps.find("**/button_rollover"),
                    self.defaultBtnMaps.find("**/button_disabled")),
            # set no relief
            relief = 1,
            frameColor = (0,0,0,0),
            # No sink in when press
            pressEffect = False,
            # position on the window
            pos = (0, 0, .25),
            # the event which is thrown on clickSound
            command = self.btnReturn_Click,
            # sounds that should be played
            rolloverSound = None,
            clickSound = None)
        self.btnReturn.setTransparency(1)

        self.btnQuit = DirectButton(
            scale = (0.25, 0.25, 0.25),
            # some temp text
            text = "ABC",
            text_scale = (0.5*screenResMultiplier, 0.5, 0.5),
            # set the alignment to center
            text_align = TextNode.ACenter,
            # put the text on the left side of the button
            text_pos = (0, -0.15),
            # set the text color to white
            text_fg = (1,1,1,1),
            # set the font of the text
            text_font = self.defaultFont,
            # set the buttons images
            geom = (self.defaultBtnMaps.find("**/button_ready"),
                    self.defaultBtnMaps.find("**/button_click"),
                    self.defaultBtnMaps.find("**/button_rollover"),
                    self.defaultBtnMaps.find("**/button_disabled")),
            relief = 1,
            frameColor = (0,0,0,0),
            pressEffect = False,
            pos = (0, 0, -.25),
            command = self.btnQuit_Click,
            rolloverSound = None,
            clickSound = None)
        self.btnQuit.setTransparency(1)

        self.btnReturn.reparentTo(self.frameMain)
        self.btnQuit.reparentTo(self.frameMain)

        self.recalcAspectRatio()

        # set the text of all GUI elements
        self.setText()

        self.hideBase()

    def show(self):
        self.title["text"] = _("Pause")
        self.showBase()

    def hide(self):
        self.hideBase()

    def setText(self):
        self.title["text"] = _("Main Menu")
        self.btnReturn["text"] = _("Return")
        self.btnQuit["text"] = _("Quit")

    def btnReturn_Click(self):
        base.messenger.send("PauseMenu_return")

    def btnQuit_Click(self):
        base.messenger.send("PauseMenu_quit")


import os
import json
import statistics
from imutils import face_utils
from copy import deepcopy

import numpy as np
import dlib
import cv2

from config import LANDS_PATH, RATIOS_PATH,\
    RATIOS_AVG_PATH, RATIOS_STDS_PATH, ANGLES_PATH, ANGLES_AVG_PATH
from helpers import throw, is_number
from utils import cumul, match_hists

NUM_OF_RATIOS = 100232
NUM_OF_ANGLES = 50116

def faces(img):
    """ Returns dlib.rectangles containing rectangles of faces """
    if not isinstance(img, np.ndarray):
        throw("Expected numpy.array got " + type(img).__name__)
    if not os.path.isfile(LANDS_PATH):
        throw(LANDS_PATH + " doesn't exist")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_detector = dlib.get_frontal_face_detector()
    faces = face_detector(gray, 1)

    return faces

def landmarks(img, face):
    if not isinstance(face, dlib.rectangle):
        throw("Expected dlib.rectangle got " + type(face).__name__)
    predictor = dlib.shape_predictor(LANDS_PATH)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    dlib_shape = predictor(gray, face)
    return face_utils.shape_to_np(dlib_shape)

def all_distances(lands, as_dict=True, allow_repeats=False):
    """ Returns distances from each point of landmarks to every other point.
    Returning format:
    if as_dict is True
    {1: {2: 3}} - from 1st landmark to 2nd the distance is 3
    if as_dict is False
    [(1, 2, 3)]
    """
    if as_dict:
        distances = {}
    else:
        distances = []
    length = len(lands)
    for i in range(length):
        pfrom = lands[i]
        if as_dict:
            distances[i] = {}
        from_index = 0 if allow_repeats else i + 1
        for j in range(from_index, length):
            pto = lands[j]
            dis = _distance(pfrom, pto)
            if as_dict:
                distances[i][j] = dis
            else:
                distances.append((i, j, dis))
    return distances

def _distance(p1, p2):
    """ Returns distance between p1 point and p2 """
    return ((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5

def extract_rect(img, face):
    npf = _dlib_rect_to_np(face)
    return img[npf[1]:npf[1] + npf[3], npf[0]:npf[0] + npf[2]]

def stdeviation(data):
    """ Returns standard deviation of the data """
    return statistics.stdev(data)

def feed_ratio(val):
    """ Writes into ratio the new value - val
    The file is a line by line json objects
    """
    with open(RATIOS_PATH, 'a') as ratios:
        json.dump(val, ratios)
        ratios.write("\n")

def feed_angle(ang):
    """ Writes into angles the new angle - ang
    The file is a line by line json objects
    """
    with open(ANGLES_PATH, 'a') as angles:
        json.dump(ang, angles)
        angles.write("\n")

def ratios(all_dists):
    if not isinstance(all_dists, dict):
        throw("Expected dict got " + type(all_dists).__name__)

    length = len(all_dists)
    ratios = []
    for i in range(length):
        for j in range(i + 1, length):
            dist_i_j = all_dists[i][j]
            if i == j:
                continue
            for k in range(j + 1, length):
                if i == k or j == k:
                    continue
                dist_j_k = all_dists[j][k]
                dist_i_k = all_dists[i][k]

                by_j = dist_i_j / dist_j_k if dist_j_k != 0 else 0
                by_k = dist_j_k / dist_i_k if dist_i_k != 0 else 0

                ratios.append((i, k, j, by_k))
                ratios.append((i, j, k, by_j))
    return ratios

def calc_avg_angles():
    # TODO: To Be Continued...
    sums = np.zeros(NUM_OF_ANGLES)
    with open(ANGLES_PATH, 'r') as angles:
        for line_count, line in enumerate(angles):
            curr_angles = json.loads(line)
            for n, angle in enumerate(curr_angles):
                if is_number(angle[3]):
                    sums[n] += angle[3]
    avgs = []
    for i, j in enumerate(sums):
        avgs.append(j / (line_count + 1))
    with open(ANGLES_AVG_PATH, 'w') as avgs_file:
        json.dump(avgs, avgs_file)

def calc_avg_ratios():
    sums = np.zeros(NUM_OF_RATIOS)
    with open(RATIOS_PATH, 'r') as ratios:
        for line_count, line in enumerate(ratios):
            curr_ratios = json.loads(line)
            for n, ratio in enumerate(curr_ratios):
                sums[n] += ratio[3]

    avgs = [0 for i in range(NUM_OF_RATIOS)]
    for i, j in enumerate(sums):
        avgs[i] = j / (line_count + 1)
    with open(RATIOS_AVG_PATH, 'w') as avgs_file:
        json.dump(avgs, avgs_file)

def calc_ratio_stds():
    """ Calculate standard deviations of ratios and write them into file.
    Will be total (global NUM_OF_RATIOS)
    """
    data = {i: [] for i in range(NUM_OF_RATIOS)}
    with open(RATIOS_PATH, 'r') as ratios:
        for line_count, line in enumerate(ratios):
            curr_data = [x[3] for x in json.loads(line)]
            for n, ratio in enumerate(curr_data):
                data[n].append(ratio)
            print("Num of lines processed: " + str(line_count))
    print("Calculating standard deviations...")
    stds = [stdeviation(data[i]) for i in range(NUM_OF_RATIOS)]
    with open(RATIOS_STDS_PATH, 'w') as stds_file:
        json.dump(stds, stds_file)

def _dlib_rect_to_np(rect):
    """ Format: (x, y, width, height) - tuple """
    return face_utils.rect_to_bb(rect)

def show_img(img, name='default'):
    cv2.imshow(name, img)
    cv2.wait()

def compare_with_avg(img):
    fs = faces(img)
    lands = landmarks(img, fs[0])
    all_dists = all_distances(lands, allow_repeats=True)
    rs = [x[3] for x in ratios(all_dists)]

    with open(RATIOS_AVG_PATH, 'r') as avgs_file:
        means = json.load(avgs_file)
    with open(RATIOS_STDS_PATH, 'r') as stds_file:
        stds = json.load(stds_file)
    z_scores = [z_score(rs[i], means[i], stds[i])
                for i in range(NUM_OF_RATIOS)]
    mean_score = sum(z_scores) / len(z_scores)
    return mean_score

def z_score(x, u, o):
    return abs((x - u) / o)

def lands_with_names(lands):
    return [{
        "chin": lands[0:17],
        "left_eyebrow": lands[17:22],
        "right_eyebrow": lands[22:27],
        "nose_bridge": lands[27:31],
        "nose_tip": lands[31:36],
        "left_eye": lands[36:42],
        "right_eye": lands[42:48],
        "top_lip": lands[48:55] + [lands[64]] + [lands[63]] +
        [lands[62]] + [lands[61]] + [lands[60]],
        "bottom_lip": lands[54:60] + [lands[48]] + [lands[60]] +
        [lands[67]] + [lands[66]] + [lands[65]] + [lands[64]]
    }]

def lands_from_img(img):
    dlib_face = faces(img)[0]
    return landmarks(img, dlib_face)

def angles(lands):
    """ Returns dict of angles made by joining
    all combinations of 3 points from landmarks
    """
    length = len(lands)  # 68
    angles = {}
    for i in range(length):
        for j in range(i + 1, length):
            for k in range(j + 1, length):
                ang = angle(lands[i], lands[j], lands[k])
                angles[(i, j, k)] = ang
    return angles

def angle(a, b, c):
    """ Returns angle in degrees between
    line 1 and line 2 created by joining p1 to p2 and p2 to p3
    """
    ba = a - b
    bc = c - b
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    angle = np.arccos(cosine_angle)
    return np.degrees(angle)

def to_hsv(img):
    return cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

def to_rgb(img):
    return cv2.cvtColor(img, cv2.COLOR_HSV2BGR)

def equalize_red(img1, img2):
    height1, width1, _ = img1.shape
    height2, width2, _ = img2.shape
    if height1 != height2 or width1 != width2:
        throw('img1 and img2 should have the same size')

    red_hist1 = [0 for _ in range(255)]
    red_hist2 = [0 for _ in range(255)]

    # compute red color histograms
    for i in range(height1):
        for j in range(width1):
            b1, g1, r1 = img1[i, j]
            b2, g2, r2 = img2[i, j]
            red_hist1[r1] += 1
            red_hist2[r2] += 1
            # hsv1 = to_hsv(img1)
            # h, s, v = hsv[i, j]
            # hsv[i, j] = np.array([h, s, v], dtype=np.uint8)
    cumul1 = cumul(red_hist1)
    cumul2 = cumul(red_hist2)

    adjusted = match_hists(cumul1, cumul2)

    return repaint_red(img1, adjusted)

def repaint_red(img, mapping):
    h, w, c = img.shape
    for i in range(h):
        for j in range(w):
            b, g, r = img[i, j]
            r = mapping[r]
            img[i, j] = b, g, r
    return img

def bot_0(x):
    return 0.9848 * x - 6.7474

def top_0(x):
    return - 0.0009 * x * x + 1.1917 * x - 4.0146

def bot_1(x):
    return - 0.0009 * x * x + 1.1917 * x - 4.0146

def top_1(x):
    return - 0.0011 * x * x + 1.2262 * x + 4.0264

def bot_2(x):
    return - 0.0011 * x * x + 1.2262 * x + 4.0264

def top_2(x):
    return - 0.0013 * x * x + 1.2608 * x + 12.067

def bot_3(x):
    return - 0.0013 * x * x + 1.2608 * x + 12.067

def top_3(x):
    return - 0.0026 * x * x + 1.5713 * x + 14.8

def layer(img, layer_no):
    tmp = deepcopy(img)
    top_func = eval('top_' + str(layer_no))
    bot_func = eval('bot_' + str(layer_no))
    h, w, c = tmp.shape
    for i in range(h):
        for j in range(w):
            px = tmp[i, j]
            b, g, r = px
            rb = (int(r) + int(b)) / 2

            if b < g and g < r and rb >= bot_func(g) and rb <= top_func(g):
                tmp[i, j] = 0, 0, 0
            else:
                tmp[i, j] = 255, 255, 255
    return tmp


import os
import time
import yaml
import dotenv
import argparse
from flask import Flask, request, jsonify
from flask_restx import Api, Resource, fields, reqparse
from werkzeug.datastructures import FileStorage
from utils import decodebase64, load_model

app = Flask(__name__)
api = Api(app, version='1.0-alpha', title='Emotion Recognizer API',
          description="Recognize Emotion of a given image.", default="Endpoints", default_label='Emotion Recognizer')
namespace = api.namespace('emotion', description='Emotion Recognizer')

upload_parser = reqparse.RequestParser()
upload_parser.add_argument('imgfile', type=FileStorage,
                           location='files', required=True)

secrets = dotenv.dotenv_values('.env')
is_production = secrets["SERVER_ENV"] == "PRODUCTION"
print("Running on Production: {}".format(is_production))

'''
Gunicorn doesnt run on __main__, so we need to initialize the inference here
'''
if is_production:
    with open('default.config.yaml', 'r') as stream:
        config = yaml.safe_load(stream)
    labels = config['labels']
    os.makedirs('temp', exist_ok=True)

    emotion_recognizer = load_model(
        namespace.logger, config['model_xml'], config['weights_bin'])

@api.route('/emotion/base64', methods=["POST"], endpoint="base64")
@api.doc(responses={
    200: 'Successful Emotion Detection',
    400: 'Bad request',
    500: 'Internal server error.'},
    description="Detect Emotion of an Image in Base64 format")
class EmotionDetectionBase64(Resource):
    base64_request = api.model("Base64 Emotion Detection Endpoint", {
        "img": fields.String(required=True, description="Base64 encoded image")
    })

    @api.expect(base64_request)
    def post(self):
        image = request.json["img"]
        try:
            # decoding the base64 image and saving
            image = decodebase64(image)
            img_path = "temp/temp_{}.png".format(time.time())
            image.save(img_path)
            # predicting the emotion
            output = emotion_recognizer.predict(img_path)
            os.remove(img_path)
        except Exception as e:
            namespace.logger.error(e)
            return jsonify({"error": str(e)}, 400)
        return jsonify({"emotion": labels[output]})

@api.route('/emotion', methods=["POST"])
@api.doc(responses={
    200: 'Successful Emotion Detection',
    400: 'Bad request',
    500: 'Internal server error.'},
    description="Detect Emotion of an Uploaded Image")
class EmotionDetection(Resource):

    @api.expect(upload_parser)
    def post(self):
        uploads = upload_parser.parse_args()
        img_path = "temp/temp_{}.png".format(time.time())
        try:
            uploads["imgfile"].save(img_path)
            output = emotion_recognizer.predict(img_path)
            os.remove(img_path)
        except Exception as e:
            namespace.logger.error(e)
            return jsonify({"error": str(e)}, 400)
        return jsonify({"emotion": labels[output]})

if (__name__ == "__main__") & (not is_production):
    '''Following Lines run only in development'''
    argparser = argparse.ArgumentParser()
    argparser.add_argument('-p', '--port', default=5000,
                           type=int, help='port to listen on')
    argparser.add_argument(
        '--config', default="default.config.yaml", type=str, help='path to config file')
    argparser.add_argument('-d', '--debug', default=True,
                           action='store_true', help='enable debug mode')
    args = argparser.parse_args()

    with open(args.config, 'r') as stream:
        config = yaml.safe_load(stream)
    labels = config['labels']  # Label mappings
    os.makedirs('temp', exist_ok=True)

    emotion_recognizer = load_model(
        namespace.logger, config['model_xml'], config['weights_bin'])

    app.run(port=args.port, debug=args.debug)


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 28 10:35:20 2019

@author: rdamseh
"""

from GraphContraction import GraphContraction
import numpy as np
import networkx as nx
from util import isSklNodes, fixG
import itk

class ScanLines:

    def __init__(self, raw, scales=1, tolerance=.1):

        # read array
        self.__Im=itk.GetImageFromArray(raw.astype('float32'))
        self.__Im.UpdateOutputData()

        self.tolerance=tolerance
        self.scales=scales

        self.__VessMap=[]
        self.__BinaryMap=None
        self.__Skeleton=None
        self.__Graph=None
        self.__Lines=None

        #####
        # private
        #####

    def __GetVessImage(self, im, sigma=1):

        # apply hessian
        dimension = im.GetImageDimension()
        InputImageType = type(im)
        OutputPixelType = itk.UC
        OutputImageType = itk.Image[OutputPixelType, dimension]
        cast = itk.CastImageFilter[InputImageType, OutputImageType].New(im)
        hessianFilter = itk.HessianRecursiveGaussianImageFilter.New(cast.GetOutput())
        hessianFilter.SetSigma(sigma)
        hessianFilter.Update()

        # apply vesselness filter
        hessOutput = hessianFilter.GetOutput()
        dimension = hessOutput.GetImageDimension()
        InputImageType = type(hessOutput)
        OutputPixelType = itk.F
        OutputImageType = itk.Image[OutputPixelType, dimension]
        vess=itk.HessianToObjectnessMeasureImageFilter[InputImageType,
                                                         OutputImageType].New(hessOutput)
        vess.Update()
        vessOutput = vess.GetOutput()
        self.__VessMap.append(vessOutput)
        return vess

    def __GetBinaryImage(self, im, sigma=1):

        vess=self.__GetVessImage(im, sigma=sigma)
        vessOutput = vess.GetOutput()
        dimension = vessOutput.GetImageDimension()
        InputImageType = type(vessOutput)
        OutputPixelType = itk.F
        OutputImageType = itk.Image[OutputPixelType, dimension]
        vessOutput = itk.CastImageFilter[InputImageType, OutputImageType].New(vessOutput)

        # norm and thresholding from vessOutput
        norm=itk.NormalizeImageFilter.New(vessOutput)
        norm.Update()
        binary=itk.BinaryThresholdImageFilter.New(norm)
        binary.SetLowerThreshold(0)
        binary.Update()

        return binary

    def __RemoveSmallObjects(self, binary, thr_size=75):

        # norm and thresholding from vessOutput
        binary=itk.BinaryThresholdImageFilter.New(binary.GetOutput())
        binary.SetLowerThreshold(1)
        binary.Update()

        # get labels and remove small ones
        binaryOutput=binary.GetOutput()
        labels=itk.BinaryImageToLabelMapFilter.New(binaryOutput)
        labels.Update()
        labelsOutput=labels.GetOutput()
        nLabels=labelsOutput.GetNumberOfLabelObjects()

        size=[]
        for i in range(nLabels):
            size.append(int(labelsOutput.GetNthLabelObject(i).Size()))

        labelsToRemove=np.where(np.array(size)<thr_size)[0]

        for i in labelsToRemove:
            labelsOutput.RemoveLabel(int(i+1))

        seg=itk.LabelMapToBinaryImageFilter.New(labelsOutput)
        seg.Update()

        return seg

    def __GetBinaryMap(self):

        sigmas=[1,3,5,7,9,13]
        fillRadius=3

        # get binary images at different ranks
        binaryImages=[]
        for i in range(self.scales):
            binaryImages.append(self.__GetBinaryImage(self.__Im,
                                        sigma=sigmas[i]))

        # add binary maps
        if self.scales>1:

            for i in range(len(binaryImages)):
                if i>0:
                    if i==1:
                        add=itk.AddImageFilter.New(binaryImages[i-1], binaryImages[i])
                        add.Update()
                    else:
                        add=itk.AddImageFilter.New(add.GetOutput(), binaryImages[i])
                        add.Update()
        else:

            add=binaryImages[0]

        # remove small objects
        refined=self.__RemoveSmallObjects(add, thr_size=75)

        # fill holes
        refined=itk.VotingBinaryIterativeHoleFillingImageFilter.New(refined.GetOutput());
        refined.SetRadius(fillRadius)
        refined.Update()

        self.__BinaryMap=refined.GetOutput()

    def __GetSkeleton(self):

        if self.__BinaryMap:

            # get skeleton
            skl=itk.BinaryThinningImageFilter.New(self.__BinaryMap)
            skl.Update()

            self.Skeleton=skl.GetOutput()

    def __GetGraphFromSkl(self):
        return

    def __GetGraphFromSeg(self):
        print('Generating a graph ...')
        if self.__BinaryMap is not None:

            binmap=itk.GetArrayFromImage(self.__BinaryMap)
            binmap=np.swapaxes(binmap,0,1).astype('bool')
            label=np.array(binmap[None,:])
            label=np.rollaxis(label,0,3)

            G=GraphContraction(label=label)
            G.generateGraph()
            G.contractGraph(speed_param=.1, #[low_speed, high_speed]
                             dis_param=.1,
                             deg_param=0,
                             med_param=1,
                             thr=5,
                             n_iter=10,
                             stop_thr=.01)

            G.refineGraph(diam_refine=5)
            graph=G.G_refined

            self.__Graph=graph

    def __GetLines(self):

        diam=self.diam
        length=self.length
        if self.__Graph:

            graph=self.__Graph.copy()

            # get potential lines
            tolerance=self.tolerance
            smooth_itr=1
            seg_len=4

            nb=[graph.neighbors(i) for i in graph.nodes()]
            nodesToRemove1=[k for k,i in enumerate(nb) if len(i)>2]
            nodesToProcess=[k for k,i in enumerate(nb) if len(i)==2]

            # average position
            for i in range(smooth_itr):
                ownpos=np.array([graph.node[i]['pos'] for i in nodesToProcess])
                otherpos=np.array([[ graph.node[k]['pos']  for k in graph.neighbors(i)] for i in nodesToProcess])
                posNew=(ownpos.astype('float')+otherpos[:,0,:]+otherpos[:,1,:])/3
                for ind, node in enumerate(nodesToProcess):
                    graph.node[node]['pos']=posNew[ind]

            otherpos=np.array([[ graph.node[k]['pos']  for k in graph.neighbors(i)] for i in nodesToProcess])
            indNodesToRemove2=isSklNodes(posNew, otherpos, thr=tolerance)==False
            nodesToProcess=np.array(nodesToProcess)
            nodesToRemove2=nodesToProcess[indNodesToRemove2]

            graph.remove_nodes_from(nodesToRemove1)
            graph.remove_nodes_from(nodesToRemove2)
            graph=fixG(graph)

            # get graph segments
            components=list(nx.connected_components(graph))
            components=[i for i in components if len(i)> seg_len]

            segments=[[i for i in j if len(graph.neighbors(i))==1] for j in components]
            segments=[i for i in segments if len(i)==2] # ensure line-like segments

            print('Number of segments after tolerance constraint: '+ str(len(segments)))
            if len(segments)==0:
                return

            segments=np.array(segments)

            # remove segments with big diam
            segmentsToRecall=[]
            for idx, i in enumerate(segments):
                try:
                    if graph.node[i[0]]['d']<diam and graph.node[i[1]]['d']<diam:
                        segmentsToRecall.append(idx)
                except:
                    print('Error in graph segment :'+str(i))
                    pass

            segments=segments[segmentsToRecall,:]

            print('Number of segments after radius constraint: '+ str(len(segments)))
            if len(segments)==0:
                return

            lines=[[np.array(graph.node[j]['pos']) for j in i] for i in segments]
            lines=np.array(lines)

            # retrieve with length constrains
            norm=np.linalg.norm(lines[:,0,:]-lines[:,1,:], axis=1)
            mid=np.mean(lines, axis=1)
            unitV=(lines[:,0,:]-lines[:,1,:])/norm[:,None]

            IndLinesToReduce=np.where(norm>length)
            LinesReduced=np.array([mid[IndLinesToReduce]-length*unitV[IndLinesToReduce]/2.0,
                                    mid[IndLinesToReduce]+length*unitV[IndLinesToReduce]/2.0])
            LinesReduced=np.swapaxes(LinesReduced,0,1)
            lines[IndLinesToReduce]=LinesReduced

            print('Number of segments after length constraint: ', len(lines))
            if len(segments)==0:
                return

            self.__Lines=lines

    #######
    # Public
    #######

    def CalcBinaryMap(self, scales=1):
        self.scales=scales
        self.__GetBinaryMap()

    def CalcSkeleton(self):
        self.__GetSkeleton()

    def CalcGraphFromSeg(self):
        self.__GetGraphFromSeg()

    def CalcGraphFromSkl(self):
        self.__GetGraphFromSkl()

    def CalcLines(self, diam=5, length=15, tolerance=.1):
        self.diam=diam
        self.length=length
        self.tolerance=tolerance
        self.__GetLines()

    def GetOutputBinaryMap(self):

        if self.__BinaryMap is not None:

            im=itk.GetArrayFromImage(self.__BinaryMap)
            im=np.swapaxes(im,0,1).astype('bool')

            return im

    def GetOutputVessMap(self):

        if len(self.__VessMap)>0:

            vessmap=[]

            for i in self.__VessMap:

                vess=itk.GetArrayFromImage(i)
                vessmap.append(np.swapaxes(vess,0,1).astype('float'))

            return vessmap

    def GetOutputSkeleton(self):

        if self.Skeleton is not None:

            skl=itk.GetArrayFromImage(self.Skeleton)
            skl=np.swapaxes(skl,0,1).astype('bool')

            return skl

    def GetOutputGraph(self):

        if self.__Graph is not None: return self.__Graph

    def GetOutputLines(self):

        if self.__Lines is not None: return self.__Lines

if __name__=='__main__':

    pass


"""
Annotator

Annotator for .nlp files. Facilitates the addition of comments
to nlp pass files to prepare for code generation clm training.
"""

import pandas as pandas
import math
import re
import os
import typing
from typing import Union, List, Tuple, Dict
from tqdm import tqdm
import logging
import warnings
import pickle
from datasets import Dataset

logger = logging.getLogger(__name__)

class Annotator:
    def __init__(self, dataset: Union[List[str], dict]):
        """
        Initialize annotator with nlp++ dataset.
        """
        self.dataset = dataset
        self.current_index = 0
        self.finished = False
        warnings.warn('utils.annotator.Annotator is deprecated.', DeprecationWarning, stacklevel=2)

class Dataset:
    def __init__(self, dataset: Union[List[str], dict], text_column_name: str="text"):
        if isinstance(dataset, list):
            self._dataset = {text_column_name: dataset}
        else:
            self._dataset = dataset

        self.text_column_name = text_column_name
        self.column_names = [text_column_name]

        warnings.warn('utils.annotator.Dataset is deprecated.', DeprecationWarning, stacklevel=2)

    def __getitem__(self, key: Union[slice, int]):
        if isinstance(key, slice):
            result = {}
            for k, v in self._dataset:
                result[k] = v[key]
        elif isinstance(key, int):
            result = {}
            for k, v in self._dataset:
                result[k] = v[key]
        else:
            raise ValueError

        return result

    def __setitem__(self, key: int, value: Union[str, dict]):
        if isinstance(value, str):
            if len(self._dataset) > 1:
                logger.warn(f"Columns in value (1) must match dataset ({len(self._dataset)}).")
                raise ValueError
            self._dataset[self.text_column_name][key]

        elif isinstance(value, dict):
            if set(value.keys()) != set(self._dataset):
                logger.warn(f"Columns in value (1) must match dataset ({len(self._dataset)}).")
                raise ValueError

        else:
            raise ValueError("Value must be type str or dict.")

    def __len__(self):
        return len(self._dataset[self.text_column_name])

    def _set_column_names(self):
        """
        Update column_names.
        """
        self.column_names = list(self._dataset.keys())

    def rename_columns(self, column_mapping: Dict[str,str]):
        """
        Rename column(s) using mapping.
        """
        for k, v in column_mapping.items():
            if k not in self._dataset:
                raise KeyError(f"Column '{k}' not in dataset.")
            if v in self.dataset:
                raise ValueError(f"Column '{v}' already in dataset.")
            self._dataset[v] = self._dataset[k]
            del self._dataset[k]

        self._set_column_names()

    def to_hf_dataset(self):
        """
        Convert dataset to datasets.Dataset
        """
        return Dataset.from_dict(self._dataset)

    def save_dataset(self, save_path: str):
        """
        pickle dataset
        """
        with open(save_path, "wb") as f:
            pickle.dump(self.__dict__, f)

    @classmethod
    def load_dataset(cls, dataset_path: str) -> Dataset:
        """
        Load saved dataset from file.
        """
        with open(save_path, "rb") as f:
            dataset = pickle.load(f)
            return cls(dataset["_dataset"], dataset["text_column_name"])