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

import numpy as np
import pandas as pd

from ... import opcodes as OperandDef
from ...serialize import KeyField
from ...tiles import TilesError
from ...tensor import tensor as astensor
from ...tensor.core import TENSOR_TYPE
from ...tensor.utils import decide_unify_split, validate_axis, recursive_tile
from ..core import DATAFRAME_TYPE, SERIES_TYPE, IndexValue
from ..operands import DataFrameOperand, DataFrameOperandMixin, ObjectType
from ..utils import parse_index

class DataFrameDot(DataFrameOperand, DataFrameOperandMixin):
    _op_type_ = OperandDef.DOT

    _lhs = KeyField('lhs')
    _rhs = KeyField('rhs')

    def __init__(self, object_type=None, lhs=None, rhs=None, **kw):
        super().__init__(_object_type=object_type, _lhs=lhs, _rhs=rhs, **kw)

    @property
    def lhs(self):
        return self._lhs

    @property
    def rhs(self):
        return self._rhs

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        self._lhs = self._inputs[0]
        self._rhs = self._inputs[1]

    def __call__(self, lhs, rhs):
        if not isinstance(rhs, (DATAFRAME_TYPE, SERIES_TYPE)):
            rhs = astensor(rhs)
            test_rhs = rhs
        else:
            test_rhs = rhs.to_tensor()

        test_ret = lhs.to_tensor().dot(test_rhs)
        if test_ret.ndim == 0:
            if isinstance(lhs, SERIES_TYPE) and isinstance(rhs, TENSOR_TYPE):
                # return tensor
                return test_ret
            self._object_type = ObjectType.scalar
            return self.new_scalar([lhs, rhs], dtype=test_ret.dtype)
        elif test_ret.ndim == 1:
            self._object_type = ObjectType.series
            if lhs.ndim == 1:
                if hasattr(rhs, 'columns_value'):
                    index_value = rhs.columns_value
                else:
                    # tensor
                    length = -1 if np.isnan(rhs.shape[1]) else rhs.shape[1]
                    pd_index = pd.RangeIndex(length)
                    index_value = parse_index(pd_index, store_data=True)
            else:
                assert rhs.ndim == 1
                index_value = lhs.index_value
            return self.new_series([lhs, rhs], shape=test_ret.shape,
                                   dtype=test_ret.dtype, index_value=index_value)
        else:
            self._object_type = ObjectType.dataframe
            if isinstance(rhs, TENSOR_TYPE):
                dtypes = pd.Series(np.repeat(test_ret.dtype, test_ret.shape[1]),
                                   index=pd.RangeIndex(test_ret.shape[1]))
                columns_value = parse_index(dtypes.index, store_data=True)
            else:
                dtypes = pd.Series(np.repeat(test_ret.dtype, test_ret.shape[1]),
                                   index=rhs.columns_value.to_pandas())
                columns_value = rhs.columns_value
            return self.new_dataframe([lhs, rhs], shape=test_ret.shape,
                                      index_value=lhs.index_value,
                                      columns_value=columns_value,
                                      dtypes=dtypes)

    @classmethod
    def _align(cls, lhs, rhs):
        if isinstance(rhs, TENSOR_TYPE):
            # no need to align when rhs is a tensor
            return lhs, rhs

        is_lhs_range_index = False
        if isinstance(lhs, DATAFRAME_TYPE) and \
                isinstance(lhs.columns_value.value, IndexValue.RangeIndex):
            is_lhs_range_index = True
        if isinstance(lhs, SERIES_TYPE) and \
                isinstance(lhs.index_value.value, IndexValue.RangeIndex):
            is_lhs_range_index = True

        is_rhs_range_index = False
        if isinstance(rhs.index_value.value, IndexValue.RangeIndex):
            is_rhs_range_index = True

        if not is_lhs_range_index or not is_rhs_range_index:
            # TODO: e.g. use rhs.loc[lhs.columns_value.to_pandas()]
            # when lhs is a DataFrame and lhs.columns is not a RangeIndex,
            # so does Series
            raise NotImplementedError

        return lhs, rhs

    @classmethod
    def tile(cls, op):
        from ..datasource.from_tensor import dataframe_from_tensor, series_from_tensor

        lhs, rhs = op.lhs, op.rhs
        lhs, rhs = cls._align(lhs, rhs)
        out = op.outputs[0]

        if any(np.isnan(ns) for ns in lhs.nsplits[-1]):
            raise TilesError('lhs has unknown chunk shape on last axis')
        if any(np.isnan(ns) for ns in rhs.nsplits[0]):
            raise TilesError('rhs has unknown chunk shape on first axis')

        nsplit = decide_unify_split(lhs.nsplits[-1], rhs.nsplits[0])
        lhs_axis = validate_axis(lhs.ndim, -1)
        lhs = lhs.rechunk({lhs_axis: nsplit})._inplace_tile()
        rhs = rhs.rechunk({0: nsplit})._inplace_tile()

        # delegate computation to tensor
        lhs_tensor = lhs if isinstance(lhs, TENSOR_TYPE) else lhs.to_tensor()
        rhs_tensor = rhs if isinstance(rhs, TENSOR_TYPE) else rhs.to_tensor()
        ret = lhs_tensor.dot(rhs_tensor)

        if isinstance(out, TENSOR_TYPE):
            pass
        elif ret.ndim == 1:
            ret = series_from_tensor(ret)
            if isinstance(lhs, DATAFRAME_TYPE):
                # lhs DataFrame
                ret._index_value = lhs.index_value
            elif isinstance(rhs, DATAFRAME_TYPE):
                # lhs Series, rhs DataFrame
                ret._index_value = rhs.columns_value
        elif ret.ndim == 2:
            ret = dataframe_from_tensor(ret)
            ret._index_value = lhs.index_value
            if isinstance(rhs, DATAFRAME_TYPE):
                ret._columns_value = rhs.columns_value
                ret._dtypes = rhs.dtypes

        return [recursive_tile(ret)]

def dot(df_or_seris, other):
    op = DataFrameDot(lhs=df_or_seris, rhs=other)
    return op(df_or_seris, other)


from __future__ import print_function, division

from keras.applications.vgg16 import VGG16
from keras.applications.inception_v3 import InceptionV3
from keras.applications.xception import Xception
from keras.applications.resnet50 import ResNet50

from keras.applications import vgg16, inception_v3, xception, resnet50
import keras.applications.imagenet_utils as imagenet_utils

from keras.models import Model
from keras.layers import Dense, Flatten
from keras.optimizers import SGD, Adam

import copy

import numpy as np
import abc
from .train import TrainClassifierBase
from .train_inception_base import TrainClassifierInceptionBase

class TrainClassifierResnet50(TrainClassifierBase):
    '''
    Fine-tune resnet models based on the keras_resnet package
    '''

    def __init__(self, root_dir, batch_size=32, weights_file = None, model_file = None, epochs = 5,
                    n_classes = 2, img_shape = (224, 224), class_mode = 'binary'):

        super(TrainClassifierResnet50, self).__init__(root_dir, batch_size, weights_file = weights_file, model_file = model_file,
                            epochs = epochs, n_classes = n_classes, img_shape = img_shape, class_mode = class_mode)

    def configure_model(self, multi_gpu = False):
        '''
        Download & configure the appropriate resnet
        '''
        self.model = ResNet50(weights = 'imagenet', include_top=False, input_shape=self.img_shape + (3,))

        # freeze all layers
        for layer in self.model.layers:
            layer.trainable = True

        units = len(self.n_classes) if self.class_mode == 'categorical' else 1

        x = Flatten()(self.model.layers[-1].output)
        predictor = Dense(units=units, activation='sigmoid')(x)
        self.model = Model(self.model.input, predictor)

        self.optimizer = Adam(lr=1e-5) if self.optimizer is None else self.optimizer

        self.model = self.get_parallel_model(self.model) if multi_gpu else self.model

    def pretrain(self):
        self.model.compile(optimizer=self.optimizer, loss='{}_crossentropy'.format(self.class_mode), metrics=['accuracy'])

    @property
    def network_name(self):
        return "resnet50"

    @property
    def rescale(self):
        return 1.

    def get_network_module(self):
        return resnet50

    def pretrain(self):
        pass

class TrainClassifierVgg16(TrainClassifierBase):
    """Fine-tule the VGG16 net"""

    def __init__(self, root_dir, batch_size=32, weights_file = None, model_file = None, epochs = 5,
                    n_classes = 2, img_shape = (224, 224), class_mode = 'binary'):

        super(TrainClassifierVgg16, self).__init__(root_dir, batch_size, weights_file = weights_file, model_file = model_file,
                                                        img_shape = img_shape, epochs = epochs, n_classes = n_classes, class_mode = class_mode)

    def configure_model(self, multi_gpu=False):
        '''
        Configure model for training (VGG16)
        '''
        self.vgg16 = VGG16(weights='imagenet')
        fc2 = self.vgg16.get_layer('fc2').output
        prediction = Dense(units=1, activation='sigmoid', name='logit')(fc2)
        self.model = Model(inputs=self.vgg16.input, outputs=prediction)

        for layer in self.model.layers:
            if layer.name in ['fc1', 'fc2', 'logit']:
                continue
            layer.trainable = False

        sgd = SGD(lr=1e-4, momentum=0.9)

        # if we have multiple gpus - train on both
        self.model = self.get_parallel_model(self.model) if multi_gpu else self.model

        self.model.compile(optimizer=sgd, loss='{}_crossentropy'.format(self.class_mode), metrics=['accuracy'])

    @property
    def network_name(self):
        return "vgg16"

    @property
    def rescale(self):
        return 1.

    def get_network_module(self):
        return vgg16

    def pretrain(self):
        pass

class TrainClassifierInceptionV3(TrainClassifierInceptionBase):
    """Fine-tule InceptionV3 net"""

    def __init__(self, root_dir, batch_size=32, epochs=3, weights_file = None, model_file = None,
                    n_classes = 2, img_shape = (299, 299), class_mode='binary'):

        super(TrainClassifierInceptionV3, self).__init__(root_dir, batch_size, epochs, n_classes = n_classes, weights_file = weights_file,
                                                                img_shape = img_shape, model_file = model_file, class_mode=class_mode)

    @property
    def network_name(self):
        return "inception_v3"

    @property
    def Network(self):
        return InceptionV3

    @property
    def train_layer_cutoff(self):
        return 172

    def get_network_module(self):
        return inception_v3

class TrainClassifierXception(TrainClassifierInceptionBase):
    """Fine-tule Xception net"""

    def __init__(self, root_dir, batch_size = 32, epochs = 3, weights_file = None, model_file = None, n_classes = 2, img_shape = (299, 299), class_mode='binary'):
        super(TrainClassifierXception, self).__init__(root_dir, batch_size, epochs, n_classes = n_classes, weights_file = weights_file,
                                                            img_shape = img_shape, model_file = model_file, class_mode=class_mode)

    @property
    def network_name(self):
        return "xception"

    def get_network_module(self):
        return xception

    @property
    def Network(self):
        return Xception

    @property
    def train_layer_cutoff(self):
        return 126


import cv2
import numpy as np
import copy
import math
import main as detect
import predict as predict
import os
import time
# from tkinter import *

#from appscript import app

# Environment:
# OS    : Mac OS EL Capitan
# python: 3.5
# opencv: 2.4.13

# parameters

threshold = 12  #  BINARY threshold
blurValue = 41  # GaussianBlur parameter
bgSubThreshold = 50
learningRate = 0
traverse_point = []
# variables
isBgCaptured = 0   # bool, whether the background captured
triggerSwitch = False  # if true, keyborad simulator works
bgModel = None
def printThreshold(thr):
    print("! Changed threshold to "+str(thr))

def actionPerform(val):
        if val=='f' or val=='F':
            print("opening Firefox")
            os.system("gnome-terminal -e 'firefox www.iiita.ac.in'")
        elif val=='c' or val=='C':
            print("opening Chrome")
            os.system("gnome-terminal -e 'google-chrome www.iiita.ac.in'")
        elif val=='v' or val=='V' or val=='y' or val=='Y':
            print("Opening VLC")
            os.system("gnome-terminal -e 'vlc /home/sudomaniac/Desktop/Afreen\ Afreen.mp4'")
        elif val=='1':
            print("Vol MAX")
            os.system('pactl set-sink-volume 0 90%')
        elif val=='2':
            print("Vol MID")
            os.system('pactl set-sink-volume 0 50%')
        elif val=='3':
            print("Vol MIN")
            os.system('pactl set-sink-volume 0 10%')
        else:
            return False
        return True

def removeBG(frame):
    fgmask = bgModel.apply(frame,learningRate=learningRate)
    # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    # res = cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, kernel)
    kernel = np.ones((3, 3), np.uint8)
    fgmask = cv2.erode(fgmask, kernel, iterations=1)
    res = cv2.bitwise_and(frame, frame, mask=fgmask)
    return res

def calculateFingers(res,drawing):  # -> finished bool, cnt: finger count
    #  convexity defect
    global traverse_point
    hull = cv2.convexHull(res, returnPoints=False)
    if len(hull) > 3:
        defects = cv2.convexityDefects(res, hull)
        if type(defects) != type(None):  # avoid crashing.   (BUG not found)

            cnt = 0
            top = []
            for i in range(defects.shape[0]):  # calculate the angle
                s, e, f, d = defects[i][0]
                start = tuple(res[s][0])
                end = tuple(res[e][0])
                far = tuple(res[f][0])
                a = math.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2)
                b = math.sqrt((far[0] - start[0]) ** 2 + (far[1] - start[1]) ** 2)
                c = math.sqrt((end[0] - far[0]) ** 2 + (end[1] - far[1]) ** 2)
                angle = math.acos((b ** 2 + c ** 2 - a ** 2) / (2 * b * c))  # cosine theorem
                if angle <= math.pi / 2:  # angle less than 90 degree, treat as fingers
                    cnt += 1
                    cv2.circle(drawing, far, 8, [211, 84, 0], -1)
                top.append(start)
            point = min(top,key=lambda item:item[1])
            point = (point[0]+300,point[1])
            if startDetection:
                if(cnt==0):
                    traverse_point.append(point)
                else:
                    traverse_point.append((-1,-1))
            return True, cnt
    return False, 0

startDetection = False
# Camera
def main():
    global startDetection,traverse_point,triggerSwitch,isBgCaptured,learningRate,bgSubThreshold,threshold,blurValue,bgModel
    camera = cv2.VideoCapture(0)

    camera.set(10,200)
    cv2.namedWindow('trackbar')
    cv2.moveWindow('trackbar', 300, 680)
    cv2.createTrackbar('trh1', 'trackbar', threshold, 100, printThreshold)

    finale = cv2.imread('template.jpg',0)

    s = None
    while camera.isOpened():
        ret, frame = camera.read()
        threshold = cv2.getTrackbarPos('trh1', 'trackbar')
        frame = cv2.bilateralFilter(frame, 5, 50, 100)  # smoothing filter
        frame = cv2.flip(frame, 1)  # flip the frame horizontally
        cv2.rectangle(frame, (300,0),
                    (600,300), (255, 0, 0), 2)
        # print(int(cap_region_x_begin * frame.shape[1])
        # print(frame.shape[1])
        #  Main operation
        if isBgCaptured == 1:  # this part wont run until background captured
            img = removeBG(frame)
            img = img[0:300,
                        300:600]  # clip the ROI
            cv2.imshow('mask', img)
            cv2.moveWindow('mask', 860, 100)

            # convert the image into binary image
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            blur = cv2.GaussianBlur(gray, (blurValue, blurValue), 0)
            cv2.imshow('blur', blur)
            cv2.moveWindow('blur', 1300, 100)

            ret, thresh = cv2.threshold(blur, threshold, 255, cv2.THRESH_BINARY)
            cv2.imshow('ori', thresh)
            cv2.moveWindow('ori', 1300, 550)

            # get the coutours
            thresh1 = copy.deepcopy(thresh)
            contours, hierarchy = cv2.findContours(thresh1, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            length = len(contours)

            maxArea = -1
            if length > 0:
                for i in range(length):  # find the biggest contour (according to area)
                    temp = contours[i]
                    area = cv2.contourArea(temp)
                    if area > maxArea:
                        maxArea = area
                        ci = i

                res = contours[ci]
                hull = cv2.convexHull(res)
                drawing = np.zeros(img.shape, np.uint8)
                cv2.drawContours(drawing, [res], 0, (0, 255, 0), 2)
                cv2.drawContours(drawing, [hull], 0, (0, 0, 255), 3)

            # //////// VOLUME CONTROL FUNCTIONALITY
                isFinishCal,cnt = calculateFingers(res,drawing)
                cnt=cnt+1

                if cnt == 1:
                    actionPerform('1')
                elif cnt==2:
                    actionPerform('2')
                elif cnt==3:
                    actionPerform('3')
                if triggerSwitch is True:
                    if isFinishCal is True and cnt <= 2:
                        print(cnt)

        # //////// VOLUME CONTROL FUNCTIONALITY

                        #app('System Events').keystroke(' ')  # simulate pressing blank space

            for i in range(0,len(traverse_point)):
                    cv2.circle(frame, traverse_point[i], 2, [0, 0, 255], -1)
                    # cv2.circle(finale,(int((traverse_point[i][0]-128)),int(traverse_point[i][1]/2)), 2, [0, 0, 0], -1)
                    # print(traverse_point[i])
            for i in range(1,len(traverse_point)):
                    if(traverse_point[i][0]!=-1 and traverse_point[i-1][0]!=-1):
                        cv2.line(frame,traverse_point[i],traverse_point[i-1],[0,255,0],15,lineType = cv2.LINE_AA)
                        cv2.line(finale,(int((traverse_point[i][0]-128)),int(traverse_point[i][1]/2)),(int((traverse_point[i-1][0]-128)),int(traverse_point[i-1][1]/2)),[0,0,0],15,lineType = cv2.LINE_AA)
                    else:
                        i = i+1
            cv2.imshow('output', drawing)
            cv2.moveWindow('output', 860, 550)

        if s != None:
            cv2.putText(frame,s, (0, 480), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (0, 0, 0), lineType=cv2.LINE_AA)

        # frame.set(cv2.CAP_PROP_FRAME_WIDTH, 400)
        # frame.set(cv2.CAP_PROP_FRAME_HEIGHT, 600)
        cv2.imshow('original', frame)
        cv2.moveWindow('original', 200, 100)
        # Keyboard OP
        k = cv2.waitKey(10)
        if k == 27:  # press ESC to exit
            break
        elif k == ord('b'):  # press 'b' to capture the background
            bgModel = cv2.createBackgroundSubtractorMOG2(0, bgSubThreshold)
            isBgCaptured = 1
            print( '!!!Background Captured!!!')
        elif k == ord('r'):  # press 'r' to reset the background
            bgModel = None
            triggerSwitch = False
            isBgCaptured = 0
            cv2.imwrite("ok.jpg",finale)
            finale = cv2.imread("template.jpg",0)

            if s == None:
                s = predict.main()
                cv2.putText(frame,s, (0,0), cv2.FONT_HERSHEY_SIMPLEX, 2, 255)
                print(s)
                x = actionPerform(s)
                if x:
                    time.sleep(3)
                    bgModel = cv2.createBackgroundSubtractorMOG2(0, bgSubThreshold)
                    isBgCaptured = 1
                    print('!!!Background Captured!!!')
                print("val : "+s)
                s=None

            traverse_point = []
            startDetection = False
            print ('!!!Reset BackGround!!!')
        elif k == ord('n'):
            triggerSwitch = True
            print ('!!!Trigger On!!!')
        elif k == ord('d'):
            startDetection = True
            print ('!!!!Detection On!!!!')

if __name__ == '__main__':
    	main()


import math
import pickle
import random

import numpy as np
from nltk.corpus import wordnet
from nltk.tokenize import word_tokenize

class Augmenter(object):
    """Data augmentation operator.

    Support both span and attribute level augmentation operators.
    """

    def __init__(self, idf_fn="data/wikitext-idf.dat"):
        # self.index = json.load(open(index_fn))
        if idf_fn is not None:
            self.idf_dict = pickle.load(open(idf_fn, "rb"))
        else:
            self.idf_dict = None

    def augment(self, tokens, labels, op="del", args=[]):
        """Performs data augmentation on a sequence of tokens

        The supported ops:
           ['del',
            'token_del_tfidf',
            'token_del',
            'shuffle',
            'token_repl',
            'token_repl_tfidf'
            'all']

        Args:
            tokens (list of strings): the input tokens
            labels (list of strings): the labels of the tokens
            op (str, optional): a string encoding of the operator to be applied

        Returns:
            list of strings: the augmented tokens
            list of strings: the augmented labels
        """
        tfidf = "tfidf" in op
        if "del" in op:
            if "token_del" in op:
                pos1 = self.sample_position(tokens, labels, tfidf)
                if pos1 < 0:
                    return tokens, labels
                new_tokens = tokens[:pos1] + tokens[pos1 + 1 :]
                new_labels = labels[:pos1] + labels[pos1 + 1 :]
                # new_tokens = tokens[:pos1] + ['[PAD]'] + tokens[pos1+1:]
                # new_labels = labels[:pos1] + ['<PAD>'] + labels[pos1+1:]
            else:
                # insert padding to keep the length consistent
                # span_len = random.randint(1, 3)
                max_len = args[0] if len(args) > 0 else 2
                span_len = random.randint(1, max_len)
                pos1, pos2 = self.sample_span(tokens, labels, span_len=span_len)
                if pos1 < 0:
                    return tokens, labels
                new_tokens = tokens[:pos1] + tokens[pos2 + 1 :]
                new_labels = tokens[:pos1] + labels[pos2 + 1 :]
        elif "insert" in op:
            pos1 = self.sample_position(tokens, labels, tfidf)
            if pos1 < 0:
                return tokens, labels
            ins_token = self.sample_token(tokens[pos1])
            if ins_token.lower() != tokens[pos1].lower():
                new_tokens = tokens[:pos1] + [ins_token] + tokens[pos1:]
                new_labels = labels[:pos1] + ["O"] + labels[pos1:]
            else:
                new_tokens, new_labels = tokens, labels
        elif "repl" in op:
            pos1 = self.sample_position(tokens, labels, tfidf)
            if pos1 < 0:
                return tokens, labels
            ins_token = self.sample_token(tokens[pos1])
            new_tokens = tokens[:pos1] + [ins_token] + tokens[pos1 + 1 :]
            new_labels = labels[:pos1] + ["O"] + labels[pos1 + 1 :]
        elif "shuffle" in op:
            max_len = args[0] if len(args) > 0 else 4
            span_len = random.randint(2, max_len)
            pos1, pos2 = self.sample_span(tokens, labels, span_len=span_len)
            if pos1 < 0:
                return tokens, labels
            sub_arr = tokens[pos1 : pos2 + 1]
            random.shuffle(sub_arr)
            new_tokens = tokens[:pos1] + sub_arr + tokens[pos2 + 1 :]
            new_labels = tokens[:pos1] + ["O"] * (pos2 - pos1 + 1) + labels[pos2 + 1 :]
        else:
            new_tokens, new_labels = tokens, labels

        return new_tokens, new_labels

    def augment_sent(self, text, op="all", args=[]):
        """Performs data augmentation on a classification example.

        Similar to augment(tokens, labels) but works for sentences
        or sentence-pairs.

        Args:
            text (str): the input sentence
            op (str, optional): a string encoding of the operator to be applied
            args (list, optional): the augmentation parameters (length, etc.)

        Returns:
            str: the augmented sentence
        """
        # 50% of chance of flipping
        # if ' [SEP] ' in text and random.randint(0, 1) == 0:
        #     left, right = text.split(' [SEP] ')
        #     text = right + ' [SEP] ' + left

        # tokenize the sentence
        if " [SEP] " in text:
            left, right = text.split(" [SEP] ")
            tokens = word_tokenize(left) + ["[SEP]"] + word_tokenize(right)
        else:
            tokens = word_tokenize(text)

        # avoid the special tokens
        labels = []
        for token in tokens:
            if token in ["[CLS]", "[SEP]"]:
                labels.append("<SEP>")
            else:
                labels.append("O")

        if op == "all":
            # RandAugment: https://arxiv.org/pdf/1909.13719.pdf
            N = 3
            ops = [
                "del",
                "token_del_tfidf",
                "token_del",
                "shuffle",
                "token_repl",
                "token_repl_tfidf",
            ]
            for op in random.choices(ops, k=N):
                tokens, labels = self.augment(tokens, labels, op=op)
        if op == "corrupt" or op == "corrupt_plus":
            # remove 20% of tokens
            N = len(tokens)
            to_change = N // 3
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="token_del_tfidf")
            N = len(tokens)
            to_change = N // 5
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="token_repl_tfidf")
            to_change = N // 5
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="shuffle")
        elif op == "corrupt_minus":
            # add 20% of tokens
            N = len(tokens)
            to_change = N // 3
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="insert")
            N = len(tokens)
            to_change = N // 5
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="token_repl_tfidf")
            to_change = N // 5
            for _ in range(to_change):
                tokens, labels = self.augment(tokens, labels, op="shuffle")
        elif op[:6] == "multi_":
            N = len(tokens)
            rep = 1 + N // 10
            new_op = op[6:]
            for _ in range(rep):
                tokens, labels = self.augment(tokens, labels, op=new_op)
        else:
            tokens, labels = self.augment(tokens, labels, op=op, args=args)
        results = " ".join(tokens)
        return results

    def sample_span(self, tokens, labels, span_len=3):
        candidates = []
        for idx, token in enumerate(tokens):
            if (
                idx + span_len - 1 < len(labels)
                and "".join(labels[idx : idx + span_len]) == "O" * span_len
            ):
                candidates.append((idx, idx + span_len - 1))
        if len(candidates) <= 0:
            return -1, -1
        return random.choice(candidates)

    def sample_position(self, tokens, labels, tfidf=False):
        # candidates = []
        # for idx, token in enumerate(tokens):
        #     if labels[idx] == 'O':
        #         candidates.append(idx)
        # if len(candidates) <= 0:
        #     return -1
        # return random.choice(candidates)

        candidates = []
        for idx, token in enumerate(tokens):
            if labels[idx] == "O":
                candidates.append(idx)

        if len(candidates) <= 0:
            return -1
        if tfidf:
            oov_th = math.log(1e8)
            weight = {}
            max_weight = 0.0
            for idx, token in enumerate(tokens):
                token = token.lower()
                # if token not in index:
                #     continue
                if token not in self.idf_dict:
                    self.idf_dict[token] = oov_th
                if token not in weight:
                    weight[token] = 0.0
                # weight[token] += index[token]['idf']
                weight[token] += self.idf_dict[token]
                max_weight = max(max_weight, weight[token])

            weights = []
            for idx in candidates:
                weights.append(max_weight - weight[tokens[idx].lower()] + 1e-6)

            weights = np.array(weights) / sum(weights)

            return np.random.choice(candidates, 1, p=weights)[0]
        else:
            return random.choice(candidates)

    def sample_token(self, token, max_candidates=10):
        """Randomly sample a token's similar token stored in the index

        Args:
            token (str): the input token
            same_length (bool, optional): whether the return token should have the same
                length in BERT
            max_candidates (int, optional): the maximal number of candidates
                to be sampled

        Returns:
            str: the sampled token (unchanged if the input is not in index)
        """
        token = token.lower()
        # index = self.index['token']
        # if token in index and index[token]['similar_words'] != None:
        #     candidates = [t for t, _ in \
        #                   index[token]['similar_words'][:max_candidates]]
        candidates = []
        syns = wordnet.synsets(token)
        for syn in syns:
            for lem in syn.lemmas():
                w = lem.name().lower()
                if w != token and w not in candidates and "_" not in w:
                    candidates.append(w)

        # print(candidates)
        if len(candidates) <= 0:
            return token
        else:
            return random.choice(candidates)

if __name__ == "__main__":
    ag = Augmenter()
    text = "COL content VAL vldb conference papers 2020-01-01 COL year VAL 2020 [SEP] COL content VAL sigmod conference 2010 papers 2019-12-31 COL year VAL 2019"
    for op in ["del", "shuffle", "insert", "all"]:
        print(op)
        print(ag.augment_sent(text, op=op))
    ag = Augmenter(idf_fn="wikitext-idf.dat")
    text = "demonstrates that the director of such hollywood blockbusters as patriot games can still turn out a small , personal film with an emotional wallop"
    for op in ["token_repl", "token_repl_tfidf", "token_del", "token_del_tfidf"]:
        print(op)
        print(ag.augment_sent(text, op=op))

    text_list = [
        "adobe creative suite 3 ( cs3 ) design suite standard upgrade ( upsell ) mac",
        "advances in databases and information systems ( adbis )",
        "effective timestamping in relational databases",
        "tutorial : ldap directory services - just another database application?",
    ]

    for text in text_list:
        for op in ["token_del", "shuffle"]:
            print(op)
            print(ag.augment_sent(text, op=op))

    for op in ["corrupt_plus", "corrupt_minus"]:
        print(op)
        print(ag.augment_sent(text, op=op))

    # text = "it's a familiar story, but one that is presented with great sympathy"
    text = "Where is the Orange Bowl?"
    for op in ["corrupt_minus", "corrupt_plus"]:
        text = ag.augment_sent(text, op=op)
    print(text)


try:
    import generative_playground
except:
    import sys, os, inspect

    my_location = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
    sys.path.append('../../../../../..')
    # sys.path.append('../../../../DeepRL')
    # sys.path.append('../../../../../transformer_pytorch')
import numpy as np
# from deep_rl import *
# from generative_playground.models.problem.rl.network_heads import CategoricalActorCriticNet
# from generative_playground.train.rl.run_iterations import run_iterations
from generative_playground.molecules.rdkit_utils.rdkit_utils import num_atoms, num_aromatic_rings, NormalizedScorer
# from generative_playground.models.problem.rl.DeepRL_wrappers import BodyAdapter, MyA2CAgent
from generative_playground.molecules.model_settings import get_settings
from generative_playground.molecules.train.pg.hypergraph.main_train_policy_gradient_minimal import train_policy_gradient
from generative_playground.codec.hypergraph_grammar import GrammarInitializer

batch_size = 8# 20
drop_rate = 0.5
molecules = True
grammar_cache = 'hyper_grammar_guac_10k_with_clique_collapse.pickle'#'hyper_grammar.pickle'
grammar = 'hypergraph:' + grammar_cache
# settings = get_settings(molecules, grammar)
# max_steps = 277  # settings['max_seq_length']
invalid_value = -3.5
atom_penalty = lambda x: -0.05*(np.maximum(37, np.array(num_atoms(x)))-37)
scorer = NormalizedScorer(invalid_value=invalid_value, normalize_scores=True)
reward_fun = lambda x: np.tanh(0.1*scorer(x)) + atom_penalty(x) # lambda x: reward_aromatic_rings(x)#
# later will run this ahead of time
# gi = GrammarInitializer(grammar_cache)

max_steps = 60
root_name = 'classic_logP'
model, gen_fitter, disc_fitter = train_policy_gradient(molecules,
                                                       grammar,
                                                       EPOCHS=100,
                                                       BATCH_SIZE=batch_size,
                                                       reward_fun_on=reward_fun,
                                                       max_steps=max_steps,
                                                       lr_on=3e-5,
                                                       lr_discrim=5e-4,
                                                       discrim_wt=0.3,
                                                       p_thresh=-10,
                                                       drop_rate=drop_rate,
                                                       reward_sm=0.5,
                                                       decoder_type='attn_graph_node',  # 'attention',
                                                       plot_prefix='',
                                                       dashboard=root_name,  # 'policy gradient',
                                                       save_file_root_name=root_name,
                                                       smiles_save_file=root_name + '_smiles',
                                                       on_policy_loss_type='advantage_record',
                                                       half_float=False,
                                                       priors='conditional')
# preload_file='policy_gradient_run.h5')

while True:
    next(gen_fitter)
    for _ in range(1):
        next(disc_fitter)