# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import os
import onnx
import onnx.numpy_helper
import struct
from pathlib import Path

import numpy as np

from onnx import onnx_pb as onnx_proto
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel

from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes
from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg
from .quant_utils import QuantType

from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer, QLinearOpsRegistry, IntegerOpsRegistry

from .onnx_model import ONNXModel
from .onnx_quantizer import ONNXQuantizer
from .calibrate import CalibrationDataReader, calibrate

def optimize_model(model_path: Path):
    '''
        Generate model that applies graph optimization (constant folding,etc.)
        parameter model_path: path to the original onnx model
        return: optimized onnx model
    '''
    opt_model_path = generate_identified_filename(model_path, "-opt")
    sess_option = SessionOptions()
    sess_option.optimized_model_filepath = opt_model_path.as_posix()
    sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
    _ = InferenceSession(model_path.as_posix(), sess_option)
    optimized_model = onnx.load(opt_model_path.as_posix())
    return optimized_model

def load_model(model_path: Path, optimize=True):
    if optimize:
        #optimize the original model
        onnx_model = ONNXModel(optimize_model(Path(model_path)))
        # to support GEMM
        onnx_model.replace_gemm_with_matmul()
        return onnx_model.model

    return onnx.load(Path(model_path))

def quantize(model,
             per_channel=False,
             nbits=8,
             quantization_mode=QuantizationMode.IntegerOps,
             static=False,
             force_fusions=False,
             symmetric_activation=False,
             symmetric_weight=False,
             quantization_params=None,
             nodes_to_quantize=None,
             nodes_to_exclude=None,
             op_types_to_quantize=[]):
    '''
        Given an onnx model, create a quantized onnx model and save it into a file
    :param model: ModelProto to quantize
    :param per_channel: quantize weights per channel
    :param nbits: number of bits to represent quantized data. Currently only supporting 8-bit types
    :param quantization_mode: Can be one of the QuantizationMode types.
        IntegerOps:
            the function will use integer ops. Only ConvInteger and MatMulInteger ops are supported now.
        QLinearOps:
            the function will use QLinear ops. Only QLinearConv and QLinearMatMul ops are supported now.
    :param static:
        True: The inputs/activations are quantized using static scale and zero point values
              specified through quantization_params.
        False: The inputs/activations are quantized using dynamic scale and zero point values
               computed while running the model.
    :param symmetric_activation:
        True: activations are quantized into signed integers.
        False: activations are quantized into unsigned integers.
    :param symmetric_weight:
        True: weights are quantized into signed integers.
        False: weights are quantized into unsigned integers.
    :param quantization_params:
        Dictionary to specify the zero point and scale values for inputs to conv and matmul nodes.
        Should be specified when static is set to True.
        The quantization_params should be specified in the following format:
            {
                "input_name": [zero_point, scale]
            }.
        zero_point should be of type np.uint8 and scale should be of type np.float32.
        example:
            {
                'resnet_model/Relu_1:0': [np.uint8(0), np.float32(0.019539741799235344)],
                'resnet_model/Relu_2:0': [np.uint8(0), np.float32(0.011359662748873234)]
            }
    :param nodes_to_quantize:
        List of nodes names to quantize. When this list is not None only the nodes in this list
        are quantized.
        example:
        [
            'Conv__224',
            'Conv__252'
        ]
    :param nodes_to_exclude:
        List of nodes names to exclude. The nodes in this list will be excluded from quantization
        when it is not None.
    :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default.
    :return: ModelProto with quantization
    '''
    print("Warning: onnxruntime.quantization.quantize is deprecated.\n\
         Please use quantize_static for static quantization, quantize_dynamic for dynamic quantization.")
    if nbits == 8 or nbits == 7:
        input_qType = onnx_proto.TensorProto.INT8 if symmetric_activation else onnx_proto.TensorProto.UINT8
        weight_qType = onnx_proto.TensorProto.INT8 if symmetric_weight else onnx_proto.TensorProto.UINT8
        mode = quantization_mode
        copy_model = onnx_proto.ModelProto()
        copy_model.CopyFrom(model)

        if not op_types_to_quantize or len(op_types_to_quantize) == 0:
            op_types_to_quantize = list(QLinearOpsRegistry.keys()) if static else list(IntegerOpsRegistry.keys())

        quantizer = ONNXQuantizer(copy_model, per_channel, nbits == 7, mode, static, weight_qType, input_qType,
                                  quantization_params, nodes_to_quantize, nodes_to_exclude, op_types_to_quantize)

        quantizer.quantize_model()
        return quantizer.model.model
    else:
        raise ValueError('Only 8 and 7 bit quantization is currently supported')

def quantize_static(model_input,
                    model_output,
                    calibration_data_reader: CalibrationDataReader,
                    op_types_to_quantize=[],
                    per_channel=False,
                    reduce_range=False,
                    activation_type=QuantType.QUInt8,
                    weight_type=QuantType.QUInt8,
                    nodes_to_quantize=[],
                    nodes_to_exclude=[],
                    optimize_model=True,
                    use_external_data_format=False):
    '''
        Given an onnx model and calibration data reader, create a quantized onnx model and save it into a file
    :param model_input: file path of model to quantize
    :param model_output: file path of quantized model
    :param calibration_data_reader: a calibration data reader. It enumerates calibration data and generates inputs for the original model.
    :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default.
    :param op_types: operators to quantize
    :param per_channel: quantize weights per channel
    :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
    :param activation_type: quantization data type of activation
    :param weight_type: quantization data type of weight
    :param nodes_to_quantize:
        List of nodes names to quantize. When this list is not None only the nodes in this list
        are quantized.
        example:
        [
            'Conv__224',
            'Conv__252'
        ]
    :param nodes_to_exclude:
        List of nodes names to exclude. The nodes in this list will be excluded from quantization
        when it is not None.
    :param optimize_model: optimize model before quantization.
    :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default.
    '''

    if activation_type != QuantType.QUInt8:
        raise ValueError("Static quantization only support uint8 for activation now.")

    input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    mode = QuantizationMode.QLinearOps

    if not op_types_to_quantize or len(op_types_to_quantize) == 0:
        op_types_to_quantize = list(QLinearOpsRegistry.keys())

    model = load_model(Path(model_input), optimize_model)

    quantization_params_dict = calibrate(model, calibration_data_reader, op_types_to_quantize, nodes_to_quantize,
                                         nodes_to_exclude)

    quantizer = ONNXQuantizer(
        model,
        per_channel,
        reduce_range,
        mode,
        True,  # static
        weight_qType,
        input_qType,
        quantization_params_dict,
        nodes_to_quantize,
        nodes_to_exclude,
        op_types_to_quantize)

    quantizer.quantize_model()
    quantizer.model.save_model_to_file(model_output, use_external_data_format)

def quantize_dynamic(model_input: Path,
                     model_output: Path,
                     op_types_to_quantize=[],
                     per_channel=False,
                     reduce_range=False,
                     activation_type=QuantType.QUInt8,
                     weight_type=QuantType.QUInt8,
                     nodes_to_quantize=[],
                     nodes_to_exclude=[],
                     optimize_model=True,
                     use_external_data_format=False):
    '''
        Given an onnx model, create a quantized onnx model and save it into a file
    :param model_input: file path of model to quantize
    :param model_output: file path of quantized model
    :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default
    :param per_channel: quantize weights per channel
    :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
    :param nbits: number of bits to represent quantized data. Currently only supporting 8-bit types
    :param activation_type: quantization data type of activation
    :param weight_type: quantization data type of weight
    :param nodes_to_quantize:
        List of nodes names to quantize. When this list is not None only the nodes in this list
        are quantized.
        example:
        [
            'Conv__224',
            'Conv__252'
        ]
    :param nodes_to_exclude:
        List of nodes names to exclude. The nodes in this list will be excluded from quantization
        when it is not None.
    :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default.
    '''

    input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    mode = QuantizationMode.IntegerOps

    if not op_types_to_quantize or len(op_types_to_quantize) == 0:
        op_types_to_quantize = list(IntegerOpsRegistry.keys())

    model = load_model(Path(model_input), optimize_model)
    quantizer = ONNXQuantizer(
        model,
        per_channel,
        reduce_range,
        mode,
        False,  #static
        weight_qType,
        input_qType,
        None,
        nodes_to_quantize,
        nodes_to_exclude,
        op_types_to_quantize)

    quantizer.quantize_model()
    quantizer.model.save_model_to_file(model_output, use_external_data_format)

def quantize_qat(model_input: Path,
                 model_output: Path,
                 op_types_to_quantize=[],
                 per_channel=False,
                 reduce_range=False,
                 activation_type=QuantType.QUInt8,
                 weight_type=QuantType.QUInt8,
                 nodes_to_quantize=[],
                 nodes_to_exclude=[],
                 use_external_data_format=False):
    '''
        Given a quantize-aware traning onnx model, create a quantized onnx model and save it into a file
    :param model_input: file path of model to quantize
    :param model_output: file path of quantized model
    :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default
    :param per_channel: quantize weights per channel
    :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode
    :param activation_type: quantization data type of activation
    :param nodes_to_quantize:
        List of nodes names to quantize. When this list is not None only the nodes in this list
        are quantized.
        example:
        [
            'Conv__224',
            'Conv__252'
        ]
    :param nodes_to_exclude:
        List of nodes names to exclude. The nodes in this list will be excluded from quantization
        when it is not None.
    :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default.
    '''

    input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
    mode = QuantizationMode.IntegerOps

    #optimize the original model
    optimized_model = optimize_model(Path(model_input))

    if not op_types_to_quantize or len(op_types_to_quantize) == 0:
        op_types_to_quantize = list(IntegerOpsRegistry.keys())

    quantizer = ONNXQuantizer(
        optimized_model,
        per_channel,
        reduce_range,
        mode,
        False,  #static
        weight_qType,
        input_qType,
        None,
        nodes_to_quantize,
        nodes_to_exclude,
        op_types_to_quantize)

    quantizer.quantize_model()
    quantizer.model.save_model_to_file(model_output, use_external_data_format)


<<<<<<< HEAD
import os
import tensorflow as tf
import re
import numpy as np
from tensorflow.python import pywrap_tensorflow
model_dir = './checkpoints/resnet_pretrain/50'
inflated_model_dir = './checkpoints/resnet_pretrain/inflated_50'
checkpoint_path = os.path.join(model_dir, "ResNet-L50.ckpt")
# Read data from checkpoint file
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
# Print tensor name and values
os.environ["CUDA_VISIBLE_DEVICES"] = '4'
for name in sorted(var_to_shape_map):
    if 'weights' in name:
        print(name)
with tf.Session() as sess:
    for name in sorted(var_to_shape_map):
        variable = reader.get_tensor(name)
        new_variable = []
        if 'weights' in name:
            if 'scale1' in name:
                for i in range(5):
                    new_variable.append(variable/5)
                new_variable = np.array(new_variable)
                tf.Variable(initial_value=new_variable, name=name)
                print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
            elif 'scale2' in name:
                if 'a/' in name:
                    for i in range(3):
                        new_variable.append(variable/3)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
                else:
                    for i in range(1):
                        new_variable.append(variable/1)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
            elif 'scale3' in name:
                if 'a/' in name:
                    if int(name.split('/')[1][-1])%2:
                        inflate_num = 3
                    else:
                        inflate_num = 1
                    for i in range(inflate_num):
                        new_variable.append(variable/inflate_num)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
                else:
                    for i in range(1):
                        new_variable.append(variable/1)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
            elif 'scale4' in name:
                if 'a/' in name:
                    if int(name.split('/')[1][-1])%2:
                        inflate_num = 3
                    else:
                        inflate_num = 1
                    for i in range(inflate_num):
                        new_variable.append(variable/inflate_num)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
                else:
                    for i in range(1):
                        new_variable.append(variable/1)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
            elif 'scale5' in name:
                if 'a/' in name:
                    if int(name.split('/')[1][-1])%2:
                        inflate_num = 1
                    else:
                        inflate_num = 3
                    for i in range(inflate_num):
                        new_variable.append(variable/inflate_num)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
                else:
                    for i in range(1):
                        new_variable.append(variable/1)
                    new_variable = np.array(new_variable)
                    tf.Variable(initial_value=new_variable, name=name)
                    print('inflated %s, shape: %s to shape %s'%(name, variable.shape, new_variable.shape))
        else:
            tf.Variable(initial_value=variable, name=name)
            print('keep_variable: %s, shape: %s' % (name, variable.shape))
    init = tf.global_variables_initializer()
    sess.run(init)
    saver = tf.train.Saver()


import asyncio
import numpy as np
import datetime
import copy
import logging

from FFmpeg import FFmpeg
from NewFiles import NewFiles
from utils import *
class AutoVolume:
    # Used by class Audio, uses class FFmpeg
    def __init__(self):
        self.target_LRA = 4

    def autovol(self,
            loudness_type="short",
            mix=0.8,
            max_deviation_gain=20,
            target_lufs=-23
            ):

        self.target_lufs = target_lufs
        self.loudness_max_deviation_gain = max_deviation_gain
        self.loudness_mix = mix
        self.loudness_type = loudness_type
        if self.loudness_type == "integrated":
            self.loudness_block_size = int(3 * self.sample_rate)
        elif self.loudness_type == "short":
            self.loudness_block_size = int(3 * self.sample_rate)
        elif self.loudness_type == "momentary":
            self.loudness_block_size = int(0.4 * self.sample_rate)
        else:
            raise Exception(f'loudness type {loudness_type} not recognised as an option. Please use "integrated", "short" or "momentary"')

        if self.loudness_block_size >= self.signal.shape[1]:
            print(f"Cant autovol this short a signal with loudness type({self.loudness_type}).")
            return None

        if self.loudness_type =='integrated': self.do_integrated()
        else: self.do_short_term_or_momentary()

        self.get_lufs()

    def add_inverse(self):
        self.num_samples = self.signal.shape[1]
        # Reversed audio to beginning and end to help with loudness calculation there.
        inverse_start = self.signal[:,:self.loudness_block_size][:,::-1]
        inverse_end = self.signal[:,-self.loudness_block_size:][:,::-1]
        combined = np.append(inverse_start, self.signal, axis= 1)
        combined = np.append(combined, inverse_end, axis=1)
        self.len_inverse_start = inverse_start.shape[1]
        self.len_inverse_end = inverse_end.shape[1]
        self.signal = combined
        self.num_samples = self.signal.shape[1]

    def do_integrated(self):
        self.get_lufs()
        loudness = self.loudness_integrated[-1]
        multiplier = self.db_to_linear(self.target_lufs - loudness)
        self.signal = self.signal * multiplier

    def do_short_term_or_momentary(self):
        self.add_inverse()
        self.get_lufs()
        self.get_linear_multiplier()
        self.get_volumeride()
        self.volumeride_smoothen()
        self.volumeride_apply()
        self.remove_inverse()
        self.do_integrated()

    def remove_inverse(self):
        self.signal = self.signal[:,
            self.len_inverse_start : - self.len_inverse_end]
        self.num_samples = self.signal.shape[1]

    def get_linear_multiplier(self):
        if self.loudness_type == "momentary":
            loudness = self.loudness_momentary
        elif self.loudness_type == "short":
            loudness = self.loudness_short_term
        elif self.loudness_type == "integrated":
            loudness = self.loudness_integrated

        # Apply a limit to gain increase
        median_loudness = np.median(loudness)
        for i, loud in enumerate(loudness):
            if loud < (median_loudness -
                    self.loudness_max_deviation_gain):
                loudness[i] = loudness[i-1]

        self.linear_multiplier = self.db_to_linear(
            (self.target_lufs - loudness)*self.loudness_mix)

    def get_volumeride(self):
        # Make start/end index arrays
        offset = self.loudness_block_size//2
        self.idx_start = self.loudness_pts - offset
        i = next(x[0] for x in enumerate(self.idx_start) if x[1] > 0)
        self.idx_start = self.idx_start[i:]
        self.linear_multiplier = self.linear_multiplier[i:]
        self.idx_end = self.idx_start[1:]
        self.idx_end = np.append(self.idx_end, self.num_samples-1)

        # Make volumeride
        self.volumeride = np.ones(self.num_samples)
        for start, end, mult in zip(self.idx_start, self.idx_end, self.linear_multiplier):
            self.volumeride[start : end] = mult

    def volumeride_smoothen(self):
        # Apply fades to smoothen transitions
        for start, end in zip(self.idx_start, self.idx_end):
            self.volumeride[start:end] = np.linspace(
                start = self.volumeride[start],
                stop = self.volumeride[end],
                num = end-start)

    def volumeride_apply(self):
        self.signal = self.signal * self.volumeride

class Audio(FFmpeg, AutoVolume):
    def __init__(self,
            arbitrary_input = None,
            filename:str = '',
            filename_target:str = '',
            downloads_dir:str = 'input',
            num_channels:int = None,
            sample_rate:int = None,
            array = None,
            target_codec = None,
            target_lufs = -23
            ):

        FFmpeg.__init__(self)
        AutoVolume.__init__(self)

        self.temp_dir = None

        if arbitrary_input is not None:
            if type(arbitrary_input) is str:
                filename = arbitrary_input
            elif type(arbitrary_input) is np.ndarray:
                array = arbitrary_input
            else:
                raise Exception(f'input of type {type(arbitrary_input)} not recognised. Try str or np.ndarray')

        self.filename = filename
        self.filename_target = filename_target
        self.downloads_dir = downloads_dir
        self.num_channels = num_channels
        self.sample_rate = sample_rate
        self.array = array
        self.target_codec = target_codec
        self.target_lufs = target_lufs

        if self.is_youtube_url():
            self.setup_file_from_youtube_link()

        if self.is_audio_array():
            self.setup_file_from_array()

        ## Check filename
        # def setup_self_from_filename

        # probe will set sample_rate and num_channels from file metadata
        self.probe()

        if sample_rate: self.sample_rate = sample_rate
        if num_channels: self.num_channels = num_channels
        if filename_target != '':
            self.filename_target = filename_target
        else:
            name, ext = self.name_ext(self.filename)
            self.filename_target = f'{name}_output{ext}'
        self.signal = self.load()

    def set_signal(self, signal):
        self.signal = signal
        self.assure_healthy_signal()

    def get_signal(self):
        return self.signal

    def get_mono(self):
        self.make_mono()
        return self.signal_mono

    def getSampleRate(self):
        return self.sample_rate

    def setup_file_from_youtube_link(self):
        """ This function downloads the audio to disk,
        and sets <self.filename> to be the downloaded file.
        If the file already exists in <self.downloads_dir>,
        then it will not redownload.
        """
        if not self.is_downloaded():
            self.download_from_youtube()
            self.Exception_check_youtube_download()
        self.get_filename_of_downloaded_file()
        self.get_youtube_metadata()

    def is_youtube_url(self):
        if r'youtube.com/' in self.filename:
            return True
        else:
            return False

    def is_audio_array(self):
        if type(self.array) == type(np.array([])):
            return True
        else:
            return False

    def setup_file_from_array(self):
        self.signal = self.array
        self.assure_healthy_signal()
        self.sample_rate = self.sample_rate or 48000
        self.num_channels = self.num_channels or self.signal.shape[0]
        if self.filename == '':
            self.create_temporary_filename()
        self.write(path=self.filename)
        logging.info(f'created {self.filename}')

    def assure_correct_signal_shape(self, signal):
        if signal.ndim == 1:
            signal = signal[np.newaxis, :]
        elif signal.ndim == 2:
            shape = signal.shape
            if shape[0] > shape[1]:
                signal = signal.T
        else:
            raise Exception(f'Error: Signal needs to have one or two dimensions. not {signal.ndim}')
        return signal

    def assure_healthy_signal(self):
        self.signal = self.assure_correct_signal_shape(self.signal)
        self.num_channels = self.signal.shape[0]
        if self.signal.dtype is not np.float32:
            self.signal = self.signal.astype(np.float32)
        if self.signal.max() > 1:
            logging.info(f'peak normalizing audio {self.filename} because it exceeds 0 dBFS')
            self.normalize(type_='peak', headroom_db=2)

    def create_temporary_filename(self):
        if self.filename_target != '':
            self.filename = self.filename_target
        else:
            suffix = datetime.datetime.now().strftime("%y%m%d_%H%M%S")
            suffix += '.wav'
            self.filename = os.path.join(os.getcwd(), suffix)
            self.filename_is_temporary = True

    def is_downloaded(self):
        # Todo
        return False

    def Exception_check_youtube_download(self):
        # Todo
        if self.youtube_dl_stderr.startswith('ERROR'):
            raise Exception(
                f'Unable to download "{self.filename}"'
                f' -> {self.youtube_dl_stderr}'
                )

    def download_from_youtube(self):
        # Todo: add check if download succesfull
        cmd = [
            'youtube-dl',
            '-x',
            '--extract-audio',
            '--write-info-json',
            self.filename
            ]

        _stdout, _stderr = self.get_stdout_and_stderr(*cmd)
        self.youtube_dl_stdout = _stdout.decode('iso-8859-1')
        self.youtube_dl_stderr = _stderr.decode('iso-8859-1')

    def get_filename_of_downloaded_file(self):
        # Todo: add check if already downaloaded
        stdout_arr = self.youtube_dl_stdout.split('\n')

        start_of_line_to_find = r'[ffmpeg] Destination: '
        filename_begins = len(start_of_line_to_find)
        downloaded_file = None

        for line in stdout_arr:
            if line.startswith(start_of_line_to_find):
                downloaded_file = line[filename_begins :].strip()
                break

        self.filename = downloaded_file

        start_of_ext = downloaded_file.rfind('.')
        downloaded_file_extless = downloaded_file[:start_of_ext]
        self.youtube_metadata_file = f'{downloaded_file_extless}.info.json'

    def get_youtube_metadata(self):
        with open(self.youtube_metadata_file) as f:
            _metadata_string = f.read()
            self.youtube_metadata = json.loads(_metadata_string)

    def get_signal_as_int16(self):
        return self.signal*32.767

    def copy(self):
        return copy.deepcopy(self)

    def make_mono(self):
        self.signal_mono = np.mean(self.signal, axis=0)

    def __add__(self, y):
        self.signal += y.signal
        return self

def main():
    #audio = Audio(r'/Users/admin/Dropbox/workspace/SigPy/other/test.wav')
    file = r'/Users/admin/Dropbox/workspace/SigPy/other/forevyoungalpha.wav'
    #file = r'/Users/admin/Dropbox/workspace/SigPy/source.wav'
    audio = Audio(file)

    trg = r"/Users/admin/Dropbox/workspace/SigPy/target.wav"
    print(audio)
    #audio.matcheq(target=r'/Users/admin/Dropbox/workspace/SigPy/p1.npy')
    audio.plot_freq(filename='freq.png')
    quit()
    audio.matcheq(target=trg)
    audio.write(path='test2.wav', codec_name="wav")

    quit()
    nf = NewFiles()

    while True:
        for file in nf():
            audio = Audio(filename=file.path, filename_target=file.target_path)
            audio.plot_lufs(filename=f'{file.target_path_extless}_before.png')
            audio.autovol(loudness_type="short", mix=0.8)
            audio.autovol(loudness_type="momentary", mix=0.2)
            audio.limit(max_dbfs=-11)
            audio.plot_lufs(filename=f'{file.target_path_extless}_after.png')
            audio.write(path=self.filename, codec_name="wav")
            #audio.write(path=self.filename)
            # audio.play()

if __name__ == '__main__':
    main()


import numpy as np
from libs.detectors.jetson.pose import PoseEstimator
from libs.detectors.jetson.decoder import PifPafDecoder
import time
from libs.utils.fps_calculator import convert_infr_time_to_fps
import wget
import os

class Detector:
    """
    Perform pose estimation with Openpifpaf model. extract pedestrian's bounding boxes from key-points.
    :param config: Is a ConfigEngine instance which provides necessary parameters.
    """

    def __init__(self, config):
        self.config = config
        # Get the model name from the config
        self.model_name = self.config.DETECTOR_NAME
        # Frames Per Second
        self.fps = None
        self.w, self.h = (self.config.DETECTOR_INPUT_SIZE[0], self.config.DETECTOR_INPUT_SIZE[1])
        self.model_name = "resnet50-193-257.trt"
        self.model_path = 'libs/detectors/jetson/' + self.model_name
        if not os.path.isfile(self.model_path):
            url= "https://media.githubusercontent.com/media/neuralet/neuralet-models/master/jetson-tx2/pose-estimation-openpifpaf/" + self.model_name
            print("model does not exist under: ", self.model_path, "downloading from ", url)
            wget.download(url, self.model_path)

#        self.pose_estimator = PoseEstimator(self.model_path, (self.w, self.h))
#        self.decoder = PifPafDecoder()

    def inference(self, resized_rgb_image):
        """
        This method will perform inference and return the detected bounding boxes
        Args:
            resized_rgb_image: uint8 numpy array with shape (img_height, img_width, channels)
        Returns:
            result: a dictionary contains of [{"id": 0, "bbox": [x1, y1, x2, y2], "score":s%}, {...}, {...}, ...]
        """
        assert resized_rgb_image.shape == (193, 257, 3)
        self.pose_estimator = PoseEstimator(self.model_path, (self.w, self.h))
        self.decoder = PifPafDecoder()
        t_begin = time.perf_counter()
        heads = self.pose_estimator.inference(resized_rgb_image)
        fields = [[field.cpu().numpy() for field in head] for head in heads]
        fields = [
            [[field[i] for field in head] for head in fields]
            for i in range(1) # 1 is image_batch.shape[0] which is num images in batch
            ]
        annotations = self.decoder.decode(fields)
        inference_time = time.perf_counter() - t_begin
        self.fps = convert_infr_time_to_fps(inference_time)
        result = []
        for l in annotations:
            for annotation_object in l:
                pred = annotation_object.data
                bbox_dict = {}
		# extracting face bounding box
                if np.all(pred[[0, 1, 2, 5, 6], -1] > 0.15):
                    x_min_face = int(pred[6, 0]) / self.w
                    x_max_face = int(pred[5, 0]) / self.w
                    y_max_face = int((pred[5, 1] + pred[6, 1]) / 2) / self.h
                    y_eyes = int((pred[1, 1] + pred[2, 1]) / 2) / self.h
                    y_min_face = 2 * y_eyes - y_max_face
                    if (y_max_face - y_min_face > 0) and (x_max_face - x_min_face > 0):
                        h_crop = y_max_face - y_min_face
                        x_min_face = max(0, x_min_face - 0.2 * h_crop)
                        y_min_face = max(0, y_min_face - 0.1 * h_crop)
                        x_max_face = min(self.w, x_min_face + 1 * h_crop)
                        y_max_face = min(self.h, y_min_face + 1 * h_crop)
                        bbox_dict["bbox"] = [y_min_face, x_min_face, y_max_face, x_max_face]

                result.append(bbox_dict)
        del self.pose_estimator, self.decoder
        return result


import networkx as nx
from elementary_step_om.chem import MappedMolecule, Reaction
from itertools import combinations, product
import copy
import numpy as np

from functools import partial
from multiprocessing import Pool

from rdkit import Chem
from rdkit.Chem import rdMolDescriptors, rdchem
from rdkit.Chem.rdchem import ResonanceFlags
from rdkit.Chem.EnumerateStereoisomers import EnumerateStereoisomers
from rdkit.Chem.EnumerateStereoisomers import StereoEnumerationOptions
from rdkit.Chem import rdmolops

from numba import jit
from numba.typed import Dict
from numba import types

from .xyz2mol_local import AC2mol, get_proto_mol

key_type = types.int64
val_type = types.int32[:]

@jit(nopython=True)
def valid_product_numba(adj_matrix, atomic_num):
    """ Check that the produces product is valid according to the valence"""

    max_valence_numba = Dict.empty(key_type=key_type, value_type=val_type)

    max_valence_numba[1] = np.array([0, 1], dtype=np.int32)
    max_valence_numba[5] = np.array([0, 3], dtype=np.int32)
    max_valence_numba[6] = np.array([0, 4], dtype=np.int32)
    max_valence_numba[7] = np.array([0, 4], dtype=np.int32)
    max_valence_numba[8] = np.array([0, 3], dtype=np.int32)
    max_valence_numba[9] = np.array([0, 1], dtype=np.int32)
    max_valence_numba[14] = np.array([0, 4], dtype=np.int32)
    max_valence_numba[15] = np.array([0, 5], dtype=np.int32)
    max_valence_numba[16] = np.array([0, 6], dtype=np.int32)
    max_valence_numba[17] = np.array([0, 1], dtype=np.int32)
    max_valence_numba[35] = np.array([0, 1], dtype=np.int32)
    max_valence_numba[53] = np.array([0, 1], dtype=np.int32)

    natoms = adj_matrix.shape[0]
    product_valence = np.empty(natoms, dtype=np.int64)
    for i in  range(natoms):
        valence = 0
        for conection in adj_matrix[i]:
            valence += conection
        product_valence[i] = valence

    for atom, valence in zip(atomic_num, product_valence):
        if valence > max_valence_numba[atom][1] or valence == max_valence_numba[atom][0]:
            return False
    return True

# TODO: make a version that can be compiled
def not_parallel_reactions(I, frag_idx):
    """I a parallel reaction?

    A parallel reaction is when two fragments A B becomes A' B'
    but doesn't create a bond between the two fragments.
    However, when A B becomes A' B or A B' it's not a parrallel reaction.

    # TODO: This is not doing what it supposed to do. I think.
    """
    if len(frag_idx) == 1:  # Only one fragment.
        return True

    num_frags = len(frag_idx)
    frag_change_matrix = np.zeros((num_frags, num_frags), dtype=np.int32)
    for i, frag_i in enumerate(frag_idx):
        for j, frag_j in enumerate(frag_idx):
            m = I[np.ix_(frag_i, frag_j)]
            frag_change_matrix[i, j] = not np.all(m == 0)

    off_diagonal = frag_change_matrix[
        ~np.eye(frag_change_matrix.shape[0], dtype=bool)
    ].sum()
    diagonal = frag_change_matrix[np.eye(frag_change_matrix.shape[0], dtype=bool)].sum()
    #print(off_diagonal, diagonal, off_diagonal == 0 and diagonal > 1)
    if off_diagonal == 0 and diagonal > 1: # Think < has to change.
        return False

    return True

def reassign_atom_idx(mol):
    """ Assigns RDKit mol atom id to atom mapped id """
    renumber = [(atom.GetIdx(), atom.GetAtomMapNum()) for atom in mol.GetAtoms()]
    new_idx = [idx[0] for idx in sorted(renumber, key=lambda x: x[1])]
    mol = Chem.RenumberAtoms(mol, new_idx)
    rdmolops.AssignStereochemistry(mol, force=True)
    return mol

def get_most_rigid_resonance(mol):
    """ Return the most rigid resonance structure """
    all_resonance_structures = [
        res
        for res in rdchem.ResonanceMolSupplier(mol, ResonanceFlags.UNCONSTRAINED_ANIONS)
    ]

    min_rot_bonds = 9999
    most_rigid_res = copy.deepcopy(mol)
    if len(all_resonance_structures) <= 1:  # 0 is kind weird
        return most_rigid_res

    for res in all_resonance_structures:
        Chem.SanitizeMol(res)
        num_rot_bonds = rdMolDescriptors.CalcNumRotatableBonds(res)
        if num_rot_bonds < min_rot_bonds:
            most_rigid_res = copy.deepcopy(res)
            min_rot_bonds = num_rot_bonds

    return most_rigid_res

class CreateValidIs:

    def __init__(
        self, mapped_molecule, max_num_bonds: int = 2, cd: int = 4, inactive_atoms = []
    ) -> None:

        self._max_bonds = max_num_bonds
        self._cd = cd
        self._inactive_atoms = inactive_atoms

        self._mapped_molecule = mapped_molecule
        self._ac_matrix = self._mapped_molecule.ac_matrix
        self._frag_idx = []
        for frag in Chem.GetMolFrags(mapped_molecule.rd_mol, asMols=True, sanitizeFrags=False):
            self._frag_idx.append([atom.GetAtomMapNum() - 1 for atom in frag.GetAtoms()])

        self._atom_number = np.array([atom.GetAtomicNum() for atom in self._mapped_molecule.rd_mol.GetAtoms()])

    def _make_simple_conversion_matrices(self) -> None:
        """ """
        # Create one-bond conversion matrices
        num_atoms = len(self._mapped_molecule.atom_symbols)

        # inital inactive atoms is atom map numbers.
        # convert to 0 based idx.
        inactive_atoms = [int(x) - 1 for x in self._inactive_atoms]

        make1, break1 = [], []
        for i in range(num_atoms):
            for j in range(num_atoms):
                # Check if i or j is an inactive atoms don't add the conversion matrix.
                if (i in inactive_atoms) or (j in inactive_atoms):
                    continue

                conversion_matrix = np.zeros(self._ac_matrix.shape, np.int8)
                if j > i:
                    if self._ac_matrix [i, j] == 0:
                        conversion_matrix[i, j] = conversion_matrix[j, i] = 1
                        make1.append(conversion_matrix)
                    else:
                        conversion_matrix[i, j] = conversion_matrix[j, i] = -1
                        break1.append(conversion_matrix)

        self._make1 = make1
        self._break1 = break1

    def _make_break_combinations(self):
        """
        """
        product_combination = []
        for make_break in product(range(self._max_bonds + 1), repeat=2):
            if sum(make_break) <= self._cd and max(make_break) <= self._max_bonds:
                product_combination.append(make_break)
        return product_combination[1:]  # skip 0,0 = reactant.

    def __iter__(self):
        """
        """
        valid_valence_filter = partial(valid_product_numba, atomic_num=self._atom_number)
        parallel_reaction_filter = partial(not_parallel_reactions, frag_idx=self._frag_idx)

        self._make_simple_conversion_matrices()

        for num_make, num_break in self._make_break_combinations():
            make_combs = combinations(self._make1, num_make)
            break_combs = combinations(self._break1, num_break)

            for conv_matrix in product(make_combs, break_combs):
                conv_matrix = np.array(sum(conv_matrix, ())).sum(axis=0)
                product_ac_matrix = self._ac_matrix  + conv_matrix

                if valid_valence_filter(product_ac_matrix):
                    if parallel_reaction_filter(conv_matrix):
                        yield product_ac_matrix

class TakeElementaryStep:

    def __init__(
        self, mapped_molecule, max_num_bonds: int = 2, cd: int = 4, inactive_atoms = []
    ) -> None:

        self._max_bonds = max_num_bonds
        self._cd = cd
        self._inactive_atoms = inactive_atoms

        self._mapped_molecule = mapped_molecule
        self._atom_number = [atom.GetAtomicNum() for atom in self._mapped_molecule.rd_mol.GetAtoms()]

    def _ac_to_mapped_products(self, product_ac_matrix):
        """ """
        formal_charge = Chem.GetFormalCharge(self._mapped_molecule.rd_mol)

        proto_mol = get_proto_mol(self._atom_number)
        mol = AC2mol(
            proto_mol,
            product_ac_matrix,
            self._atom_number,
            charge=formal_charge,
            allow_charged_fragments=True,
            use_graph=True,
            use_atom_maps=True,
        )

        if int(Chem.GetFormalCharge(mol)) != formal_charge:
           return None

        return get_most_rigid_resonance(mol)

    def get_isomers(self, product):
        """Produce all combinations of isomers (R/S and cis/trans). But force
        product atoms with unchanged neighbors to the same label chirality as
        the reactant"""

        product = reassign_atom_idx(product)
        reactant = reassign_atom_idx(copy.deepcopy(self._mapped_molecule.rd_mol))

        # Find chiral atoms - including label chirality
        chiral_atoms_product = Chem.FindMolChiralCenters(product, includeUnassigned=True)

        unchanged_atoms = []
        for atom, _ in chiral_atoms_product:
            product_neighbors = [
                a.GetIdx() for a in product.GetAtomWithIdx(atom).GetNeighbors()
            ]
            reactant_neighbors = [
                a.GetIdx() for a in reactant.GetAtomWithIdx(atom).GetNeighbors()
            ]

            if sorted(product_neighbors) == sorted(reactant_neighbors):
                unchanged_atoms.append(atom)

        # make combinations of isomers.
        opts = StereoEnumerationOptions(onlyUnassigned=False, unique=False)
        rdmolops.AssignStereochemistry(
            product, cleanIt=True, flagPossibleStereoCenters=True, force=True
        )

        product_isomers_mols = []
        for product_isomer in EnumerateStereoisomers(product, options=opts):
            rdmolops.AssignStereochemistry(product_isomer, force=True)
            for atom in unchanged_atoms:
                reactant_global_tag = reactant.GetAtomWithIdx(atom).GetProp("_CIPCode")

                # TODO make sure that the _CIPRank is the same for atom in reactant and product.
                product_isomer_global_tag = product_isomer.GetAtomWithIdx(atom).GetProp(
                    "_CIPCode"
                )
                if reactant_global_tag != product_isomer_global_tag:
                    product_isomer.GetAtomWithIdx(atom).InvertChirality()

            #isomer = Chem.MolFromMolBlock(Chem.MolToMolBlock(product_isomer), sanitize=False)
            #Chem.SanitizeMol(isomer)

            mapped_isomer = MappedMolecule(molblock=Chem.MolToMolBlock(product_isomer))
            if mapped_isomer not in product_isomers_mols:
                product_isomers_mols.append(mapped_isomer)

        return product_isomers_mols

    def _AC2mappedisomers(self, AC):
        """ """
        mol = self._ac_to_mapped_products(AC)
        if mol is None:
            return []

        return self.get_isomers(mol)

    def get_products(self, nprocs=1):

        valid_Is = CreateValidIs(
            self._mapped_molecule,
            max_num_bonds=self._max_bonds,
            cd=self._cd,
            inactive_atoms = self._inactive_atoms
        )

        with Pool(processes=nprocs) as pool:
            mapped_products = pool.map(self._AC2mappedisomers, valid_Is)

        tmp_mapped_products = []
        for sublist in mapped_products:
            tmp_mapped_products += sublist

        return tmp_mapped_products

    # TODO: make network nodes/edges in parallel.
    def get_network(self, spin: int = 1, remove_self_loop: bool =True) -> nx.MultiDiGraph:
        """ """
        network = nx.MultiDiGraph()
        start_unmapped_molecule = self._mapped_molecule.get_unmapped_molecule()
        network.add_node(
            start_unmapped_molecule.__hash__(),
            canonical_reactant = start_unmapped_molecule,
            is_run = True
        )

        nodes = []
        edges = []
        # Run this in parallel.
        for mapped_product in self.get_products():
            canonical_product = mapped_product.get_unmapped_molecule()
            reaction = Reaction(
                reactant=self._mapped_molecule,
                product=mapped_product,
                charge=Chem.GetFormalCharge(self._mapped_molecule.rd_mol),
                spin=spin
            )

            # Make new node
            node_name = canonical_product.__hash__()
            node_data = {
                "canonical_reactant": canonical_product,
                "mapped_reactant": mapped_product,
                "is_run": False
            }
            nodes.append((node_name, node_data))

            # Make new reaction
            in_node = start_unmapped_molecule.__hash__()
            out_node = canonical_product.__hash__()
            reaction_key = reaction.__hash__()

            edge_data = {"reaction": reaction}
            edges.append((in_node, out_node, reaction_key, edge_data))

        #this is not super pretty - but it seems to work.
        add_nodes = []
        for node in nodes:
            if node[0] not in network:
                add_nodes.append(node)

        network.add_nodes_from(add_nodes)
        network.add_edges_from(edges)

        if remove_self_loop:
            network.remove_edges_from(nx.selfloop_edges(network))

        return network


import numpy as np
import cv2
import math
import ImageHelpers
from FingerFinder import FingerFinder
from FaceFinder import FaceFinder
import MazeMaker

def nothing(x):
    pass

def main(camera):
    faceFinder = FaceFinder()
    # fingerFinder = FingerFinder()
    # maze_image = MazeMaker.makeBlankImage()
    # MazeMaker.generate_maze_helper(maze_image)
    while(camera.isOpened()):
        ret, img = cap.read()
        img = cv2.flip(img, 1)
        grey = ImageHelpers.apply_filter(img)
        grey = ImageHelpers.brighter(grey, 3.0)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
        faceFinder.drawFaces(grey, img)
        cv2.imshow('Video', img)
        # cv2.imshow('Grey', grey)
        if cv2.waitKey(100) & 0xFF == ord('q'):
            break

if __name__ == '__main__':
    cap = cv2.VideoCapture(0)
    main(cap)


import json
from pathlib import Path

import numpy as np
import pandas as pd

class Index:
    def __init__(self, index_path, index_type, articles_path, mapping, metadata, k, num_workers):
        self.index = self.load_index(index_path, index_type)
        self.index_type = index_type
        self.articles_path = articles_path
        self.mapping = mapping
        self.metadata = metadata
        self.k = k
        self.num_workers = num_workers

    def load_index(self, index_path, index_type):
        if index_type == 'nmslib':
            import nmslib
            index = nmslib.init(method='hnsw', space='cosinesimil')
            index.loadIndex(index_path)
        elif index_type == 'faiss':
            import faiss
            index = faiss.read_index(index_path)
        else:
            raise TypeError('Index type can only be faiss or nmslib.')
        return index

    def search_index(self, sentences, search_embeddings, return_batch_ids=False):
        if self.index_type == 'nmslib':
            batch = self.index.knnQueryBatch(search_embeddings,
                                             k=self.k,
                                             num_threads=self.num_workers)
            batch = np.array(batch)
            batch_ids = batch[:, 0].astype(np.int)
            batch_distances = batch[:, 1].astype(np.float32)
        elif self.index_type == 'faiss':
            batch_distances, batch_ids = self.index.search(np.array(search_embeddings), k=self.k)
        else:
            raise TypeError('Index type can only be faiss or nmslib.')

        results = self._format_results(batch_ids=batch_ids,
                                       batch_distances=batch_distances,
                                       sentences=sentences,
                                       articles_path=self.articles_path,
                                       mapping=self.mapping)
        if return_batch_ids:
            return results, batch_ids
        return results

    def _load_article(self, articles_path, paper_id):
        json_path = Path(articles_path) / (paper_id + '.json')
        with json_path.open() as f:
            article = json.load(f)
        return article

    def _find_metadata(self, paper_id):
        metadata = self.metadata[self.metadata['sha'] == paper_id]
        if len(metadata) == 1:
            metadata = metadata.iloc[0].to_dict()
            return {
                'doi': metadata['doi'] if not pd.isna(metadata['doi']) else 'N/A',
                'url': metadata['url'] if not pd.isna(metadata['url']) else 'N/A',
                'journal': metadata['journal'] if not pd.isna(metadata['journal']) else 'N/A',
                'publish_time': metadata['publish_time'] if not pd.isna(metadata['publish_time']) else 'N/A',
            }
        else:
            return None  # No metadata was found

    def _extract_k_hits(self, ids, distances, sentence, articles_path, sent_article_mapping):
        extracted = {
            "query": sentence,
            "hits": []
        }

        for id, distance in zip(ids, distances):
            mapping = sent_article_mapping[id]
            paragraph_idx = mapping["paragraph_idx"]
            sentence_idx = mapping["sentence_idx"]
            paper_id = mapping["paper_id"]
            article = self._load_article(articles_path=articles_path,
                                         paper_id=paper_id)
            hit = {
                'title': article['metadata']['title'],
                'authors': article['metadata']['authors'],
                'paragraph': article['body_text'][paragraph_idx],
                'sentence': article['body_text'][paragraph_idx]["sentences"][sentence_idx],
                'abstract': article['abstract'],
                'distance': float(distance),
            }
            metadata = self._find_metadata(paper_id)
            if metadata:
                hit['metadata'] = metadata
            extracted["hits"].append(hit)
        return extracted

    def _format_results(self, batch_ids, batch_distances, sentences, articles_path, mapping):
        return [self._extract_k_hits(ids=batch_ids[x],
                                     distances=batch_distances[x],
                                     sentence=query_sentence,
                                     articles_path=articles_path,
                                     sent_article_mapping=mapping) for x, query_sentence in enumerate(sentences)]

def search_args(parser):
    parser.add_argument('--index_path', default="index",
                        help='Path to the created index')
    parser.add_argument('--index_type', default="nmslib", type=str, choices=["nmslib", "faiss"],
                        help='Type of index')
    parser.add_argument('--dataset_path', default="cord_19_dataset_formatted/",
                        help='Path to the extracted dataset')
    parser.add_argument('--model_name_or_path', default='bert-base-nli-mean-tokens')
    parser.add_argument('--batch_size', default=8, type=int,
                        help='Batch size for the transformer model encoding')
    parser.add_argument('--num_workers', default=8, type=int,
                        help='Number of workers to use when parallelizing the index search')
    parser.add_argument('--k', default=10, type=int,
                        help='The top K hits to return from the index')
    parser.add_argument('--device', default='cpu',
                        help='Set to cuda to use the GPU')
    parser.add_argument('--silent', action="store_true",
                        help='Turn off progress bar when searching')
    return parser

def paths_from_dataset_path(dataset_path):
    """
    Creates paths to the files required for searching the index.
    :param dataset_path: The path to the extracted dataset.
    :return: Paths to various important files/folders for searching the index.
    """
    dataset_path = Path(dataset_path)
    articles_path = dataset_path / 'articles/'
    sentences_path = dataset_path / 'cord_19_sentences.txt'
    mapping_path = dataset_path / 'cord_19_sent_to_article_mapping.json'
    metadata_path = dataset_path / 'metadata.csv'
    return articles_path, sentences_path, mapping_path, metadata_path


from __future__ import division
import numpy as np
import cPickle
import os
from copy import deepcopy
import time
from sklearn.metrics import precision_recall_fscore_support
from sklearn.metrics import classification_report
import gensim
from nltk.tag import CRFTagger
import re

from deep_disfluency.language_model.ngram_language_model \
    import KneserNeySmoothingModel
from deep_disfluency.utils.tools \
    import convert_from_inc_disfluency_tags_to_eval_tags
from deep_disfluency.utils.tools import \
    verify_dialogue_data_matrices_from_folder
from deep_disfluency.utils.tools import \
    dialogue_data_and_indices_from_matrix
from deep_disfluency.load.load import load_tags
from deep_disfluency.rnn.elman import Elman
from deep_disfluency.rnn.lstm import LSTM
from deep_disfluency.rnn.test_if_using_gpu import test_if_using_GPU
from deep_disfluency.decoder.hmm import FirstOrderHMM
from deep_disfluency.decoder.noisy_channel import SourceModel
from deep_disfluency.embeddings.load_embeddings import populate_embeddings
from deep_disfluency.feature_extraction.feature_utils import \
    load_data_from_disfluency_corpus_file
from deep_disfluency.evaluation.disf_evaluation import \
    get_tag_data_from_corpus_file
from utils import process_arguments, get_last_n_features
from deep_disfluency.feature_extraction.feature_utils import \
    sort_into_dialogue_speakers

class IncrementalTagger(object):
    """A generic incremental tagging object which can deal with incremental
    word-by-word (and revokable) input. It doesn't do anything.
    """

    def __init__(self, config_file, config_number, saved_model_folder=None):
        print "Initializing Tagger"
        # initialize (word, pos) tuple graph
        self.window_size = 0
        self.word_graph = []
        self.output_tags = []
        self.timing_model = None
        self.timing_model_scaler = None

    def load_model_from_folder(self, saved_model_folder, model_type):
        """Assumes some machine learning model needs to loaded.
        """
        raise NotImplementedError

    def tag_new_word(self, word, pos=None, timing=None, rollback=0):
        """Tags current word and adds to output tags.
        """
        self.rollback(rollback)

    def rollback(self, backwards):
        """Revoke the right frontier of the input and labels back backwards.
        """
        self.word_graph = self.word_graph[:len(self.word_graph)-backwards]
        self.output_tags = self.output_tags[:len(self.output_tags)-backwards]

    def tag_new_prefix(self, prefix, rollback=0):
        self.rollback(rollback)
        for w, p, t in prefix:
            self.tag_new_word(w, p, t)

    def reset(self):
        self.output_tags = []

class DeepDisfluencyTagger(IncrementalTagger):
    """A deep-learning driven incremental disfluency tagger
    (and optionally utterance-segmenter).

    Tags each word with the following:
    <f/> - a fluent word
    <e/> - an edit term word, not necessarily inside a repair structure
    <rms id="N"/> - reparandum start word for repair with ID number N
    <rm id="N"/> - mid-reparandum word for repair N
    <i id="N"/> - interregnum word for repair N
    <rps id="N"/> - repair onset word for repair N
    <rp id="N"/> - mid-repair word for repair N
    <rpn id="N"/> - repair end word for substitution or repetition repair N
    <rpnDel id="N"/> - repair end word for a delete repair N

    If in joint utterance segmentation mode
    according to the config file,
    the following utterance segmentation tags are used:

    <cc/> - a word which continues the current utterance and whose
            following word will continue it
    <ct/> - a word which continues the current utterance and is the
            last word of it
    <tc/> - a word which is the beginning of an utterance and whose following
            word will continue it
    <tt/> - a word constituting an entire utterance
    """
    def __init__(self, config_file=None,
                 config_number=None,
                 saved_model_dir=None,
                 pos_tagger=None,
                 language_model=None,
                 pos_language_model=None,
                 edit_language_model=None,
                 timer=None,
                 timer_scaler=None,
                 use_timing_data=False,
                 use_decoder=True):

        if not config_file:
            config_file = "experiments/experiment_configs.csv"
            config_number = 35
            print "No config file, using default", config_file, config_number

        super(DeepDisfluencyTagger, self).__init__(config_file,
                                                   config_number,
                                                   saved_model_dir)
        print "Processing args from config number {} ...".format(config_number)
        self.args = process_arguments(config_file,
                                      config_number,
                                      use_saved=False,
                                      hmm=True)
        #  separate manual setting
        setattr(self.args, "use_timing_data", use_timing_data)
        print "Intializing model from args..."
        self.model = self.init_model_from_config(self.args)

        # load a model from a folder if specified
        if saved_model_dir:
            print "Loading saved weights from", saved_model_dir
            self.load_model_params_from_folder(saved_model_dir,
                                               self.args.model_type)
        else:
            print "WARNING no saved model params, needs training."
            print "Loading original embeddings"
            self.load_embeddings(self.args.embeddings)

        if pos_tagger:
            print "Loading POS tagger..."
            self.pos_tagger = pos_tagger
        elif self.args.pos:
            print "No POS tagger specified,loading default CRF switchboard one"
            self.pos_tagger = CRFTagger()
            tagger_path = os.path.dirname(os.path.realpath(__file__)) +\
                "/../feature_extraction/crfpostagger"
            self.pos_tagger.set_model_file(tagger_path)

        if self.args.n_language_model_features > 0 or \
                'noisy_channel' in self.args.decoder_type:
            print "training language model..."
            self.init_language_models(language_model,
                                      pos_language_model,
                                      edit_language_model)

        self.timing_model = None
        self.timing_model_scaler = None
        if timer:
            print "loading timer..."
            self.timing_model = timer
            self.timing_model_scaler = timer_scaler
        elif self.args.use_timing_data:
            # self.timing_model = None
            # self.timing_model_scaler = None
            print "No timer specified, using default switchboard one"
            timer_path = os.path.dirname(os.path.realpath(__file__)) +\
                '/../decoder/timing_models/' + \
                'LogReg_balanced_timing_classifier.pkl'
            with open(timer_path, 'rb') as fid:
                self.timing_model = cPickle.load(fid)
            timer_scaler_path = os.path.dirname(os.path.realpath(__file__)) +\
                '/../decoder/timing_models/' + \
                'LogReg_balanced_timing_scaler.pkl'
            with open(timer_scaler_path, 'rb') as fid:
                self.timing_model_scaler = cPickle.load(fid)
                # TODO a hack
                # self.timing_model_scaler.scale_ = \
                #    self.timing_model_scaler.std_.copy()
        else:
            print "Not using timing data"

        print "Loading decoder..."
        self.hmm_dict = deepcopy(self.tag_to_index_map)
        # add the interegnum tag
        if "disf" in self.args.tags:
            intereg_ind = len(self.hmm_dict.keys())
            interreg_tag = \
                "<i/><cc/>" if "uttseg" in self.args.tags else "<i/>"
            self.hmm_dict[interreg_tag] = intereg_ind  # add the interregnum tag

        # decoder_file = os.path.dirname(os.path.realpath(__file__)) + \
        #     "/../decoder/model/{}_tags".format(self.args.tags)
        noisy_channel = None
        if 'noisy_channel' in self.args.decoder_type:
            noisy_channel = SourceModel(self.lm, self.pos_lm,
                                        uttseg=self.args.do_utt_segmentation)
        self.decoder = None
        if use_decoder:
            self.decoder = FirstOrderHMM(
                                self.hmm_dict,
                                markov_model_file=self.args.tags,
                                timing_model=self.timing_model,
                                timing_model_scaler=self.timing_model_scaler,
                                constraint_only=True,
                                noisy_channel=noisy_channel
                                )

        # getting the states in the right shape
        self.state_history = []
        self.softmax_history = []
        # self.convert_to_output_tags = get_conversion_method(self.args.tags)
        self.reset()

    def init_language_models(self, language_model=None,
                             pos_language_model=None,
                             edit_language_model=None):
        clean_model_dir = os.path.dirname(os.path.realpath(__file__)) +\
            "/../data/lm_corpora"
        if language_model:
            self.lm = language_model
        else:
            print "No language model specified, using default switchboard one"
            lm_corpus_file = open(clean_model_dir +
                                  "/swbd_disf_train_1_clean.text")
            lines = [line.strip("\n").split(",")[1] for line in lm_corpus_file
                     if "POS," not in line and not line.strip("\n") == ""]
            split = int(0.9 * len(lines))
            lm_corpus = "\n".join(lines[:split])
            heldout_lm_corpus = "\n".join(lines[split:])
            lm_corpus_file.close()
            self.lm = KneserNeySmoothingModel(
                                        order=3,
                                        discount=0.7,
                                        partial_words=self.args.partial_words,
                                        train_corpus=lm_corpus,
                                        heldout_corpus=heldout_lm_corpus,
                                        second_corpus=None)
        if pos_language_model:
            self.pos_lm = pos_language_model
        elif self.args.pos:
            print "No pos language model specified, \
            using default switchboard one"
            lm_corpus_file = open(clean_model_dir +
                                  "/swbd_disf_train_1_clean.text")
            lines = [line.strip("\n").split(",")[1] for line in lm_corpus_file
                     if "POS," in line and not line.strip("\n") == ""]
            split = int(0.9 * len(lines))
            lm_corpus = "\n".join(lines[:split])
            heldout_lm_corpus = "\n".join(lines[split:])
            lm_corpus_file.close()
            self.pos_lm = KneserNeySmoothingModel(
                                        order=3,
                                        discount=0.7,
                                        partial_words=self.args.partial_words,
                                        train_corpus=lm_corpus,
                                        heldout_corpus=heldout_lm_corpus,
                                        second_corpus=None)
        if edit_language_model:
            self.edit_lm = edit_language_model
        else:
            edit_lm_corpus_file = open(clean_model_dir +
                                       "/swbd_disf_train_1_edit.text")
            edit_lines = [line.strip("\n").split(",")[1]
                          for line in edit_lm_corpus_file
                          if "POS," not in line and not line.strip("\n") == ""]
            edit_split = int(0.9 * len(edit_lines))
            edit_lm_corpus = "\n".join(edit_lines[:edit_split])
            heldout_edit_lm_corpus = "\n".join(edit_lines[edit_split:])
            edit_lm_corpus_file.close()
            self.edit_lm = KneserNeySmoothingModel(
                                        train_corpus=edit_lm_corpus,
                                        heldout_corpus=heldout_edit_lm_corpus,
                                        order=2,
                                        discount=0.7)
            # TODO an object for getting the lm features incrementally
            # in the language model

    def init_model_from_config(self, args):
        # for feat, val in args._get_kwargs():
        #     print feat, val, type(val)
        if not test_if_using_GPU():
            print "Warning: not using GPU, might be a bit slow"
            print "\tAdjust Theano config file ($HOME/.theanorc)"
        print "loading tag to index maps..."
        label_path = os.path.dirname(os.path.realpath(__file__)) +\
            "/../data/tag_representations/{}_tags.csv".format(args.tags)
        word_path = os.path.dirname(os.path.realpath(__file__)) +\
            "/../data/tag_representations/{}.csv".format(args.word_rep)
        pos_path = os.path.dirname(os.path.realpath(__file__)) +\
            "/../data/tag_representations/{}.csv".format(args.pos_rep)
        self.tag_to_index_map = load_tags(label_path)
        self.word_to_index_map = load_tags(word_path)
        self.pos_to_index_map = load_tags(pos_path)
        self.model_type = args.model_type
        vocab_size = len(self.word_to_index_map.keys())
        emb_dimension = args.emb_dimension
        n_hidden = args.n_hidden
        n_extra = args.n_language_model_features + args.n_acoustic_features
        n_classes = len(self.tag_to_index_map.keys())
        self.window_size = args.window
        n_pos = len(self.pos_to_index_map.keys())
        update_embeddings = args.update_embeddings
        lr = args.lr
        print "Initializing model of type", self.model_type, "..."
        if self.model_type == 'elman':
            model = Elman(ne=vocab_size,
                          de=emb_dimension,
                          nh=n_hidden,
                          na=n_extra,
                          n_out=n_classes,
                          cs=self.window_size,
                          npos=n_pos,
                          update_embeddings=update_embeddings)
            self.initial_h0_state = model.h0.get_value()
            self.initial_c0_state = None

        elif self.model_type == 'lstm':
            model = LSTM(ne=vocab_size,
                         de=emb_dimension,
                         n_lstm=n_hidden,
                         na=n_extra,
                         n_out=n_classes,
                         cs=self.window_size,
                         npos=n_pos,
                         lr=lr,
                         single_output=True,
                         cost_function='nll')
            self.initial_h0_state = model.h0.get_value()
            self.initial_c0_state = model.c0.get_value()
        else:
            raise NotImplementedError('No model init for {0}'.format(
                self.model_type))
        return model

    def load_model_params_from_folder(self, model_folder, model_type):
        if model_type in ["lstm", "elman"]:
            self.model.load_weights_from_folder(model_folder)
            self.initial_h0_state = self.model.h0.get_value()
            if model_type == "lstm":
                self.initial_c0_state = self.model.c0.get_value()
        else:
            raise NotImplementedError('No weight loading for {0}'.format(
                model_type))

    def load_embeddings(self, embeddings_name):
        # load pre-trained embeddings
        embeddings_dir = os.path.dirname(os.path.realpath(__file__)) +\
                                "/../embeddings/"
        pretrained = gensim.models.Word2Vec.load(embeddings_dir +
                                                 embeddings_name)
        print "emb shape", pretrained[pretrained.index2word[0]].shape
        # print pretrained[0].shape
        # assign and fill in the gaps
        emb = populate_embeddings(self.args.emb_dimension,
                                  len(self.word_to_index_map.items()),
                                  self.word_to_index_map,
                                  pretrained)
        self.model.load_weights(emb=emb)

    def standardize_word_and_pos(self, word, pos=None,
                                 proper_name_pos_tags=["NNP", "NNPS",
                                                       "CD", "LS",
                                                       "SYM", "FW"]):
        word = word.lower().replace("'", "")  # no punctuation
        if not pos and self.pos_tagger:
            pos = self.pos_tagger.tag([])  # TODO
        if pos:
            pos = pos.upper()
            if pos in proper_name_pos_tags and "$unc$" not in word:
                word = "$unc$" + word
            if self.pos_to_index_map.get(pos) is None:
                # print "unknown pos", pos
                pos = "<unk>"
        if self.word_to_index_map.get(word) is None:
            # print "unknown word", word
            word = "<unk>"
        return word, pos

    def tag_new_word(self, word, pos=None, timing=None, extra=None,
                     diff_only=True, rollback=0):
        """Tag new incoming word and update the word and tag graphs.

        :param word: the word to consume/tag
        :param pos: the POS tag to consume/tag (optional)
        :param timing: the duration of the word (optional)
        :param diff_only: whether to output only the diffed suffix,
        if False, outputs entire output tags
        :param rollback: the number of words to rollback
        in the case of changed word hypotheses from an ASR
        """
        self.rollback(rollback)
        if pos is None and self.args.pos:
            # if no pos tag provided but there is a pos-tagger, tag word
            test_words = [unicode(x) for x in
                          get_last_n_features(
                                              "words",
                                              self.word_graph,
                                              len(self.word_graph)-1,
                                              n=4
                                              )
                          ] + [unicode(word.lower())]
            pos = self.pos_tagger.tag(test_words)[-1][1]
            # print "tagging", word, "as", pos
        # 0. Add new word to word graph
        word, pos = self.standardize_word_and_pos(word, pos)
        # print "New word:", word, pos
        self.word_graph.append((word, pos, timing))
        # 1. load the saved internal rnn state
        # TODO these nets aren't (necessarily) trained statefully
        # The internal state in training self.args.bs words back
        # are the inital ones in training, however here
        # They are the actual state reached.
        if self.state_history == []:
            c0_state = self.initial_c0_state
            h0_state = self.initial_h0_state
        else:
            if self.model_type == "lstm":
                c0_state = self.state_history[-1][0][-1]
                h0_state = self.state_history[-1][1][-1]
            elif self.model_type == "elman":
                h0_state = self.state_history[-1][-1]

        if self.model_type == "lstm":
            self.model.load_weights(c0=c0_state,
                                    h0=h0_state)
        elif self.model_type == "elman":
            self.model.load_weights(h0=h0_state)
        else:
            raise NotImplementedError("no history loading for\
                             {0} model".format(self.model_type))

        # 2. do the softmax output with converted inputs
        word_window = [self.word_to_index_map[x] for x in
                       get_last_n_features("words", self.word_graph,
                                           len(self.word_graph)-1,
                                           n=self.window_size)
                       ]
        pos_window = [self.pos_to_index_map[x] for x in
                      get_last_n_features("POS", self.word_graph,
                                          len(self.word_graph)-1,
                                          n=self.window_size)
                      ]
        # print "word_window, pos_window", word_window, pos_window
        if self.model_type == "lstm":
            h_t, c_t, s_t = self.model.\
                soft_max_return_hidden_layer([word_window], [pos_window])
            self.softmax_history.append(s_t)
            if len(self.state_history) == 20:  # just saving history
                self.state_history.pop(0)  # pop first one
            self.state_history.append((c_t, h_t))
        elif self.model_type == "elman":
            h_t, s_t = self.model.soft_max_return_hidden_layer([word_window],
                                                               [pos_window])
            self.softmax_history.append(s_t)
            if len(self.state_history) == 20:
                self.state_history.pop(0)  # pop first one
            self.state_history.append(h_t)
        else:
            raise NotImplementedError("no softmax implemented for\
                                 {0} model".format(self.model_type))
        softmax = np.concatenate(self.softmax_history)

        # 3. do the decoding on the softmax
        if "disf" in self.args.tags:
            edit_tag = "<e/><cc/>" if "uttseg" in self.args.tags else "<e/>"
            # print self.tag_to_index_map[edit_tag]
            adjustsoftmax = np.concatenate((
                    softmax,
                    softmax[
                        :,
                        self.tag_to_index_map[edit_tag]].reshape(
                            softmax.shape[0],
                            1)), 1)
        else:
            adjustsoftmax = softmax
        last_n_timings = None if ((not self.args.use_timing_data) or
                                  not timing) \
            else get_last_n_features("timings", self.word_graph,
                                     len(self.word_graph)-1,
                                     n=3)
        if not self.decoder:
            # no decoder, just get the arg max
            max_idx = np.argmax(adjustsoftmax[-1])
            # print max_idx
            max_tag = self.hmm_dict.keys()[
                self.hmm_dict.values().index(max_idx)]
            new_tags = [max_tag]
        else:
            new_tags = self.decoder.viterbi_incremental(
                adjustsoftmax, a_range=(len(adjustsoftmax)-1,
                                        len(adjustsoftmax)),
                changed_suffix_only=True,
                timing_data=last_n_timings,
                words=[word])
        # print "new tags", new_tags
        prev_output_tags = deepcopy(self.output_tags)
        self.output_tags = self.output_tags[:len(self.output_tags) -
                                            (len(new_tags) -
                                             1)] + new_tags

        # 4. convert to standardized output format
        if "simple" in self.args.tags:
            for p in range(
                    len(self.output_tags) -
                    (len(new_tags) + 1),
                    len(self.output_tags)):
                rps = self.output_tags[p]
                self.output_tags[p] = rps.replace(
                    'rm-0',
                    'rps id="{}"'.format(p))
                if "<i" in self.output_tags[p]:
                    self.output_tags[p] = self.output_tags[p].\
                        replace("<e/>", "").replace("<i", "<e/><i")
        else:
            # new_words = [word]
            words = get_last_n_features("words", self.word_graph,
                                        len(self.word_graph)-1,
                                        n=len(self.word_graph) -
                                        (self.window_size-1))
            simple_conversion = False
            if simple_conversion:
                if "<e" in new_tags[-1]:
                    self.output_tags[-1] = "<e"
                elif "rm-" in new_tags[-1]:
                    self.output_tags[-1] = "<rps"
                else:
                    self.output_tags[-1] = "<f"
                if "rm-" in new_tags[-1]:
                    rps = re.findall("<rm-[0-9]+\/>", new_tags[-1], re.S)
                    for r in rps:  # should only be one
                        dist = int(r[r.find("-")+1:-2])
                        for o in range(len(self.output_tags)-2, -1, -1):
                            dist -= 1
                            if dist < 0:
                                break
                            if "<e" not in self.output_tags[o]:
                                self.output_tags[o] += "<rm"
            else:
                self.output_tags = \
                    convert_from_inc_disfluency_tags_to_eval_tags(
                        self.output_tags,
                        words,
                        start=len(self.output_tags) -
                        (len(new_tags)),
                        representation=self.args.tags)
        if diff_only:
            for i, old_new in enumerate(zip(prev_output_tags,
                                            self.output_tags)):
                old, new = old_new
                if old != new:
                    return self.output_tags[i:]
            return self.output_tags[len(prev_output_tags):]
        return self.output_tags

    def tag_utterance(self, utterance):
        """Tags entire utterance, only possible on models
        trained on unsegmented data.
        """
        if not self.args.utts_presegmented:
            raise NotImplementedError("Tagger trained on unsegmented data,\
            please call tag_prefix(words) instead.")
        # non segmenting
        self.reset()  # always starts in initial state
        if not self.args.pos:  # no pos tag model
            utterance = [(w, None, t) for w, p, t in utterance]
            # print "Warning: not using pos tags as not pos tag model"
        if not self.args.use_timing_data:
            utterance = [(w, p, None) for w, p, t in utterance]
            # print "Warning: not using timing durations as no timing model"
        for w, p, t in utterance:
            if self.args.pos:
                self.tag_new_word(w, pos=p, timing=t)
        return self.output_tags

    def rollback(self, backwards):
        super(DeepDisfluencyTagger, self).rollback(backwards)
        self.state_history = self.state_history[:len(self.state_history) -
                                                backwards]
        self.softmax_history = self.softmax_history[:
                                                    len(self.softmax_history) -
                                                    backwards]
        if self.decoder:
            self.decoder.rollback(backwards)

    def init_deep_model_internal_state(self):
        if self.model_type == "lstm":
            self.model.load_weights(c0=self.initial_c0_state,
                                    h0=self.initial_h0_state)
        elif self.model_type == "elman":
            self.model.load_weights(h0=self.initial_h0_state)

    def reset(self):
        super(DeepDisfluencyTagger, self).reset()
        self.word_graph = [("<s>", "<s>", 0)] * \
            (self.window_size - 1)
        self.state_history = []
        self.softmax_history = []
        if self.decoder:
            self.decoder.viterbi_init()
        self.init_deep_model_internal_state()

    def evaluate_fast_from_matrices(self, validation_matrices, tag_file,
                                    idx_to_label_dict):
        output = []
        true_y = []
        for v in validation_matrices:
            words_idx, pos_idx, extra, y, indices = v
            if extra:
                output.extend(self.model.classify_by_index(words_idx, indices,
                                                           pos_idx,
                                                           extra))
            else:
                output.extend(self.model.classify_by_index(words_idx, indices,
                                                           pos_idx))
            true_y.extend(y)
        p_r_f_tags = precision_recall_fscore_support(true_y,
                                                     output,
                                                     average='macro')
        tag_summary = classification_report(
                    true_y, output,
                    labels=[i for i in xrange(len(idx_to_label_dict.items()))],
                    target_names=[
                        idx_to_label_dict[i]
                        for i in xrange(len(idx_to_label_dict.items()))
                                  ]
                                            )
        print tag_summary
        results = {"f1_rmtto": p_r_f_tags[2], "f1_rm": p_r_f_tags[2],
                   "f1_tto1": p_r_f_tags[2], "f1_tto2": p_r_f_tags[2]}

        results.update({
                    'f1_tags': p_r_f_tags[2],
                    'tag_summary': tag_summary
        })
        return results

    def train_net(self, train_dialogues_filepath=None,
                  validation_dialogues_filepath=None,
                  model_dir=None,
                  tag_accuracy_file_path=None):
        """Train the internal deep learning model
        from a list of dialogue matrices.
        """
        tag_accuracy_file = open(tag_accuracy_file_path, "a")
        print "Verifying files..."
        for filepath in [train_dialogues_filepath,
                         validation_dialogues_filepath]:
            if not verify_dialogue_data_matrices_from_folder(
                            filepath,
                            word_dict=self.word_to_index_map,
                            pos_dict=self.pos_to_index_map,
                            tag_dict=self.tag_to_index_map,
                            n_lm=self.args.n_language_model_features,
                            n_acoustic=self.args.n_acoustic_features):
                raise Exception("Dialogue vectors in wrong format!\
                See README.md.")
        lr = self.args.lr  # even if decay, start with specific lr
        n_extra = self.args.n_language_model_features + \
            self.args.n_acoustic_features
        # validation matrices filepath much smaller so can store these
        # and preprocess them all:
        validation_matrices = [np.load(
                                    validation_dialogues_filepath + "/" + fp)
                               for fp in os.listdir(
                                validation_dialogues_filepath)]
        validation_matrices = [dialogue_data_and_indices_from_matrix(
                                  d_matrix,
                                  n_extra,
                                  pre_seg=self.args.utts_presegmented,
                                  window_size=self.window_size,
                                  bs=self.args.bs,
                                  tag_rep=self.args.tags,
                                  tag_to_idx_map=self.tag_to_index_map,
                                  in_utterances=self.args.utts_presegmented)
                               for d_matrix in validation_matrices
                               ]
        idx_2_label_dict = {v: k for k, v in self.tag_to_index_map.items()}
        if not os.path.exists(model_dir):
            os.mkdir(model_dir)
        start = 1  # by default start from the first epoch
        best_score = 0
        best_epoch = 0
        print "Net training started..."
        for e in range(start, self.args.n_epochs + 1):
            tic = time.time()
            epoch_folder = model_dir + "/epoch_{}".format(e)
            if not os.path.exists(epoch_folder):
                os.mkdir(epoch_folder)
            train_loss = 0
            # TODO IO is slow, where the memory allows do in one
            load_separately = True
            test = False
            if load_separately:
                for i, dialogue_f in enumerate(
                                        os.listdir(train_dialogues_filepath)):
                    if test and i > 3:
                        break
                    print dialogue_f
                    d_matrix = np.load(train_dialogues_filepath + "/" +
                                       dialogue_f)
                    word_idx, pos_idx, extra, y, indices = \
                        dialogue_data_and_indices_from_matrix(
                                          d_matrix,
                                          n_extra,
                                          window_size=self.window_size,
                                          bs=self.args.bs,
                                          pre_seg=self.args.utts_presegmented
                                                              )
                    # for i in range(len(indices)):
                    #     print i, word_idx[i], pos_idx[i], \
                    #     y[i], indices[i]
                    train_loss += self.model.fit(word_idx,
                                                 y,
                                                 lr,
                                                 indices,
                                                 pos_idx=pos_idx,
                                                 extra_features=extra)
                    print '[learning] file %i >>' % (i+1),\
                        'completed in %.2f (sec) <<\r' % (time.time() - tic)
            # save the initial states we've learned to override the random
            self.initial_h0_state = self.model.h0.get_value()
            if self.args.model_type == "lstm":
                self.initial_c0_state = self.model.c0.get_value()
            # reset and evaluate simply
            self.reset()
            results = self.evaluate_fast_from_matrices(
                                        validation_matrices,
                                        tag_accuracy_file,
                                        idx_to_label_dict=idx_2_label_dict
                                        )
            val_score = results['f1_tags']  #TODO get best score type
            print "epoch training loss", train_loss
            print '[learning] epoch %i >>' % (e),\
                'completed in %.2f (sec) <<\r' % (time.time() - tic)
            print "validation score", val_score
            tag_accuracy_file.write(str(e) + "\n" + results['tag_summary'] +
                                    "\n%%%%%%%%%%\n")
            tag_accuracy_file.flush()
            print "saving model..."
            self.model.save(epoch_folder)  # Epoch file dump
            # checking patience and decay, if applicable
            # stopping criterion
            if val_score > best_score:
                self.model.save(model_dir)
                best_score = val_score
                print 'NEW BEST raw labels at epoch ', e, 'best valid',\
                    best_score
                best_epoch = e
            # stopping criteria = if no improvement in 10 epochs
            if e - best_epoch >= 10:
                print "stopping, no improvement in 10 epochs"
                break
            if self.args.decay and (e - best_epoch) > 1:
                # just a steady decay if things aren't improving for 2 epochs
                # a hidden hyperparameter
                decay_rate = 0.85
                lr *= decay_rate
                print "learning rate decayed, now ", lr
            if lr < 1e-5:
                print "stopping, below learning rate threshold"
                break
            print '[learning and testing] epoch %i >>' % (e),\
                'completed in %.2f (sec) <<\r' % (time.time()-tic)

        print 'BEST RESULT: epoch', best_epoch, 'valid score', best_score
        tag_accuracy_file.close()
        return best_epoch

    def get_output_tags(self, with_words=False):
        if with_words:
            return zip(self.output_tags,
                       self.word_graph[self.window_size-1:])
        return self.output_tags

    def incremental_output_from_file(self, source_file_path,
                                     target_file_path=None,
                                     is_asr_results_file=False):
        """Return the incremental output in an increco style
        given the incoming words + POS. E.g.:

        Speaker: KB3_1

        Time: 1.50
        KB3_1:1    0.00    1.12    $unc$yes    NNP    <f/><tc/>

        Time: 2.10
        KB3_1:1    0.00    1.12    $unc$yes    NNP    <rms id="1"/><tc/>
        KB3_1:2    1.12    2.00     because    IN    <rps id="1"/><cc/>

        Time: 2.5
        KB3_1:2    1.12    2.00     because    IN    <rps id="1"/><rpndel id="1"/><cc/>

        from an ASR increco style input without the POStags:

        or a normal style disfluency dectection ground truth corpus:

        Speaker: KB3_1
        KB3_1:1    0.00    1.12    $unc$yes    NNP    <rms id="1"/><tc/>
        KB3_1:2    1.12    2.00     $because    IN    <rps id="1"/><cc/>
        KB3_1:3    2.00    3.00    because    IN    <f/><cc/>
        KB3_1:4    3.00    4.00    theres    EXVBZ    <f/><cc/>
        KB3_1:6    4.00    5.00    a    DT    <f/><cc/>
        KB3_1:7    6.00    7.10    pause    NN    <f/><cc/>

        :param source_file_path: str, file path to the input file
        :param target_file_path: str, file path to output in the above format
        :param is_asr_results_file: bool, whether the input is increco style
        """
        if target_file_path:
            target_file = open(target_file_path, "w")
        if not self.args.do_utt_segmentation:
            print "not doing utt seg, using pre-segmented file"
        if is_asr_results_file:
            return NotImplementedError
        if 'timings' in source_file_path:
            print "input file has timings"
            if not is_asr_results_file:
                dialogues = []
                IDs, timings, words, pos_tags, labels = \
                    get_tag_data_from_corpus_file(source_file_path)
                for dialogue, a, b, c, d in zip(IDs,
                                                timings,
                                                words,
                                                pos_tags,
                                                labels):
                    dialogues.append((dialogue, (a, b, c, d)))
        else:
            print "no timings in input file, creating fake timings"
            raise NotImplementedError

        for speaker, speaker_data in dialogues:
            # if "4565" in speaker: quit()
            print speaker
            self.reset()  # reset at the beginning of each dialogue
            if target_file_path:
                target_file.write("Speaker: " + str(speaker) + "\n\n")
            timing_data, lex_data, pos_data, labels = speaker_data
            # iterate through the utterances
            # utt_idx = -1
            current_time = 0
            for i in range(0, len(timing_data)):
                # print i, timing_data[i]
                _, end = timing_data[i]
                if (not self.args.do_utt_segmentation) \
                    and self.args.utts_presegmented \
                        and "<t" in labels[i]:
                    self.reset()  # reset after each utt if non pre-seg
                # utt_idx = frames[i]
                timing = None
                if 'timings' in source_file_path and self.args.use_timing_data:
                    timing = end - current_time
                word = lex_data[i]
                pos = pos_data[i]
                diff = self.tag_new_word(word, pos, timing,
                                         diff_only=True,
                                         rollback=0)
                current_time = end
                if target_file_path:
                    target_file.write("Time: " + str(current_time) + "\n")
                    new_words = lex_data[i-(len(diff)-1):i+1]
                    new_pos = pos_data[i-(len(diff)-1):i+1]
                    new_timings = timing_data[i-(len(diff)-1):i+1]
                    for t, w, p, tag in zip(new_timings,
                                            new_words,
                                            new_pos,
                                            diff):
                        target_file.write("\t".join([str(t[0]),
                                                     str(t[1]),
                                                     w,
                                                     p,
                                                     tag]))
                        target_file.write("\n")
                    target_file.write("\n")
            target_file.write("\n")

    def train_decoder(self, tag_file):
        raise NotImplementedError

    def save_decoder_model(self, dir_path):
        raise NotImplementedError