import cv2
import scipy
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from datetime import datetime
from pprint import pprint

np.set_printoptions(threshold=np.inf, linewidth=np.inf)

def detect(interpreter, x, ishape):
	'''
	'''

	batchx_4dtensor = tf.expand_dims(input=x, axis=0)
	interpreter.set_tensor(tensor_index=input_details[0]['index'], value=batchx_4dtensor)
	interpreter.invoke()

	valid_outputs = interpreter.get_tensor(output_details[0]['index'])
	detection = interpreter.get_tensor(output_details[1]['index'])

	bbox2d = detection[:valid_outputs[0], :]
	return bbox2d

test_file = 'believer.mp4'
result_file = 'believer_detected.mp4'
output_path = 'output'
ishape = [256, 256, 3]

interpreter = tf.lite.Interpreter(model_path='{}/det_model.tflite'.format(output_path))
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# pprint(input_details)
# pprint(output_details)

cap = cv2.VideoCapture(test_file)
# cap.set(cv2.CAP_PROP_POS_FRAMES, 0)

processed_images = []

while True:
	# Quit with 'q' press
	if cv2.waitKey(1) & 0xFF == ord('q'):
		break

	success, pix = cap.read()
	if success is not True:
		break

	pix = pix[60:-60, :, :]
	h, w, _ = pix.shape
	size = min(h, w)
	origin_y = 0
	origin_x = (w - h)//2
	if h > w:
		origin_y = (h - w)//2
		origin_x = 0

	pix = pix[origin_y:origin_y+size, origin_x:origin_x+size, :]
	pix = cv2.resize(pix, (ishape[1], ishape[0]))

	# print('Start: {}'.format(datetime.now().time()))
	x = tf.constant(value=pix, dtype='float32') # (h, w, 3)
	bbox2d = detect(interpreter=interpreter, x=x, ishape=ishape)
	# print('End: {}'.format(datetime.now().time()))

	for y1, x1, y2, x2, _ in bbox2d:
		cv2.rectangle(pix, (x1, y1), (x2, y2), [255, 255, 0], 1)

	# Display frame
	cv2.imshow('frame', pix)

	processed_images.append(pix)

cap.release()
cv2.destroyAllWindows()

out = cv2.VideoWriter(result_file, cv2.VideoWriter_fourcc(*'MP4V'), 24, (ishape[1], ishape[0]))
for i in range(len(processed_images)):
    out.write(processed_images[i])
out.release()


#  Copyright 2022 Deep Learning on Flink Authors
#
#  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 base64
import logging
import os
import pickle
from typing import Dict, Any, List, Optional, Type, Mapping, Union

import numpy as np
import tensorflow as tf
from pyflink.common import Row
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.ml.core.api import Estimator, Model, M, T
from pyflink.ml.core.param import Param
from pyflink.ml.util import read_write_utils
from pyflink.table import Table, StatementSet, TableSchema, TableFunction, \
    FunctionContext
from pyflink.table.types import CharType, VarCharType, IntType, \
    BigIntType, FloatType, DoubleType, DataType
from pyflink.table.udf import udtf

from dl_on_flink_tensorflow.tf_cluster_config import TFClusterConfig
from dl_on_flink_tensorflow.flink_ml.tf_estimator_constants import \
    INPUT_TYPES, INPUT_COL_NAMES, FEATURE_COLS, \
    LABEL_COL, MODEL_SAVE_PATH, BATCH_SIZE, MODEL_FACTORY_BASE64, MAX_EPOCHS
from dl_on_flink_tensorflow.flink_ml.tf_model_factory import \
    SimpleTFModelFactory, TFModelFactory
from dl_on_flink_tensorflow.flink_ml.tf_multi_worker_entry import \
    tf_multi_worker_entry
from dl_on_flink_tensorflow.tf_utils import train

logger = logging.getLogger(__file__)

FLINK_TYPE_TO_DL_ON_FLINK_TYPE: Mapping[Type, str] = {
    CharType: "STRING",
    VarCharType: "STRING",
    IntType: "INT_32",
    BigIntType: "INT_64",
    FloatType: "FLOAT_32",
    DoubleType: "FLOAT_64"
}

class TFEstimator(Estimator):

    def __init__(self, statement_set: StatementSet,
                 model: tf.keras.Model,
                 loss: Union[tf.keras.losses.Loss, str],
                 optimizer: Union[tf.keras.optimizers.Optimizer, str],
                 worker_num: int,
                 feature_cols: List[str],
                 label_col: str,
                 max_epochs: int = 1,
                 batch_size: Optional[int] = 32,
                 cluster_config_properties: Optional[Mapping[str, str]] = None):
        """
        A FlinkML Estimator implementation to distributed train a Tensorflow
        Keras Model with MultiWorkerMirroredStrategy.

        :param statement_set: The statement set created by the
        StreamTableEnvironment.
        :param model: The Keras model to train. We only support models with one
        output.
        :param loss: The loss function. It should be keras.losses.Loss instance
        or str.
        :param optimizer: The optimizer. It should be keras.optimizers.Optimizer
        instance or str.
        :param worker_num: The number of workers to run distributed model
        training.
        :param feature_cols: The name of the feature columns.
        :param label_col: The name of the label_col.
        :param max_epochs: Maximum number of epoch to train the model.
        :param batch_size: The batch size default to 32.
        :param cluster_config_properties: Extra cluster config properties.
        """
        self.max_epochs = max_epochs
        self.optimizer = optimizer
        self.loss = loss
        self.model = model
        self.batch_size = batch_size
        self.label_col = label_col
        self.feature_cols = feature_cols
        self.worker_num = worker_num
        self._statement_set = statement_set
        self.cluster_config_properties = \
            cluster_config_properties if cluster_config_properties is not None else {}

    def fit(self, *inputs: Table) -> 'TFModel':
        if len(inputs) != 1:
            raise ValueError("Only one input table is allowed.")

        self._verify_input_table(inputs[0])

        input_table = inputs[0]

        tf_cluster_config_builder = TFClusterConfig.new_builder()

        predict_col_data_type = input_table.get_schema().get_field_data_type(
            self.label_col)
        tf_cluster_config_builder.set_worker_count(self.worker_num) \
            .set_node_entry(tf_multi_worker_entry) \
            .set_property(INPUT_COL_NAMES,
                          self._get_column_names(input_table.get_schema())) \
            .set_property(FEATURE_COLS, ",".join(self.feature_cols)) \
            .set_property(LABEL_COL, self.label_col) \
            .set_property(BATCH_SIZE, str(self.batch_size)) \
            .set_property(INPUT_TYPES,
                          self._get_input_type(input_table.get_schema())) \
            .set_property(MAX_EPOCHS, str(self.max_epochs))

        for k, v in self.cluster_config_properties.items():
            tf_cluster_config_builder.set_property(k, v)

        tf_model_factory = SimpleTFModelFactory(model=self.model,
                                                loss=self.loss,
                                                optimizer=self.optimizer)

        tf_cluster_config_builder.set_property(MODEL_FACTORY_BASE64,
                                               self._pickle_model_factory(
                                                   tf_model_factory))

        return TFModel(tf_cluster_config_builder=tf_cluster_config_builder,
                       predict_col_data_type=predict_col_data_type,
                       statement_set=self._statement_set, input_table=inputs[0])

    def save(self, path: str) -> None:
        raise Exception("TFEstimator does not support save and load")

    @classmethod
    def load(cls, env: StreamExecutionEnvironment, path: str) -> 'TFEstimator':
        raise Exception("TFEstimator does not support save and load")

    def get_param_map(self) -> Dict['Param[Any]', Any]:
        return {}

    @staticmethod
    def _pickle_model_factory(model_factory: TFModelFactory) -> str:
        return base64.encodebytes(pickle.dumps(model_factory)) \
            .decode('utf-8')

    @staticmethod
    def _get_input_type(schema: TableSchema) -> str:
        data_types = schema.get_field_data_types()
        dl_on_flink_types = []
        for data_type in data_types:
            data_type = type(data_type)
            if data_type not in FLINK_TYPE_TO_DL_ON_FLINK_TYPE:
                raise TypeError(f"Unsupported type of column {data_type}")
            dl_on_flink_types.append(
                FLINK_TYPE_TO_DL_ON_FLINK_TYPE[data_type])
        return ",".join(dl_on_flink_types)

    @staticmethod
    def _get_column_names(schema: TableSchema):
        return ",".join(schema.get_field_names())

    def _verify_input_table(self, table: Table):
        for feature_col in self.feature_cols:
            assert feature_col in table.get_schema().get_field_names()

        assert self.label_col in table.get_schema().get_field_names()

class Predict(TableFunction):

    def __init__(self, tf_cluster_config: TFClusterConfig):
        super().__init__()
        self._model: Optional[tf.keras.Model] = None
        self._tf_cluster_config = tf_cluster_config
        self._model_path = tf_cluster_config.get_property(MODEL_SAVE_PATH)
        self._feature_cols = tf_cluster_config.get_property(FEATURE_COLS) \
            .split(",")
        self._label_col = tf_cluster_config.get_property(LABEL_COL)

    def open(self, function_context: FunctionContext):
        import tensorflow as tf
        if not tf.io.gfile.exists(self._model_path):
            raise FileNotFoundError(
                f"Model is not found at {self._model_path}.")
        self._model = tf.keras.models.load_model(self._model_path)

    def eval(self, row: Row):
        row_dict = row.as_dict()
        features = {}
        for feature_col in self._feature_cols:
            if feature_col not in row_dict:
                raise KeyError(f"{feature_col} not in the given row")
            features[feature_col] = np.array([row_dict[feature_col]])
        res = self._model.predict(features)[0]
        return Row(*row, res)

class TFModel(Model):

    def __init__(self,
                 predict_col_data_type: DataType,
                 tf_cluster_config_builder: TFClusterConfig.Builder = None,
                 tf_cluster_config: TFClusterConfig = None,
                 statement_set: StatementSet = None,
                 input_table=None):
        self.predict_col_data_type = predict_col_data_type
        self.input_table = input_table
        self.statement_set = statement_set
        self.tf_cluster_config_builder = tf_cluster_config_builder
        self.tf_cluster_config = tf_cluster_config

    def transform(self, *inputs: Table) -> List[Table]:
        if self.tf_cluster_config is None \
                or self.tf_cluster_config.get_property(MODEL_SAVE_PATH) is None:
            raise Exception("TFClusterConfig is missing or model path does "
                            "not exist. Please invoke save and "
                            "StatementSet#exeucte before transform.")
        table = inputs[0]
        row_data_type = table.get_schema().to_row_data_type()
        field_types = row_data_type.field_types()
        field_names = row_data_type.field_names()

        label_col = self.tf_cluster_config.get_property(LABEL_COL)
        table = table.flat_map(
            udtf(f=Predict(self.tf_cluster_config),
                 result_types=field_types + [self.predict_col_data_type]))
        table = table.alias(*field_names, label_col)
        return [table]

    def save(self, path: str) -> None:
        self.tf_cluster_config_builder \
            .set_property(MODEL_SAVE_PATH, os.path.join(path, "model.h5"))
        self.tf_cluster_config = self.tf_cluster_config_builder.build()
        max_epoch = int(self.tf_cluster_config.get_property(MAX_EPOCHS))

        read_write_utils.save_metadata(self, path, {
            "tf_cluster_config": self.tf_cluster_config,
            "predict_col_data_type": self.predict_col_data_type
        })
        train(self.statement_set, self.tf_cluster_config, self.input_table,
              max_epoch=max_epoch)

    @classmethod
    def load(cls, env: StreamExecutionEnvironment, path: str) -> 'TFModel':
        meta = read_write_utils.load_metadata(path)
        model = TFModel(tf_cluster_config=meta["tf_cluster_config"],
                        predict_col_data_type=meta["predict_col_data_type"])
        return model

    def get_param_map(self) -> Dict['Param[Any]', Any]:
        return {}


# import the necessary packages
from imutils.video import VideoStream
import numpy as np
import imutils
import time
import cv2
import tensorflow as tf

from keras.preprocessing import image
#use this latter for bounding box
resMap = {
        1 : 'Mask On',
        0 : 'Kindly Wear Mask'
    }

colorMap = {
        1 : (0,255,0),
        0 : (0,0,255)
    }

def prepImg(pth):
    return cv2.resize(pth,(224,224)).reshape(1,224,224,3)/255.0

#defining model paths
caffe_model = "Face detector/res10_300x300_ssd_iter_140000_fp16.caffemodel"
prototext_path = "Face detector/deploy.prototxt.txt"
mask_path = 'mask_model/face_mask_model.h5'

#Load Models
print("Loading models...................\n")

#Loading mask classification model
print("[INFO] Loading model 1...\n")
model = tf.keras.models.load_model(mask_path)
# model.load_weights(mask_path)

#Loading face detction model
print("[INFO] Loading model 2...\n")
net = cv2.dnn.readNetFromCaffe(prototext_path,caffe_model)

# initialize the video stream to get the video frames
vs_num = input('[INPUT] Enter your video streaming devices unique id: ')
print("[INFO] starting video stream...")
vs = VideoStream(src=int(vs_num)).start()
time.sleep(2.0)

#loop the frames from the  VideoStream
while True :
    #Get the frames from the video stream and resize to 400 px
    frame = vs.read()
    frame = imutils.resize(frame,width=400)

    # extract the dimensions , Resize image into 300x300 and converting image into blobFromImage
    (h, w) = frame.shape[:2]
    # blobImage convert RGB (104.0, 177.0, 123.0)
    blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                 (300, 300), (104.0, 177.0, 123.0))

    # passing blob through the network to detect and pridiction
    net.setInput(blob)
    detections = net.forward()

    # loop over the detections
    for i in range(0, detections.shape[2]):
        # extract the confidence and prediction

        confidence = detections[0, 0, i, 2]

        # filter detections by confidence greater than the minimum confidence
        if confidence > 0.8 :

            # Determine the (x, y)-coordinates of the bounding box for the face
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")

            (startX, startY, endX, endY) = (startX-15, startY-10, endX+15, endY+15)

            #Determining bounding box color and text
            text = "{:.2f}%".format(confidence * 100)
            y = startY - 10 if startY - 10 > 10 else startY + 10

            #get the image to pass to the mask detction model
            img = frame[startX:endX, startY:endY]

            pred = model.predict(prepImg(img))

            pred = np.round(pred)

            #draw the bounding box of the face along with the associated text
            cv2.rectangle(frame,(startX,startY),(endX,endY),
                            colorMap[pred[0][0]],2)
            cv2.putText(frame, resMap[pred[0][0]],(startX,y),
                        cv2.FONT_HERSHEY_SIMPLEX,0.8,(255,255,255),2)

    # show the output frame
    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF

    # if the `q` key was pressed, break from the loop
    if key == ord("q"):
        break

# do a bit of cleanup
cv2.destroyAllWindows()
vs.stop()


import os
import numpy as np
import pickle
import re

def load_dict(dict_path):
    """Loads a dictionary in pickle format."""
    with open(dict_path,'rb') as fh:
        d = pickle.load(fh)
    return d

def parse_search(workspace,corpus_name,node,pos,case_sensitive,regexp):
    """Parses the input search and returns a tuple with encoded word and pos node.
    """
    # load dictionaries for encoding words to numbers
    if case_sensitive == False:
        dict_words = load_dict(workspace+corpus_name+'/data/idx/wordslc.pickle') # lowercase
    else:
        dict_words = load_dict(workspace+corpus_name+'/data/idx/words.pickle')
        dict_words = {v: k for k, v in dict_words.items()}
    dict_tags = load_dict(workspace+corpus_name+'/data/idx/tags.pickle')
    dict_tags = {v: k for k, v in dict_tags.items()}
    # parse node
    n = []
    if case_sensitive == False:
        node = node.lower()
    for w in node.strip().split(' '):
        if len(w.strip()) != 0:
            n.append(w.strip())
    # parse pos
    if pos != None:
        p = []
        if type(pos) == list:
            for t in pos:
                if len(t.strip())!= 0:
                    p.append(t.strip())
        else:
            for t in pos.strip().split(' '):
                if len(t.strip())!=0:
                    p.append(t.strip())
    # encode node
    encoded_node = []
    if regexp == False:
        for w in n:
            if w in dict_words:
                arr = np.array(dict_words[w],dtype=int)
                encoded_node.append(arr)
            else: # if any word has no match, return None
                encoded_node = None
                break
    else: # use regexp:
        patterns = []
        for ptrn in n:
            patterns.append(re.compile(ptrn))
        for ptrn in patterns:
            a = []
            for k in dict_words:
                if ptrn.match(k) != None:
                    a+=dict_words[k]
            if len(a)!=0:
                arr = np.array(a,dtype=int)
                encoded_node.append(arr)
            else: # if any pattern has no match, return None
                encoded_node = None
                break
    # encode pos
    if pos != None:
        encoded_pos = []
        for t in p:
            if t in dict_tags:
                arr = np.array(dict_tags[t],dtype=int)
                encoded_pos.append(arr)
            else: # if any tag has no match, return None
                encoded_pos = None
                break
    else:
        encoded_pos = None
    return (encoded_node, encoded_pos)

def search_node(npy_path,encoded_node,encoded_pos,horizon):
    """Searches the node in every text file and gets contexts."""
    size = len(encoded_node)
    files = os.listdir(npy_path)
    contexts = []
    for filename in files:
        # maker arr concatenating with arrays of zeros for preventing index out of bounds
        arr = np.concatenate([np.zeros((horizon+size,4),dtype=int),np.load(npy_path + filename),np.zeros((horizon+size,4),dtype=int)])
        # get matching indexes
        ix = arr_search_indexes(arr,encoded_node,encoded_pos)
        # get contexts from indexes
        contexts += get_context(arr,ix,horizon,size)
    return contexts

def arr_search_indexes(arr,search_w,search_t):
    """Searches the node in arr and returns the matching indexes. (Max. 3 words)"""
    size = len(search_w)
    if search_t == None:
        # search only words
        if size == 1:
            ix = np.isin(arr[:,[0]],search_w)
            ix = np.where(ix)[0]
        elif size == 2:
            # w1
            x1 = np.isin(arr[:,[0]],search_w[0])
            if np.any(x1)==False:
                return np.where(x1)[0]
            # w2
            x2 = np.isin(arr[:,[0]],search_w[1])
            if np.any(x2)==False:
                return np.where(x2)[0]
            # match indexes
            x = np.arange(0,x1.shape[0]-1)
            y = np.arange(1,x1.shape[0])
            # return indexes
            ix = np.where(x1[x] & x2[y])[0]
            return ix
        elif size == 3:
            # w1
            x1 = np.isin(arr[:,[0]],search_w[0])
            if np.any(x1)==False:
                return np.where(x1)[0]
            # w2
            x2 = np.isin(arr[:,[0]],search_w[1])
            if np.any(x2)==False:
                return np.where(x2)[0]
            # w3
            x3 = np.isin(arr[:,[0]],search_w[2])
            if np.any(x3)==False:
                return np.where(x3)[0]
            # match indexes
            x = np.arange(0,x1.shape[0]-2)
            y = np.arange(1,x1.shape[0]-1)
            z = np.arange(2,x1.shape[0])
            # return indexes
            ix = np.where(x1[x] & x2[y] & x3[z])[0]
    else:
        # search words and POS
        if size == 1:
            iw = np.isin(arr[:,[0]],search_w)
            it = np.isin(arr[:,[1]],search_t)
            ix = np.where(iw & it)[0]
        elif size == 2:
            # w1 t1
            x1 = np.isin(arr[:,[0]],search_w[0])
            t1 = np.isin(arr[:,[1]],search_t[0])
            if np.any(x1)==False:
                return np.where(x1)[0]
            if np.any(t1)==False:
                return np.where(t1)[0]
            # w2 t2
            x2 = np.isin(arr[:,[0]],search_w[1])
            t2 = np.isin(arr[:,[1]],search_t[1])
            if np.any(x2)==False:
                return np.where(x2)[0]
            if np.any(t2)==False:
                return np.where(t2)[0]
            # match indexes
            x = np.arange(0,x1.shape[0]-1)
            y = np.arange(1,x1.shape[0])
            # return indexes
            ix = np.where(x1[x] & t1[x] & x2[y] & t2[y])[0]
        elif size == 3:
            # w1 t1
            x1 = np.isin(arr[:,[0]],search_w[0])
            t1 = np.isin(arr[:,[1]],search_t[0])
            if np.any(x1)==False:
                return np.where(x1)[0]
            if np.any(t1)==False:
                return np.where(t1)[0]
            # w2 t2
            x2 = np.isin(arr[:,[0]],search_w[1])
            t2 = np.isin(arr[:,[1]],search_t[1])
            if np.any(x2)==False:
                return np.where(x2)[0]
            if np.any(t2)==False:
                return np.where(t2)[0]
            # w3 t3
            x3 = np.isin(arr[:,[0]],search_w[2])
            t3 = np.isin(arr[:,[1]],search_t[2])
            if np.any(x3)==False:
                return np.where(x3)[0]
            if np.any(t3)==False:
                return np.where(t3)[0]
            # match indexes
            x = np.arange(0,x1.shape[0]-2)
            y = np.arange(1,x1.shape[0]-1)
            z = np.arange(2,x1.shape[0])
            # return indexes
            ix = np.where(x1[x] & t1[x] & x2[y] & t2[y] & x3[z] & t3[z])[0]
    return ix

def get_context(arr,indexes,horizon,size):
    """Gets the contexts from indexes."""
    neg = np.arange(horizon,0,-1)
    pos = np.arange(0,horizon,1)
    nsize = np.arange(0,size)
    contexts=[]
    for i in indexes:
        l = arr[i-neg,0]         # left horizon
        n = arr[i+nsize,0]       # node
        r = arr[(i+size)+pos,0]  # right horizon
        t = i                    # token id
        s = arr[i,2]             # sentence id
        f = arr[i,3]             # file id
        contexts.append((l,n,r,t,s,f))
    return contexts

def translate(contexts,dict_words,dict_files):
    """Translate numbers back to words in a list format for kitconc KIWC object."""
    dict_files = {v: k for k, v in dict_files.items()} # reverse keys
    kwic = []
    i = 0
    for context in contexts:
        i+=1
        l = []
        n = []
        r = []
        # left
        for i in context[0]:
            if i != 0:
                l.append(dict_words[i])
        # node
        for i in context[1]:
            if i != 0:
                n.append(dict_words[i])
        # right
        for i in context[2]:
            if i != 0:
                r.append(dict_words[i])
        # add to list
        kwic.append((i, ' '.join(l), ' '.join(n),' '.join(r), dict_files[context[5]][0],str(context[3]),str(context[4]),str(context[5])))
    return kwic

def make_kwic(workspace,corpus_name,node,pos,case_sensitive,regexp,horizon):
    encoded_node,encoded_pos = parse_search(workspace,corpus_name,node,pos,case_sensitive,regexp)
    if encoded_node != None: # search node is OK
        nl = len(encoded_node)
        contexts = search_node(workspace+corpus_name + '/data/npy/',encoded_node,encoded_pos,horizon)
        if len(contexts)!=0:
            dict_words = load_dict(workspace+corpus_name + '/data/idx/words.pickle')
            dict_files = load_dict(workspace+corpus_name + '/data/idx/filenames.pickle')
            kwic = translate(contexts,dict_words,dict_files)
        else:
            kwic=[]
    else: # no match in search node, set empty
        nl = 0
        kwic = []
    return (nl,kwic)


# Desc: orig(v0) + gt of target style image

import os
import sys
import time
import random
import string
import argparse
import glob
import math
import re

import torch
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.optim as optim
import torch.utils.data
from torch.optim import lr_scheduler
import numpy as np
import torchvision.models as models
from torchvision import transforms, utils
import torch.nn.functional as F
import torch.autograd as autograd
import pdb
import html_visual as html
from skimage.metrics import peak_signal_noise_ratio, mean_squared_error, structural_similarity
from nltk.metrics.distance import edit_distance
from skimage.transform import resize
from skimage.io import imsave
import skimage
import cv2

from utils import CTCLabelConverter, AttnLabelConverter, Averager
from dataset import hierarchical_dataset, AlignPairCollate, AlignPairImgCollate, AlignPairImgCollate_Test, AlignSynthTextCollate, Batch_Balanced_Dataset, tensor2im, save_image, phoc_gen, text_gen, text_gen_synth, LmdbStyleDataset, LmdbTestStyleContentDataset
from model import ModelV1, GlobalContentEncoder, VGGPerceptualEmbedLossModel, VGGFontModel
from modules.feature_extraction import ResNet_StyleExtractor, VGG_ContentExtractor, ResNet_StyleExtractor_WIN

import tflib as lib
import tflib.plot

sys.path.append('/private/home/pkrishnan/codes/st-scribe/stylegan2-pytorch/')

from model_word import GeneratorM2V4_2 as styleGANGen
from model_word import EncDiscriminator as styleGANDis
from non_leaking import augment
from distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle)

    if shuffle:
        return torch.utils.data.RandomSampler(dataset)

    else:
        return torch.utils.data.SequentialSampler(dataset)

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

def sample_data(loader, data_sampler, is_distributed):
    epochCntr=0
    while True:
        if is_distributed:
            data_sampler.set_epoch(epochCntr)

        for batch in loader:
            yield batch
        epochCntr+=1

def make_noise(z_code, batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.cat((z_code, torch.randn(batch, latent_dim, device=device)),1)

    noises = torch.cat((z_code.repeat(2,1,1), torch.randn(n_noise, batch, latent_dim, device=device)),2).unbind(0)

    return noises

def make_noise_style(batch, latent_dim, n_noise, device):
    if n_noise == 1:
        return torch.randn(batch, latent_dim, device=device)

    noises = torch.randn(n_noise, batch, latent_dim, device=device).unbind(0)

    return noises

def mixing_noise(batch, latent_dim, prob, device, z_code=None):
    if prob > 0 and random.random() < prob:
        if z_code is None:
            return make_noise_style(batch, latent_dim, 2, device)
        else:
            return make_noise(z_code, batch, latent_dim, 2, device)
    else:
        if z_code is None:
            return [make_noise_style(batch, latent_dim, 1, device)]
        else:
            return [make_noise(z_code, batch, latent_dim, 1, device)]

def mixing_noise_style(batch, latent_dim, prob, device):
    if prob > 0 and random.random() < prob:
        return make_noise_style(batch, latent_dim, 2, device)

    else:
        return [make_noise_style(batch, latent_dim, 1, device)]

def set_grad_none(model, targets):
    for n, p in model.named_parameters():
        if n in targets:
            p.grad = None

def test(opt):
    lib.print_model_settings(locals().copy())

    if 'Attn' in opt.Prediction:
        converter = AttnLabelConverter(opt.character)
        text_len = opt.batch_max_length+2
    else:
        converter = CTCLabelConverter(opt.character)
        text_len = opt.batch_max_length

    opt.classes = converter.character

    """ dataset preparation """
    if not opt.data_filtering_off:
        print('Filtering the images containing characters which are not in opt.character')
        print('Filtering the images whose label is longer than opt.batch_max_length')
        # see https://github.com/clovaai/deep-text-recognition-benchmark/blob/6593928855fb7abb999a99f428b3e4477d4ae356/dataset.py#L130

    # AlignCollate_valid = AlignPairImgCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
    AlignCollate_valid = AlignPairImgCollate_Test(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)

    valid_dataset = LmdbTestStyleContentDataset(root=opt.test_data, opt=opt, dataMode=opt.realVaData)
    test_data_sampler = data_sampler(valid_dataset, shuffle=False, distributed=False)
    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=opt.batch_size,
        shuffle=False,  # 'True' to check training progress with validation function.
        sampler=test_data_sampler,
        num_workers=int(opt.workers),
        collate_fn=AlignCollate_valid, pin_memory=True, drop_last=False)

    print('-' * 80)

    AlignCollate_text = AlignSynthTextCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
    text_dataset = text_gen_synth(opt)
    text_data_sampler = data_sampler(text_dataset, shuffle=True, distributed=False)
    text_loader = torch.utils.data.DataLoader(
        text_dataset, batch_size=opt.batch_size,
        shuffle=False,
        sampler=text_data_sampler,
        num_workers=int(opt.workers),
        collate_fn=AlignCollate_text,
        drop_last=True)
    opt.num_class = len(converter.character)

    text_loader = sample_data(text_loader, text_data_sampler, False)

    c_code_size = opt.latent
    if opt.cEncode == 'mlp':
        cEncoder = GlobalContentEncoder(opt.num_class, text_len, opt.char_embed_size, c_code_size).to(device)
    elif opt.cEncode == 'cnn':
        if opt.contentNorm == 'in':
            cEncoder = ResNet_StyleExtractor_WIN(1, opt.latent).to(device)
        else:
            cEncoder = ResNet_StyleExtractor(1, opt.latent).to(device)
    if opt.styleNorm == 'in':
        styleModel = ResNet_StyleExtractor_WIN(opt.input_channel, opt.latent).to(device)
    else:
        styleModel = ResNet_StyleExtractor(opt.input_channel, opt.latent).to(device)

    ocrModel = ModelV1(opt).to(device)

    g_ema = styleGANGen(opt.size, opt.latent, opt.latent, opt.n_mlp, content_dim=c_code_size, channel_multiplier=opt.channel_multiplier).to(device)
    g_ema.eval()

    bestModelError=1e5

    ## Loading pre-trained files
    print(f'loading pretrained ocr model from {opt.saved_ocr_model}')
    checkpoint = torch.load(opt.saved_ocr_model, map_location=lambda storage, loc: storage)
    ocrModel.load_state_dict(checkpoint)

    print(f'loading pretrained synth model from {opt.saved_synth_model}')
    checkpoint = torch.load(opt.saved_synth_model, map_location=lambda storage, loc: storage)

    cEncoder.load_state_dict(checkpoint['cEncoder'])
    styleModel.load_state_dict(checkpoint['styleModel'])
    g_ema.load_state_dict(checkpoint['g_ema'])

    iCntr=0
    evalCntr=0
    fCntr=0

    valMSE=0.0
    valSSIM=0.0
    valPSNR=0.0
    c1_s1_input_correct=0.0
    c2_s1_gen_correct=0.0
    c1_s1_input_ed_correct=0.0
    c2_s1_gen_ed_correct=0.0

    ims, txts = [], []

    for vCntr, (image_input_tensors, image_output_tensors, labels_gt, labels_z_c, labelSynthImg, synth_z_c, input_1_shape, input_2_shape) in enumerate(valid_loader):
        print(vCntr)

        if opt.debugFlag and vCntr >10:
            break

        image_input_tensors = image_input_tensors.to(device)
        image_output_tensors = image_output_tensors.to(device)

        if opt.realVaData and opt.outPairFile=="":
            # pdb.set_trace()
            labels_z_c, synth_z_c = next(text_loader)

        labelSynthImg = labelSynthImg.to(device)
        synth_z_c = synth_z_c.to(device)
        synth_z_c = synth_z_c[:labelSynthImg.shape[0]]

        text_z_c, length_z_c = converter.encode(labels_z_c, batch_max_length=opt.batch_max_length)
        text_gt, length_gt = converter.encode(labels_gt, batch_max_length=opt.batch_max_length)

        # print(labels_z_c)

        cEncoder.eval()
        styleModel.eval()
        g_ema.eval()

        with torch.no_grad():
            if opt.cEncode == 'mlp':
                z_c_code = cEncoder(text_z_c)
                z_gt_code = cEncoder(text_gt)
            elif opt.cEncode == 'cnn':
                z_c_code = cEncoder(synth_z_c)
                z_gt_code = cEncoder(labelSynthImg)

            style = styleModel(image_input_tensors)

            if opt.noiseConcat or opt.zAlone:
                style = mixing_noise(opt.batch_size, opt.latent, opt.mixing, device, style)
            else:
                style = [style]

            fake_img_c1_s1, _ = g_ema(style, z_gt_code, input_is_latent=opt.input_latent)
            fake_img_c2_s1, _ = g_ema(style, z_c_code, input_is_latent=opt.input_latent)

            currBatchSize = fake_img_c1_s1.shape[0]
            # text_for_pred = torch.LongTensor(opt.batch_size, opt.batch_max_length + 1).fill_(0).to(device)
            length_for_pred = torch.IntTensor([opt.batch_max_length] * currBatchSize).to(device)
            #Run OCR prediction
            if 'CTC' in opt.Prediction:
                preds = ocrModel(fake_img_c1_s1, text_gt, is_train=False, inAct = opt.taskActivation)
                preds_size = torch.IntTensor([preds.size(1)] * currBatchSize)
                _, preds_index = preds.max(2)
                preds_str_fake_img_c1_s1 = converter.decode(preds_index.data, preds_size.data)

                preds = ocrModel(fake_img_c2_s1, text_z_c, is_train=False, inAct = opt.taskActivation)
                preds_size = torch.IntTensor([preds.size(1)] * currBatchSize)
                _, preds_index = preds.max(2)
                preds_str_fake_img_c2_s1 = converter.decode(preds_index.data, preds_size.data)

                preds = ocrModel(image_input_tensors, text_gt, is_train=False)
                preds_size = torch.IntTensor([preds.size(1)] * image_input_tensors.shape[0])
                _, preds_index = preds.max(2)
                preds_str_gt_1 = converter.decode(preds_index.data, preds_size.data)

                preds = ocrModel(image_output_tensors, text_z_c, is_train=False)
                preds_size = torch.IntTensor([preds.size(1)] * image_output_tensors.shape[0])
                _, preds_index = preds.max(2)
                preds_str_gt_2 = converter.decode(preds_index.data, preds_size.data)

            else:

                preds = ocrModel(fake_img_c1_s1, text_gt[:, :-1], is_train=False, inAct = opt.taskActivation)  # align with Attention.forward
                _, preds_index = preds.max(2)
                preds_str_fake_img_c1_s1 = converter.decode(preds_index, length_for_pred)
                for idx, pred in enumerate(preds_str_fake_img_c1_s1):
                    pred_EOS = pred.find('[s]')
                    preds_str_fake_img_c1_s1[idx] = pred[:pred_EOS]  # prune after "end of sentence" token ([s])

                preds = ocrModel(fake_img_c2_s1, text_z_c[:, :-1], is_train=False, inAct = opt.taskActivation)  # align with Attention.forward
                _, preds_index = preds.max(2)
                preds_str_fake_img_c2_s1 = converter.decode(preds_index, length_for_pred)
                for idx, pred in enumerate(preds_str_fake_img_c2_s1):
                    pred_EOS = pred.find('[s]')
                    preds_str_fake_img_c2_s1[idx] = pred[:pred_EOS]  # prune after "end of sentence" token ([s])

                preds = ocrModel(image_input_tensors, text_gt[:, :-1], is_train=False)  # align with Attention.forward
                _, preds_index = preds.max(2)
                preds_str_gt_1 = converter.decode(preds_index, length_for_pred)
                for idx, pred in enumerate(preds_str_gt_1):
                    pred_EOS = pred.find('[s]')
                    preds_str_gt_1[idx] = pred[:pred_EOS]  # prune after "end of sentence" token ([s])

                preds = ocrModel(image_output_tensors, text_z_c[:, :-1], is_train=False)  # align with Attention.forward
                _, preds_index = preds.max(2)
                preds_str_gt_2 = converter.decode(preds_index, length_for_pred)
                for idx, pred in enumerate(preds_str_gt_2):
                    pred_EOS = pred.find('[s]')
                    preds_str_gt_2[idx] = pred[:pred_EOS]  # prune after "end of sentence" token ([s])

        pathPrefix = os.path.join(opt.valDir, opt.exp_iter)
        os.makedirs(os.path.join(pathPrefix), exist_ok=True)

        for trImgCntr in range(image_output_tensors.shape[0]):

            if opt.outPairFile!="":
                labelId = 'label-' + valid_loader.dataset.pairId[fCntr] + '-' + str(fCntr)
            else:
                labelId = 'label-%09d' % valid_loader.dataset.filtered_index_list[fCntr]
            #evaluations
            valRange = (-1,+1)
            # pdb.set_trace()
            # inpTensor = skimage.img_as_ubyte(resize(tensor2im(image_input_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0])))
            # gtTensor = skimage.img_as_ubyte(resize(tensor2im(image_output_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0])))
            # predTensor = skimage.img_as_ubyte(resize(tensor2im(fake_img_c2_s1[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0])))

            # inpTensor = resize(tensor2im(image_input_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]), anti_aliasing=True)
            # gtTensor = resize(tensor2im(image_output_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]), anti_aliasing=True)
            # predTensor = resize(tensor2im(fake_img_c2_s1[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]), anti_aliasing=True)

            inpTensor = F.interpolate(image_input_tensors[trImgCntr].unsqueeze(0),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]))
            gtTensor = F.interpolate(image_output_tensors[trImgCntr].unsqueeze(0),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]))
            predTensor = F.interpolate(fake_img_c2_s1[trImgCntr].unsqueeze(0),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]))
            predGTTensor = F.interpolate(fake_img_c1_s1[trImgCntr].unsqueeze(0),(input_1_shape[trImgCntr][1], input_1_shape[trImgCntr][0]))

            # inpTensor = cv2.resize(tensor2im(image_input_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr]))
            # gtTensor = cv2.resize(tensor2im(image_output_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr]))
            # predTensor = cv2.resize(tensor2im(fake_img_c2_s1[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr]))
            # predGTTensor = cv2.resize(tensor2im(fake_img_c1_s1[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr]))

            # inpTensor = cv2.medianBlur(cv2.resize(tensor2im(image_input_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr])),5)
            # gtTensor = cv2.medianBlur(cv2.resize(tensor2im(image_output_tensors[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr])),5)
            # predTensor = cv2.medianBlur(cv2.resize(tensor2im(fake_img_c2_s1[trImgCntr].clone().clamp_(min=valRange[0], max=valRange[1])),tuple(input_1_shape[trImgCntr])),5)
            # pdb.set_trace()

            if not(opt.realVaData):
                evalMSE = mean_squared_error(tensor2im(gtTensor.squeeze())/255, tensor2im(predTensor.squeeze())/255)
                # evalMSE = mean_squared_error(gtTensor/255, predTensor/255)
                evalSSIM = structural_similarity(tensor2im(gtTensor.squeeze())/255, tensor2im(predTensor.squeeze())/255, multichannel=True)
                # evalSSIM = structural_similarity(gtTensor/255, predTensor/255, multichannel=True)
                evalPSNR = peak_signal_noise_ratio(tensor2im(gtTensor.squeeze())/255, tensor2im(predTensor.squeeze())/255)
                # evalPSNR = peak_signal_noise_ratio(gtTensor/255, predTensor/255)
                # print(evalMSE,evalSSIM,evalPSNR)

                valMSE+=evalMSE
                valSSIM+=evalSSIM
                valPSNR+=evalPSNR

            #ocr accuracy
            # for gt, pred, pred_max_prob in zip(labels, preds_str, preds_max_prob):
            c1_s1_input_gt = labels_gt[trImgCntr]
            c1_s1_input_ocr = preds_str_gt_1[trImgCntr]
            c2_s1_gen_gt = labels_z_c[trImgCntr]
            c2_s1_gen_ocr = preds_str_fake_img_c2_s1[trImgCntr]

            if c1_s1_input_gt == c1_s1_input_ocr:
                c1_s1_input_correct += 1
            if c2_s1_gen_gt == c2_s1_gen_ocr:
                c2_s1_gen_correct += 1

            # ICDAR2019 Normalized Edit Distance
            if len(c1_s1_input_gt) == 0 or len(c1_s1_input_ocr) == 0:
                c1_s1_input_ed_correct += 0
            elif len(c1_s1_input_gt) > len(c1_s1_input_ocr):
                c1_s1_input_ed_correct += 1 - edit_distance(c1_s1_input_ocr, c1_s1_input_gt) / len(c1_s1_input_gt)
            else:
                c1_s1_input_ed_correct += 1 - edit_distance(c1_s1_input_ocr, c1_s1_input_gt) / len(c1_s1_input_ocr)

            if len(c2_s1_gen_gt) == 0 or len(c2_s1_gen_ocr) == 0:
                c2_s1_gen_ed_correct += 0
            elif len(c2_s1_gen_gt) > len(c2_s1_gen_ocr):
                c2_s1_gen_ed_correct += 1 - edit_distance(c2_s1_gen_ocr, c2_s1_gen_gt) / len(c2_s1_gen_gt)
            else:
                c2_s1_gen_ed_correct += 1 - edit_distance(c2_s1_gen_ocr, c2_s1_gen_gt) / len(c2_s1_gen_ocr)

            evalCntr+=1

            #save generated images
            if opt.visFlag and iCntr>500:
                pass
            else:
                try:
                    if iCntr == 0:
                        # update website
                        webpage = html.HTML(pathPrefix, 'Experiment name = %s' % 'Test')
                        webpage.add_header('Testing iteration')

                    iCntr += 1
                    img_path_c1_s1 = os.path.join(labelId+'_pred_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_fake_img_c1_s1[trImgCntr]+'.png')
                    img_path_gt_1 = os.path.join(labelId+'_gt_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_gt_1[trImgCntr]+'.png')
                    img_path_gt_2 = os.path.join(labelId+'_gt_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_gt_2[trImgCntr]+'.png')
                    img_path_c2_s1 = os.path.join(labelId+'_pred_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_fake_img_c2_s1[trImgCntr]+'.png')

                    ims.append([img_path_gt_1, img_path_c1_s1, img_path_gt_2, img_path_c2_s1])

                    content_c1_s1 = 'PSTYLE-1 '+'Text-1:' + labels_gt[trImgCntr]+' OCR:' + preds_str_fake_img_c1_s1[trImgCntr]
                    content_gt_1 = 'OSTYLE-1 '+'GT:' + labels_gt[trImgCntr]+' OCR:' + preds_str_gt_1[trImgCntr]
                    content_gt_2 = 'OSTYLE-1 '+'GT:' + labels_z_c[trImgCntr]+' OCR:'+preds_str_gt_2[trImgCntr]
                    content_c2_s1 = 'PSTYLE-1 '+'Text-2:' + labels_z_c[trImgCntr]+' OCR:'+preds_str_fake_img_c2_s1[trImgCntr]

                    txts.append([content_gt_1, content_c1_s1, content_gt_2, content_c2_s1])

                    utils.save_image(predGTTensor,os.path.join(pathPrefix,labelId+'_pred_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_fake_img_c1_s1[trImgCntr]+'.png'),nrow=1,normalize=True,range=(-1, 1))
                    # cv2.imwrite(os.path.join(pathPrefix,labelId+'_pred_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_fake_img_c1_s1[trImgCntr]+'.png'), predGTTensor)

                    # pdb.set_trace()
                    if not opt.zAlone:
                        utils.save_image(inpTensor,os.path.join(pathPrefix,labelId+'_gt_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_gt_1[trImgCntr]+'.png'),nrow=1,normalize=True,range=(-1, 1))
                        utils.save_image(gtTensor,os.path.join(pathPrefix,labelId+'_gt_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_gt_2[trImgCntr]+'.png'),nrow=1,normalize=True,range=(-1, 1))
                        utils.save_image(predTensor,os.path.join(pathPrefix,labelId+'_pred_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_fake_img_c2_s1[trImgCntr]+'.png'),nrow=1,normalize=True,range=(-1, 1))

                        # imsave(os.path.join(pathPrefix,labelId+'_gt_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_gt_1[trImgCntr]+'.png'), inpTensor)
                        # imsave(os.path.join(pathPrefix,labelId+'_gt_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_gt_2[trImgCntr]+'.png'), gtTensor)
                        # imsave(os.path.join(pathPrefix,labelId+'_pred_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_fake_img_c2_s1[trImgCntr]+'.png'), predTensor)

                        # cv2.imwrite(os.path.join(pathPrefix,labelId+'_gt_val_c1_s1_'+labels_gt[trImgCntr]+'_ocr_'+preds_str_gt_1[trImgCntr]+'.png'), inpTensor)
                        # cv2.imwrite(os.path.join(pathPrefix,labelId+'_gt_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_gt_2[trImgCntr]+'.png'), gtTensor)
                        # cv2.imwrite(os.path.join(pathPrefix,labelId+'_pred_val_c2_s1_'+labels_z_c[trImgCntr]+'_ocr_'+preds_str_fake_img_c2_s1[trImgCntr]+'.png'), predTensor)
                except:
                    print('Warning while saving validation image')

            fCntr += 1

    webpage.add_images(ims, txts, width=256, realFlag=opt.realVaData)
    webpage.save()

    avg_valMSE = valMSE/float(evalCntr)
    avg_valSSIM = valSSIM/float(evalCntr)
    avg_valPSNR = valPSNR/float(evalCntr)
    avg_c1_s1_input_wer = c1_s1_input_correct/float(evalCntr)
    avg_c2_s1_gen_wer = c2_s1_gen_correct/float(evalCntr)
    avg_c1_s1_input_cer = c1_s1_input_ed_correct/float(evalCntr)
    avg_c2_s1_gen_cer = c2_s1_gen_ed_correct/float(evalCntr)

    # if not(opt.realVaData):
    with open(os.path.join(opt.exp_dir,opt.exp_name,'log_test.txt'), 'a') as log:
        # training loss and validation loss
        if opt.realVaData:
            loss_log = f'Test Input Word Acc: {avg_c1_s1_input_wer:0.5f}, Test Gen Word Acc: {avg_c2_s1_gen_wer:0.5f}, Test Input Char Acc: {avg_c1_s1_input_cer:0.5f}, Test Gen Char Acc: {avg_c2_s1_gen_cer:0.5f}'
        else:
            loss_log = f'Test MSE: {avg_valMSE:0.5f}, Test SSIM: {avg_valSSIM:0.5f}, Test PSNR: {avg_valPSNR:0.5f}, Test Input Word Acc: {avg_c1_s1_input_wer:0.5f}, Test Gen Word Acc: {avg_c2_s1_gen_wer:0.5f}, Test Input Char Acc: {avg_c1_s1_input_cer:0.5f}, Test Gen Char Acc: {avg_c2_s1_gen_cer:0.5f}'

        print(loss_log)
        log.write(loss_log+"\n")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_dir', default='/checkpoint/pkrishnan/experiments/scribe/Exp18/', help='Where to store logs and models')
    parser.add_argument('--exp_name', default='debug', help='Where to store logs and models')
    parser.add_argument('--exp_iter', default='0', help='Where to store logs and models')
    parser.add_argument('--test_data', required=True, help='path to validation dataset')
    parser.add_argument('--words_file', required=True, default='', help="path to words file. phoc will be sampled from this file")
    parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
    parser.add_argument('--batch_size', type=int, default=32, help='input batch size')
    parser.add_argument('--saved_ocr_model', default='', help="path to model to continue training")
    parser.add_argument('--saved_synth_model', default='', help="path to model to continue training")

    """ Data processing """
    parser.add_argument('--select_data', type=str, default='ST',
                        help='select training data (default is MJ-ST, which means MJ and ST used as training data)')
    parser.add_argument('--batch_ratio', type=str, default='1.0',
                        help='assign ratio for each selected data in the batch')
    parser.add_argument('--total_data_usage_ratio', type=str, default='1.0',
                        help='total data usage ratio, this ratio is multiplied to total number of data.')
    parser.add_argument('--batch_max_length', type=int, default=24, help='maximum-label-length')
    parser.add_argument('--batch_min_length', type=int, default=1, help='minimum-label-length')
    parser.add_argument('--fixedString', action='store_true', help='use fixed length data')
    parser.add_argument('--batch_exact_length', type=int, default=4, help='exact-label-length')
    parser.add_argument('--imgH', type=int, default=64, help='the height of the input image')
    parser.add_argument('--imgW', type=int, default=256, help='the width of the input image')
    parser.add_argument('--ocr_imgH', type=int, default=32, help='the height of the input image')
    parser.add_argument('--ocr_imgW', type=int, default=100, help='the width of the input image')
    parser.add_argument('--imgH_filt', type=int, default=48, help='the height of the input image')
    parser.add_argument('--rgb', action='store_true', help='use rgb input')
    parser.add_argument('--character', type=str,
                        default='0123456789abcdefghijklmnopqrstuvwxyz', help='character label')
    parser.add_argument('--sensitive', action='store_true', help='for sensitive character mode')
    parser.add_argument('--PAD', action='store_true', help='whether to keep ratio then pad for image resize')
    parser.add_argument('--data_filtering_off', action='store_true', help='for data_filtering_off mode')
    parser.add_argument('--pairText', action='store_true', help='use additional text for generation')
    parser.add_argument('--lexFile', default='/checkpoint/pkrishnan/datasets/vocab/english-words.txt', help='unqiue words in language')

    """ Model Architecture """
    parser.add_argument('--Transformation', type=str,  help='Transformation stage. None|TPS')
    parser.add_argument('--FeatureExtraction', type=str,
                        help='FeatureExtraction stage. VGG|RCNN|ResNet')
    parser.add_argument('--SequenceModeling', type=str, help='SequenceModeling stage. None|BiLSTM')
    parser.add_argument('--Prediction', type=str, help='Prediction stage. CTC|Attn')
    parser.add_argument('--num_fiducial', type=int, default=20, help='number of fiducial points of TPS-STN')
    parser.add_argument('--input_channel', type=int, default=1,
                        help='the number of input channel of Feature extractor')
    parser.add_argument('--ocr_input_channel', type=int, default=1,
                        help='the number of input channel of Feature extractor')
    parser.add_argument('--output_channel', type=int, default=512,
                        help='the number of output channel of Feature extractor')
    parser.add_argument('--hidden_size', type=int, default=256, help='the size of the LSTM hidden state')
    parser.add_argument('--char_embed_size', type=int, default=60, help='character embedding for content encoder')
    parser.add_argument('--vggNoMean', action='store_true', help='if yes; No mean normalization is applied to vgg input')
    parser.add_argument('--styleDetach', action='store_true', help='whether to detach style')
    parser.add_argument('--gan_type', default='lsgan', help='lsgan/nsgan/wgan')
    parser.add_argument('--style_input', action='store_true', help='whether target style input is given for training/validation')
    parser.add_argument('--style_content_input', action='store_true', help='whether target  input content image is given for training/validation')
    parser.add_argument('--styleNorm', default='bn', help='bn/in')
    parser.add_argument('--contentNorm', default='bn', help='bn/in')
    parser.add_argument('--sizeFilt', action='store_true', help='for applying size based image filtering')

    parser.add_argument('--debugFlag', action='store_true', help='for debugging')
    parser.add_argument('--visFlag', action='store_true', help='for visualization')

    parser.add_argument("--size", type=int, default=64)
    parser.add_argument("--mixing", type=float, default=0.9)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--latent", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--input_latent", action="store_true")
    parser.add_argument("--content_inject_index", type=int, default=1)
    parser.add_argument("--zAlone", action="store_true", help="test original style GAN")
    parser.add_argument("--noiseConcat", action="store_true", help="noise concat with style")

    parser.add_argument('--cEncode', default='mlp', help='mlp/cnn')
    parser.add_argument('--taskActivation', default=None, help='None/tanh')
    parser.add_argument("--numLatents", type=int, default=1)
    parser.add_argument("--realTrData", action="store_true", help="flag for training real data where we don't have target style GT")
    parser.add_argument("--realVaData", action="store_true", help="flag for validation/testing real data where we don't have target style GT")

    parser.add_argument('--outPairFile', default="", help='unqiue words in language')

    opt = parser.parse_args()

    if opt.rgb:
        opt.input_channel = 3

    if opt.taskActivation == 'None':
        opt.taskActivation = None

    if not opt.exp_name:
        opt.exp_name = f'{opt.Transformation}-{opt.FeatureExtraction}-{opt.SequenceModeling}-{opt.Prediction}'
        opt.exp_name += f'-Seed{opt.manualSeed}'
        # print(opt.exp_name)

    os.makedirs(os.path.join(opt.exp_dir,opt.exp_name), exist_ok=True)
    opt.valDir=os.path.join(opt.exp_dir,opt.exp_name,'testImages')

    os.makedirs(opt.valDir, exist_ok=True)

    """ vocab / character number configuration """
    if opt.sensitive:
        # opt.character += 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        opt.character = string.printable[:-6]  # same with ASTER setting (use 94 char).

    cudnn.benchmark = True
    cudnn.deterministic = True

    test(opt)


from .modules.backbone import csp_darknet53
from .modules.neck import neck_yolov4
from .modules.head import yolo_head

import tensorflow as tf
import numpy as np

class YOLO:
    """
        YOLOv4 Implementation Model.

        Takes de input shape for the constructor, and
        creates an instance of the YOLOV4 net.

        @version: 0.1v
    """
    def __init__(self, input_shape):
        anchors = np.random.rand(3, 2)

        # Define the Model phases
        self.backbone = csp_darknet53(input_shape, name = "CSPDarknet53")
        self.neck = neck_yolov4(self.backbone.outputs, name = "SPP/PAN")
        self.head = yolo_head(self.neck.outputs, anchors = anchors, num_classes = 8) # , name = "Detector") #, anchors = anchors, num_classes = 8, name = "Detector")

        # Declare Inputs
        self.inputs = tf.keras.layers.Input(input_shape, name = "Inputs")

          # Process each phase
        low = self.backbone(self.inputs)
        print("low -> mid", end=' ->')
        mid = self.neck(low)
        print(" high")
        high = self.head(mid)

        self.Model = tf.keras.Model(inputs = self.inputs, outputs = high)

    def predict(self, input):
        return self.Model.predict(input)

    def show(self):
        self.Model.summary()


import math
import sagemaker as sage
import os
import requests
from sagemaker import get_execution_role
from sagemaker import transformer
from sagemaker import model
from src.model_package_arns import ModelPackageArnProvider
from sagemaker import ModelPackage
from io import BytesIO

import boto3
import json
import ast

def compare_price(actual_price, third_party_price, json_format):
    if actual_price > third_party_price:
        percent_diff = third_party_price/actual_price

        if percent_diff < 0.7:
            if json_format["is_used_product"] == True:
                json_body = {"fake_product_message": "We compared your selling price to other retailers in the market and it seems to be significantly lower. Please update the description and the title. In your description you must include the appropriate reason to sell at a lower price."
                }
                return json_body
            elif json_format["is_used_product"] == False:
                json_body = {"fake_product_message": "We compared your selling price to other retailers in the market and it seems to be significantly lower. Is this a used product? If so please mark this product as used/refurbished. Please update the description and the title. In your description you must include the appropriate reason to sell at a lower price."}
            else:
                json_body = {"fake_product_message": "We compared your selling price to other retailers in the market and it seems to be significantly lower. Please update the description and the title. In your description you must include the appropriate reason to sell at a lower price."}

            return json_body

        else:
            return {"fake_product_message": "All Clear"}
    return {"fake_product_message": "All Clear"}

def truth_detector(text, text2):
    comprehend = boto3.client(service_name='comprehend',
                              region_name='us-east-1')
    product_val = json.dumps(
        comprehend.detect_key_phrases(Text=text, LanguageCode='en'),
        sort_keys=True, indent=8)
    product_values = ast.literal_eval(product_val)
    product_phrases = []
    for i in product_values['KeyPhrases']:
        product_phrases.append(i["Text"].lower())

    print(product_phrases)

    customer_val = json.dumps(
        comprehend.detect_key_phrases(Text=text2, LanguageCode='en'),
        sort_keys=True, indent=8)
    customer_values = ast.literal_eval(customer_val)
    customer_phrases = []
    for i in customer_values['KeyPhrases']:
        customer_phrases.append(i["Text"].lower())

    intersection = list(set(product_phrases) & set(customer_phrases))
    percent = len(intersection) / len(customer_phrases)

    if percent > 0.7:
        return 0
    else:
        return 1


# encoding: UTF-8
#!/usr/bin/env python2
import cv2 as cv
import os
import numpy as np
import time
import rospy
from visualization_msgs.msg import Marker
from moving_utils import Movement

# y轴偏移量
pump_y = -55
# x轴偏移量
pump_x = 15

class Detect_marker(Movement):
    def __init__(self):
        super(Detect_marker, self).__init__()
        # set cache of real coord
        self.cache_x = self.cache_y = 0

        # which robot
        self.robot_m5 = os.popen("ls /dev/ttyUSB*").readline()[:-1]
        self.robot_wio = os.popen("ls /dev/ttyACM*").readline()[:-1]
        self.robot_raspi = os.popen("ls /dev/ttyAMA*").readline()[:-1]
        self.robot_jes = os.popen("ls /dev/ttyTHS1").readline()[:-1]
        self.raspi = False
        if "dev" in self.robot_m5:
            self.Pin = [2, 5]
        elif "dev" in self.robot_wio:
            self.Pin = [20, 21]
            for i in self.move_coords:
                i[2] -= 20
        elif "dev" in self.robot_raspi or "dev" in self.robot_jes:
            import RPi.GPIO as GPIO
            self.GPIO = GPIO
            GPIO.setwarnings(False)
            GPIO.setmode(GPIO.BCM)
            GPIO.setup(20, GPIO.OUT)
            GPIO.setup(21, GPIO.OUT)
            self.raspi = True
        if self.raspi:
            self.gpio_status(False)
        else:
            self.pub_pump(False, self.Pin)
        # Creating a Camera Object
        cap_num = 0
        self.cap = cv.VideoCapture(cap_num, cv.CAP_V4L)
        # Get ArUco marker dict that can be detected.
        self.aruco_dict = cv.aruco.Dictionary_get(cv.aruco.DICT_6X6_250)
        # Get ArUco marker params.
        self.aruco_params = cv.aruco.DetectorParameters_create()
        self.calibrationParams = cv.FileStorage(
            "calibrationFileName.xml", cv.FILE_STORAGE_READ)
        # Get distance coefficient.
        self.dist_coeffs = self.calibrationParams.getNode("distCoeffs").mat()

        height = self.cap.get(4)
        focal_length = width = self.cap.get(3)
        center = [width / 2, height / 2]
        # Calculate the camera matrix.
        self.camera_matrix = np.array(
            [
                [focal_length, 0, center[0]],
                [0, focal_length, center[1]],
                [0, 0, 1],
            ],
            dtype=np.float32,
        )
        # init a node and a publisher
        rospy.init_node("encode_marker", anonymous=True)
        self.pub = rospy.Publisher('/cube', Marker, queue_size=1)

        self.marker = Marker()
        self.marker.header.frame_id = "/joint1"
        self.marker.ns = "cube"
        self.marker.type = self.marker.CUBE
        self.marker.action = self.marker.ADD
        self.marker.scale.x = 0.04
        self.marker.scale.y = 0.04
        self.marker.scale.z = 0.04
        self.marker.color.a = 1
        self.marker.color.r = 0.3
        self.marker.color.g = 0.3
        self.marker.color.b = 0.3

        # marker position initial
        self.marker.pose.position.x = 0
        self.marker.pose.position.y = 0
        self.marker.pose.position.z = 0.03
        self.marker.pose.orientation.x = 0
        self.marker.pose.orientation.y = 0
        self.marker.pose.orientation.z = 0
        self.marker.pose.orientation.w = 1.0

    # Grasping motion
    def move(self, x, y):
        if self.raspi:
            coords = [
                [145.6, -64.9, 285.2, 179.88, 7.67, 179],
                [130.1, -155.6, 243.9, 178.99, 5.38, -179.9]
            ]
        else:
            coords = [
                [135.0, -65.5, 280.1, 178.99, 5.38, -179.9],
                [136.1, -141.6, 243.9, 178.99, 5.38, -179.9]
            ]

        # publish marker
        self.marker.header.stamp = rospy.Time.now()
        self.marker.pose.position.x = (coords[0][0]-x)/1000.0
        self.marker.pose.position.y = (coords[0][1]-y)/1000.0
        self.pub.publish(self.marker)

        # send coordinates to move mycobot
        self.pub_coords(coords[0], 30, 1)
        time.sleep(2)
        self.pub_coords([coords[0][0]-x, coords[0][1]-y,
                        240, 178.99, 5.38, -179.9], 25, 1)
        time.sleep(2)
        self.pub_coords([coords[0][0]-x, coords[0][1]-y,
                        200, 178.99, 5.38, -179.9], 25, 1)
        time.sleep(2)
        if "dev" in self.robot_m5 or self.raspi:
            self.pub_coords([coords[0][0]-x, coords[0][1]-y,
                            90, 178.99, 5.38, -179.9], 25, 1)
        elif "dev" in self.robot_wio:
            self.pub_coords([coords[0][0]-x+20, coords[0][1] -
                            y-10, 70, 178.99, 5.38, -179.9], 25, 1)
        time.sleep(2)
        if self.raspi:
            self.gpio_status(True)
        else:
            self.pub_pump(True, self.Pin)
        time.sleep(1)
        self.pub_coords(coords[0], 30, 1)
        time.sleep(3)
        self.pub_coords(coords[1], 30, 1)
        time.sleep(2)
        if self.raspi:
            self.gpio_status(False)
        else:
            self.pub_pump(False, self.Pin)
        # publish marker
        time.sleep(1)
        self.marker.header.stamp = rospy.Time.now()
        self.marker.pose.position.x = coords[1][0]/1000.0
        self.marker.pose.position.y = coords[1][1]/1000.0
        self.pub.publish(self.marker)

        self.pub_coords(coords[0], 30, 1)
        time.sleep(2)

    def gpio_status(self, flag):
        if flag:
            self.GPIO.output(20, 0)
            self.GPIO.output(21, 0)
        else:
            self.GPIO.output(20, 1)
            self.GPIO.output(21, 1)

    # decide whether grab cube
    def decide_move(self, x, y):

        print(x, y)
        # detect the cube status move or run
        if (abs(x - self.cache_x) + abs(y - self.cache_y)) / 2 > 5:  # mm
            self.cache_x, self.cache_y = x, y
            return
        else:
            self.cache_x = self.cache_y = 0
            if "dev" in self.robot_jes:
                if x > -20:
                    y += 10
                if y > -25:
                    x -= 5
                x += 10

            self.move(x, y)

    # init mycobot
    def init_mycobot(self):

        for _ in range(5):
            print _
            self.pub_coords([145.6, -64.9, 285.2, 179.88, 7.67, 179], 20, 1)
            time.sleep(0.5)

    def run(self):
        global pump_y, pump_x
        self.init_mycobot()
        num = sum_x = sum_y = 0
        while cv.waitKey(1) < 0:
            success, img = self.cap.read()
            if not success:
                print("It seems that the image cannot be acquired correctly.")
                break

            # transfrom the img to model of gray
            gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            # Detect ArUco marker.
            corners, ids, rejectImaPoint = cv.aruco.detectMarkers(
                gray, self.aruco_dict, parameters=self.aruco_params
            )

            if len(corners) > 0:
                if ids is not None:
                    # get informations of aruco
                    ret = cv.aruco.estimatePoseSingleMarkers(
                        corners, 0.03, self.camera_matrix, self.dist_coeffs
                    )
                    # rvec:rotation offset,tvec:translation deviator
                    (rvec, tvec) = (ret[0], ret[1])
                    (rvec - tvec).any()
                    xyz = tvec[0, 0, :]
                    # calculate the coordinates of the aruco relative to the pump
                    xyz = [round(xyz[0]*1000+pump_y, 2), round(xyz[1]
                                                               * 1000+pump_x, 2), round(xyz[2]*1000, 2)]

                    for i in range(rvec.shape[0]):
                        # draw the aruco on img
                        cv.aruco.drawDetectedMarkers(img, corners)
                        cv.aruco.drawAxis(
                            img,
                            self.camera_matrix,
                            self.dist_coeffs,
                            rvec[i, :, :],
                            tvec[i, :, :],
                            0.03,
                        )

                        if num < 40:
                            if self.raspi:
                                sum_x -= 30
                            sum_x += xyz[1]
                            sum_y += xyz[0]
                            num += 1
                        elif num == 40:
                            self.decide_move(sum_x/40.0, sum_y/40.0)
                            num = sum_x = sum_y = 0

            cv.imshow("encode_image", img)

if __name__ == "__main__":
    detect = Detect_marker()
    detect.run()


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from model.latticelstm import LatticeLSTM

class BiLSTM(nn.Module):
    def __init__(self, data):
        super(BiLSTM, self).__init__()
        print("build batched bilstm...")
        self.use_bichar = data.use_bichar
        self.gpu = data.HP_gpu
        # self.use_char = data.HP_use_character
        self.use_gaz = data.HP_use_gaz
        self.batch_size = data.HP_batch_size
        self.char_hidden_dim = 0
        self.embedding_dim = data.char_emb_dim
        self.hidden_dim = data.HP_hidden_dim
        self.drop = nn.Dropout(data.HP_dropout)
        self.droplstm = nn.Dropout(data.HP_dropout)
        self.char_embeddings = nn.Embedding(data.char_alphabet.size(), self.embedding_dim)
        self.bichar_embeddings = nn.Embedding(data.bichar_alphabet.size(), data.bichar_emb_dim)
        self.bilstm_flag = data.HP_bilstm
        self.lstm_layer = data.HP_lstm_layer
        if data.pretrain_char_embedding is not None:
            self.char_embeddings.weight.data.copy_(torch.from_numpy(data.pretrain_char_embedding))
        else:
            self.char_embeddings.weight.data.copy_(
                # torch.from_numpy(_): numpy to  torch
                # torch_data.numpy(): torch to numpy
                torch.from_numpy(self.random_embedding(data.char_alphabet.size(), self.embedding_dim)))

        if data.pretrain_bichar_embedding is not None:
            self.bichar_embeddings.weight.data.copy_(torch.from_numpy(data.pretrain_bichar_embedding))
        else:
            self.bichar_embeddings.weight.data.copy_(
                torch.from_numpy(self.random_embedding(data.bichar_alphabet.size(), data.bichar_emb_dim)))
        # The LSTM takes word embeddings as inputs, and outputs hidden states with dimensionality hidden_dim.
        lstm_hidden = data.HP_hidden_dim // 2 if self.bilstm_flag else data.HP_hidden_dim
        lstm_input = self.embedding_dim + self.char_hidden_dim
        if self.use_bichar:
            lstm_input += data.bichar_emb_dim
        self.forward_lstm = LatticeLSTM(lstm_input, lstm_hidden, data.gaz_dropout, data.gaz_alphabet.size(),
                                        data.gaz_emb_dim, data.pretrain_gaz_embedding, True, data.HP_fix_gaz_emb,
                                        self.gpu)
        if self.bilstm_flag:
            self.backward_lstm = LatticeLSTM(lstm_input, lstm_hidden, data.gaz_dropout, data.gaz_alphabet.size(),
                                             data.gaz_emb_dim, data.pretrain_gaz_embedding, False, data.HP_fix_gaz_emb,
                                             self.gpu)
        # use biLSTM
        # self.lstm = nn.LSTM(lstm_input, lstm_hidden, num_layers=self.lstm_layer, batch_first=True, bidirectional=self.bilstm_flag)
        # The linear layer that maps from hidden state space to tag space
        self.hidden2tag = nn.Linear(data.HP_hidden_dim, data.label_alphabet_size)

        if self.gpu:
            self.drop = self.drop.cuda()
            self.droplstm = self.droplstm.cuda()
            self.char_embeddings = self.char_embeddings.cuda()
            self.bichar_embeddings = self.bichar_embeddings.cuda()
            self.forward_lstm = self.forward_lstm.cuda()
            if self.bilstm_flag:
                self.backward_lstm = self.backward_lstm.cuda()
            self.hidden2tag = self.hidden2tag.cuda()

    def random_embedding(self, vocab_size, embedding_dim):
        pretrain_emb = np.empty([vocab_size, embedding_dim])
        scale = np.sqrt(3.0 / embedding_dim)
        for index in range(vocab_size):
            pretrain_emb[index, :] = np.random.uniform(-scale, scale, [1, embedding_dim])
        return pretrain_emb

    def get_lstm_features(self, gaz_list, char_inputs, bichar_inputs, char_seq_lengths):
        """
            input:
                char_inputs: (batch_size, sent_len)
                gaz_list:
                char_seq_lengths: list of batch_size, (batch_size,1)
                character_inputs: (batch_size*sent_len, word_length)
                character_seq_lengths: list of whole batch_size for char, (batch_size*sent_len, 1)
                char_seq_recover: variable which records the char order information, used to recover char order
            output:
                Variable(sent_len, batch_size, hidden_dim)
        """
        char_embs = self.char_embeddings(char_inputs)
        if self.use_bichar:
            bichar_embs = self.bichar_embeddings(bichar_inputs)
            char_embs = torch.cat([char_embs, bichar_embs], 2)

        char_embs = self.drop(char_embs)

        hidden = None
        lstm_out, hidden = self.forward_lstm(char_embs, gaz_list, hidden)
        if self.bilstm_flag:
            backward_hidden = None
            backward_lstm_out, backward_hidden = self.backward_lstm(char_embs, gaz_list, backward_hidden)
            lstm_out = torch.cat([lstm_out, backward_lstm_out], 2)
        lstm_out = self.droplstm(lstm_out)
        return lstm_out

    def get_output_score(self, gaz_list, char_inputs, bichar_inputs, char_seq_lengths):
        lstm_out = self.get_lstm_features(gaz_list, char_inputs, bichar_inputs, char_seq_lengths)
        outputs = self.hidden2tag(lstm_out)
        return outputs

    def neg_log_likelihood_loss(self, gaz_list, char_inputs, bichar_inputs, char_seq_lengths, batch_label, mask):
        ## mask is not used
        batch_size = char_inputs.size(0)
        seq_len = char_inputs.size(1)
        total_word = batch_size * seq_len
        loss_function = nn.NLLLoss(ignore_index=0, size_average=False)
        outs = self.get_output_score(gaz_list, char_inputs, bichar_inputs, char_seq_lengths)
        outs = outs.view(total_word, -1)
        score = F.log_softmax(outs, 1)
        loss = loss_function(score, batch_label.view(total_word))
        _, tag_seq = torch.max(score, 1)
        tag_seq = tag_seq.view(batch_size, seq_len)
        return loss, tag_seq

    def forward(self, gaz_list, char_inputs, bichar_inputs, char_seq_lengths, mask):

        batch_size = char_inputs.size(0)
        seq_len = char_inputs.size(1)
        total_word = batch_size * seq_len
        outs = self.get_output_score(gaz_list, char_inputs, bichar_inputs, char_seq_lengths)
        outs = outs.view(total_word, -1)
        _, tag_seq = torch.max(outs, 1)
        tag_seq = tag_seq.view(batch_size, seq_len)
        # filter padded position with zero
        decode_seq = mask.long() * tag_seq
        return decode_seq


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Jun  3 15:36:49 2017

@author: adam
"""

import cv2
import numpy as np
from keras.models import load_model

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.1
set_session(tf.Session(config=config))

faceCascade = cv2.CascadeClassifier('haarcascade_frontalface_alt2.xml')

video_capture = cv2.VideoCapture(0)
model = load_model('keras_model/6_emo_50_epo.h5')
model.get_config()

target = ['Angry','Fear','Happy','Sad','Surprise','Neutral']
font = cv2.FONT_HERSHEY_SIMPLEX
while True:
    # Capture frame-by-frame
    ret, frame = video_capture.read()

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    faces = faceCascade.detectMultiScale(gray, scaleFactor=1.1)

    # Draw a rectangle around the faces
    for (x, y, w, h) in faces:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2, 5)
        face_crop = frame[y:y + h, x:x + w]
        face_crop = cv2.resize(face_crop, (48, 48))
        face_crop = cv2.cvtColor(face_crop, cv2.COLOR_BGR2GRAY)
        face_crop = face_crop.astype('float32') / 255
        face_crop = np.asarray(face_crop)
        face_crop = face_crop.reshape(1, face_crop.shape[0], face_crop.shape[1],1)
        result = target[np.argmax(model.predict(face_crop))]
        cv2.putText(frame, result, (x, y), font, 1, (200, 150, 100), 1, cv2.LINE_AA)

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

    if cv2.waitKey(1) & 0xFF == ord('q'):     # to quit the window press q
        break

# When everything is done, release the capture
video_capture.release()
cv2.destroyAllWindows()


"""
SCRIPT TO WRITE GROMACS GRO AND ITP FILES
FROM BOSS ZMATRIX
Created on Mon Feb 15 15:40:05 2016
@author: Leela S. Dodda leela.dodda@yale.edu
@author: William L. Jorgensen Lab

Usage: python OPM_Routines.py -z phenol.z -r PHN
REQUIREMENTS:
BOSS (need to set BOSSdir in bashrc and cshrc)
Preferably Anaconda python with following modules
pandas
argparse
numpy
"""

from LigParGen.BOSSReader import ucomb,pairing_func,bossPdbAtom2Element,bossElement2Mass
from LigParGen.BOSSReader import Refine_PDB_file,get_coos_from_pdb
import pickle
import pandas as pd
import numpy as np

def GMX_pairs(tors, ang, bond):
    atom_list = np.array(list(bond.cl1) + list(bond.cl2))
    atom_list = np.sort(atom_list)
    atom_list = np.unique(atom_list)
    dict_bond = {ano: list(bond[bond.cl1 == ano]['cl2']) +
                 list(bond[bond.cl2 == ano]['cl1']) for ano in atom_list}
    dict_ang = {}
    for ano in atom_list:
        tr_1 = []
        for bno in dict_bond[ano]:
            tr_1 += dict_bond[bno]
        tr_1 = list(set(tr_1))
        tr_1.remove(ano)
        dict_ang[ano] = tr_1
    NP_dat = []
    for a in dict_bond.keys():
        for b in dict_bond[a]:
            NP_dat.append([a, b, 1, pairing_func(a + 1, b + 1)
                           [0], pairing_func(a + 1, b + 1)[1]])
            for c in dict_ang[b]:
                NP_dat.append([a, c, 3, pairing_func(a + 1, c + 1)
                               [0], pairing_func(a + 1, c + 1)[1]])
    for a in dict_ang.keys():
        for b in dict_ang[a]:
            NP_dat.append([a, b, 2, pairing_func(a + 1, b + 1)
                           [0], pairing_func(a + 1, b + 1)[1]])
            for c in dict_bond[b]:
                NP_dat.append([a, c, 3, pairing_func(a + 1, c + 1)
                               [0], pairing_func(a + 1, c + 1)[1]])
    NP_df = pd.DataFrame(NP_dat, columns=['I', 'J', 'BSEP', 'UNQ', 'TY'])
    NP_df = NP_df[NP_df.I != NP_df.J]  # CASES WITH 3 MEMBERED RINGS
    NP_df = NP_df.sort_values(['UNQ'])
    NP_B = NP_df[NP_df.BSEP == 1]
    NP_A = NP_df[NP_df.BSEP == 2]
    NP_T = NP_df[NP_df.BSEP == 3]
    NP_T = NP_T.drop_duplicates(['UNQ'])
    NP_T = NP_T[~ NP_T.UNQ.isin(NP_B.UNQ)]
    NP_T = NP_T[~ NP_T.UNQ.isin(NP_A.UNQ)]
    return list(NP_T.TY)

def gmxDihed(df):
    '''
        PRINTING RYCKET-BELLMAN STYLE PARAMETERS BECAUSE FEP DOES NOT SUPPORT PROPER DIHEDRALS
    '''
    [f1, f2, f3, f4] = [df.V1, df.V2, df.V3, df.V4]
    cdat = [f2 + (f1 + f3) * 0.5, 1.5 * f3 - 0.5 * f1, 4.0 *
            f4 - f2, -2.0 * f3, -4.0 * f4, 0.0, 0.00]
    return ('%5s%5s%5s%5s        3      %6.3f  %6.3f  %6.3f  %6.3f  %6.3f  %6.3f\n' % (
        df['I'] + 1, df['J'] + 1, df['K'] + 1, df['L'] + 1, cdat[0], cdat[1], cdat[2], cdat[3], cdat[4], cdat[5]))

def gmxImp(df):
    odihed = []
    for pot in range(1, 5):
        if df['V' + str(pot)] != 0.00:
            odihed.append('%6d%6d%6d%6d    4     %10.3f %10.3f %5d  \n' % (
                df['I'] + 1, df['J'] + 1, df['K'] +
                1, df['L'] + 1, 180.00 * abs(pot % 2 - 1),
                float(df['V' + str(pot)]) * 0.5, pot))
    return (odihed)

def boss2opmTorsion(bnd_df, num2opls, st_no, molecule_data, itpf):
    dhd = []
    for line in molecule_data.MolData['TORSIONS']:
        dt = [float(f) for f in line]
        dhd.append(dt)
    dhd = np.array(dhd)
    dhd = dhd * 4.184  # kcal to kj conversion
    dhd = dhd  # Komm = Vopls/2
    dhd_df = pd.DataFrame(dhd, columns=['V1', 'V2', 'V3', 'V4'])
    ats = []
    for line in molecule_data.MolData['ATOMS'][3:]:
        dt = [line.split()[0], line.split()[4],
              line.split()[6], line.split()[8]]
        dt = [int(d) for d in dt]
        ats.append(dt)
    for line in molecule_data.MolData['ADD_DIHED']:
        dt = [int(l) for l in line]
        ats.append(dt)
    assert len(ats) == len(
        dhd), 'Number of Dihedral angles in Zmatrix and Out file dont match'
    ats = np.array(ats) - st_no
    for i in range(len(ats)):
        for j in range(len(ats[0])):
            if ats[i][j] < 0:
                ats[i][j] = 0
    at_df = pd.DataFrame(ats, columns=['I', 'J', 'K', 'L'])
    final_df = pd.concat([dhd_df, at_df], axis=1, join_axes=[at_df.index])
    bndlist = list(bnd_df.UR) + (list(bnd_df.UR))
    final_df['TY'] = ['Proper' if ucomb(list([final_df.I[n], final_df.J[n], final_df.K[
        n], final_df.L[n]]), bndlist) == 3 else 'Improper' for n in range(len(final_df.I))]
    final_df['SumV'] = np.abs(
        final_df.V1) + np.abs(final_df.V2) + np.abs(final_df.V3) + np.abs(final_df.V4)
    final_df['TI'] = [num2opls[j] for j in final_df.I]
    final_df['TJ'] = [num2opls[j] for j in final_df.J]
    final_df['TK'] = [num2opls[j] for j in final_df.K]
    final_df['TL'] = [num2opls[j] for j in final_df.L]
    if len(final_df.index) > 0:
        final_df['NAME'] = final_df.TI + '-' + final_df.TJ + \
            '-' + final_df.TK + '-' + final_df.TL
        final_df = final_df.sort_values(['NAME'])
        tor_bos = final_df.drop(
            ['I', 'J', 'K', 'L', 'TI', 'TJ', 'TK', 'TL'], 1)
        tor_bos = tor_bos.drop_duplicates()
        df = final_df.ix[tor_bos.index]
        return final_df, df
    else:
        return final_df, final_df

def boss2gmxBond(molecule_data, st_no, itpf):
    bdat = molecule_data.MolData['BONDS']
    bdat['cl1'] = [x - st_no if not x - st_no < 0 else 0 for x in bdat['cl1']]
    bdat['cl2'] = [x - st_no if not x - st_no < 0 else 0 for x in bdat['cl2']]
    bnd_df = pd.DataFrame(bdat)
    bnd_df['KIJ'] = bnd_df['KIJ'] * 836.80
    bnd_df['RIJ'] = bnd_df['RIJ'] * 0.10
    bnd_df['UNQ'] = [pairing_func(i + 1, j + 1)[0]
                     for i, j in zip(bnd_df.cl1, bnd_df.cl2)]
    bnd_df['UNQ'] = [pairing_func(i, j)[0]
                     for i, j in zip(bnd_df.cl1, bnd_df.cl2)]
    bnd_df['UF'] = ((bnd_df.cl1 + bnd_df.cl2) *
                    (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5) + bnd_df.cl2
    bnd_df['UR'] = ((bnd_df.cl1 + bnd_df.cl2) *
                    (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5) + bnd_df.cl1
    connects = []
    for (ai, aj) in zip(bnd_df.cl1, bnd_df.cl2):
        connects.append("{:>5}{:>5}    1".format(aj + 1, ai + 1))
    full_bnd = bnd_df.copy()
    bnd_df = bnd_df.drop_duplicates(['TIJ'])
    return full_bnd, connects

def boss2gmxAngle(anglefile, num2opls, st_no, itpf):
    adat = anglefile
    adat['cl1'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl1']]
    adat['cl2'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl2']]
    adat['cl3'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl3']]
    ang_df = pd.DataFrame(adat)
    ang_df = ang_df[ang_df.K > 0]
    ang_df['K'] = 8.3680 * ang_df['K']
    ang_df['R'] = ang_df['R']
    full_ang = ang_df.copy()
    ang_df['TY'] = np.array([num2opls[i] + '-' + num2opls[j] + '-' + num2opls[k]
                             for i, j, k in zip(ang_df.cl1, ang_df.cl2, ang_df.cl3)])
    ang_df = ang_df.drop_duplicates(['TY'])
    return full_ang

def bossData(molecule_data):
    ats_file = molecule_data.MolData['ATOMS']
    types = []
    for i in enumerate(ats_file):
        types.append([i[1].split()[1], 'opls_' + i[1].split()[2]])
    st_no = 3
    Qs = molecule_data.MolData['Q_LJ']
    assert len(Qs) == len(types), 'Please check the at_info and Q_LJ_dat files'
    num2opls = {}
    for i in range(0, len(types)):
        num2opls[i] = Qs[i][0]
    num2typ2symb = {i: types[i] for i in range(len(Qs))}
    for i in range(len(Qs)):
        num2typ2symb[i].append(bossPdbAtom2Element(
            num2typ2symb[i][0]) + num2typ2symb[i][1][-3:])
        num2typ2symb[i].append(bossPdbAtom2Element(num2typ2symb[i][0]))
        num2typ2symb[i].append(bossElement2Mass(num2typ2symb[i][3]))
        num2typ2symb[i].append(Qs[i][0])
    return (types, Qs, num2opls, st_no, num2typ2symb)

def pdb2gro(atoms, coos, resid):
    grof = open(resid + '.gro', 'w+')
    grof.write('LIGPARGEN GENERATED GRO FILE\n')
    grof.write('%5d\n' % len(atoms))
    num = 0
    for (i, j) in zip(atoms, coos):
        num += 1
        grof.write('%5d%-5s%5s%5d%8.3f%8.3f%8.3f\n' %
                   (1, resid, i, num, j[0] * 0.1, j[1] * 0.1, j[2] * 0.1))
    grof.write('%10.5f%10.5f%10.5f\n' % (1.00000, 1.00000, 1.00000))
    grof.close()
    return None

def boss2gmxAtom(resid, num2typ2symb, Qs, itpf):
    itpf.write('[ atoms ]\n')
    itpf.write(
        ';   nr       type  resnr residue  atom   cgnr     charge       mass  \n')
    for i in range(len(Qs)):
        itpf.write(' %5d %10s %6d %6s %5s %6d %10s %10.4f \n' % (
            i + 1, num2typ2symb[i][1], 1, resid, num2typ2symb[i][0], ((i + 1) / 33.0) + 1, Qs[i][1], num2typ2symb[i][4]))
    return None

def boss2gmx(resid, molecule_data, pdb_file):
    types, Qs, num2opls, st_no, num2typ2symb = bossData(
        molecule_data)
    itpf = open(resid + '.itp', 'w+')
    itpf.write("""
; Writted by Leela S. Dodda leela.dodda@yale.edu
; GENERATED BY Server
; Jorgensen Lab @ Yale University
;
""")
    itpf.write('[ atomtypes ]\n')
    atty = ['%10s %5s %10.4f     0.000    A    %10.5E   %10.5E\n' % (num2typ2symb[i][
        1], num2typ2symb[i][2], num2typ2symb[i][4],
        float(Qs[i][2]) * 0.1, float(Qs[i][3]) * 4.184) for
        i in range(len(Qs))]
    atty = list(set(atty))
    for att in atty:
        itpf.write('%s' % att)
    bnd_df, connects = boss2gmxBond(molecule_data, st_no, itpf)
    full_ang = boss2gmxAngle(molecule_data.MolData[
        'ANGLES'], num2opls, st_no, itpf)
    itpf.write('[ moleculetype ]\n')
    itpf.write('; Name               nrexcl\n')
    itpf.write('%s                   3\n' % resid)
    boss2gmxAtom(resid, num2typ2symb, Qs, itpf)
    itpf.write('[ bonds ]\n')
    for brow in bnd_df.iterrows():
        index, bnd = brow
        itpf.write('%5d %5d %5d %11.4f %10.3f\n' %
                   (bnd['cl1'] + 1, bnd['cl2'] + 1, 1, bnd['RIJ'], bnd['KIJ']))
    itpf.write('\n[ angles ]\n')
    itpf.write(
        ';  ai    aj    ak funct            c0            c1            c2            c3 \n')
    for row in full_ang.iterrows():
        index, angs = row
        itpf.write('%5d %5d %5d %5d %10.3f %10.3f\n' %
                   (angs['cl1'] + 1, angs['cl2'] + 1, angs['cl3'] + 1, 1, angs['R'], angs['K']))
    full_tor, tor_df = boss2opmTorsion(
        bnd_df, num2opls, st_no, molecule_data, itpf)
    if len(tor_df.index) != len(full_tor.index):
        itpf.write('\n[ dihedrals ]\n')
        itpf.write('; IMPROPER DIHEDRAL ANGLES \n')
        itpf.write(
            ';  ai    aj    ak    al funct            c0            c1            c2            c3            c4            c5\n')
        for row in full_tor[full_tor.TY == 'Improper'].iterrows():
            index, tors = row
            for st in gmxImp(tors):
                itpf.write('%s' % st)
        itpf.write('\n[ dihedrals ]\n')
        itpf.write('; PROPER DIHEDRAL ANGLES\n')
        itpf.write(
            ';  ai    aj    ak    al funct            c0            c1            c2            c3            c4            c5\n')
        for row in full_tor[full_tor.TY == 'Proper'].iterrows():
            index, tors = row
            for st in gmxDihed(tors):
                itpf.write('%s' % st)
    ppairs = GMX_pairs(full_tor, full_ang, bnd_df)
    itpf.write('\n[ pairs ]\n')
    for pair in ppairs:
        itpf.write('%s    1\n' % pair)
    itpf.close()
    pdblines = Refine_PDB_file(pdb_file)
    atoms, coos = get_coos_from_pdb(pdblines)
    pdb2gro(atoms, coos, resid)
    return None

def mainBOSS2GMX(resid, clu):
    mol = pickle.load(open(resid + ".p", "rb"))
    if clu:
        pdb_file = 'clu.pdb'
    else:
        pdb_file = 'plt.pdb'
    boss2gmx(resid, mol, pdb_file)
    return None