from __future__ import division
from __future__ import print_function

# import re
import sys
import math
from scipy.io import wavfile
import numpy as np
import h5py
from numpy.lib.stride_tricks import as_strided
import tensorflow as tf
from tensorflow.keras.layers import Input, Reshape, Conv2D, BatchNormalization, Softmax, Conv1D, Bidirectional
from tensorflow.keras.layers import MaxPool2D, Dropout, Permute, Flatten, Dense, MaxPool1D, AvgPool1D, Bidirectional, \
    LSTM, Lambda
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.models import Model
import librosa
import librosa.display
from encoder_2 import Encoder
import pyhocon
import os
import pandas as pd
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping

# store as a global variable, since we only support a few models for now
from data_generator import DataGenerator

models = {
    'tiny': None,
    'small': None,
    'medium': None,
    'large': None,
    'full': None
}

res_net = ResNet50(include_top=False, input_shape=(60, 60, 3))
# for layer in res_net.layers[:143]:
#     layer.trainable = False
# for layer in res_net.layers[143:]:
#     layer.trainable = True
for layer in res_net.layers:
    layer.trainable = False

# model_2 = ResNet50(include_top=False, input_shape=(60, 60, 3))
# for layer in model_2.layers:
#     layer.trainable = False

# the model is trained on 16kHz audio
# model_srate = 16000
# max_batch_size = 3000
# sequence_length = 200
# n_labels = 30
# config = pyhocon.ConfigFactory.parse_file("crepe/experiments.conf")['test']

def build_and_load_model(config, task='raga'):
    """
    Build the CNN model and load the weights

    Parameters
    ----------
    model_capacity : 'tiny', 'small', 'medium', 'large', or 'full'
        String specifying the model capacity, which determines the model's
        capacity multiplier to 4 (tiny), 8 (small), 16 (medium), 24 (large),
        or 32 (full). 'full' uses the model size specified in the paper,
        and the others use a reduced number of filters in each convolutional
        layer, resulting in a smaller model that is faster to evaluate at the
        cost of slightly reduced pitch estimation accuracy.

    Returns
    -------
    model : tensorflow.keras.models.Model
        The pre-trained keras model loaded in memory
    """
    model_capacity = config['model_capacity']
    model_srate = config['model_srate']
    hop_size = int(config['hop_size'] * model_srate)
    sequence_length = int((config['sequence_length'] * model_srate - 1024) / hop_size) + 1
    drop_rate_tonic = config['drop_rate_tonic']
    drop_rate_raga = config['drop_rate_raga']
    cutoff = config['cutoff']
    n_frames = 1 + int((model_srate * cutoff - 1024) / hop_size)
    n_seq = int(n_frames // sequence_length)

    n_labels = config['n_labels']

    note_dim = config['note_dim']

    # x = Input(shape=(1024,), name='input2', dtype='float32')
    x_batch = Input(shape=(None, 1024), name='x_input', dtype='float32')
    tonic_batch = Input(shape=(60,), name='tonic_input', dtype='float32')
    pitches_batch = Input(shape=(None, 360), name='pitches_input', dtype='float32')
    random_batch = Input(shape=(), name='random_input', dtype='int32')
    cqt_batch = Input(shape=(60,), name='cqt_input', dtype='float32')
    # transpose_by_batch = Input(shape=(), name='transpose_input', dtype='int32')

    x = x_batch[0]
    tonic_input = tf.expand_dims(tonic_batch[0], 0)
    pitches = pitches_batch[0]
    random = random_batch[0]
    cqt = cqt_batch[0]
    # transpose_by = transpose_by_batch[0]

    transpose_by = tf.random.uniform(shape=(), minval=0, maxval=60, dtype=tf.int32)
    transpose_by = transpose_by * random
    random = tf.cast(random, tf.float32)

    tonic_input = tf.roll(tonic_input, -transpose_by, axis=-1)
    y, note_emb = get_pitch_emb(x, n_seq, n_frames, model_capacity)
    pitch_model = Model(inputs=[x_batch], outputs=y)

    if task == 'pitch':
        return pitch_model
    # pitch_model.load_weights('model/hindustani_pitch_model.hdf5', by_name=True)
    pitch_model.load_weights('model/model-full.h5', by_name=True)

    note_emb = reduce_note_emb_dimensions(note_emb, note_dim)

    red_y = tf.reshape(pitches, [-1, 6, 60])
    red_y = tf.reduce_sum(red_y, axis=1)  # (None, 60)
    red_y = tf.roll(red_y, -transpose_by, axis=1)
    cqt = tf.roll(cqt, -transpose_by, axis=-1)
    # red_y_random = apply_note_random(red_y)
    # red_y = random*red_y_random + (1-random)*red_y
    # Tonic
    tonic_logits = get_tonic_emb(red_y, cqt, note_emb, note_dim, drop_rate_tonic)

    raga_emb = get_raga_emb(red_y, cqt, tonic_logits, note_emb, note_dim, drop_rate_raga)
    raga_logits = Dense(n_labels, activation='softmax', name='raga')(raga_emb)

    loss_weights = config['loss_weights']

    tonic_logits_roll = tf.roll(tonic_logits, transpose_by, axis=-1, name='tonic')

    rag_model = Model(inputs=[pitches_batch, random_batch, cqt_batch], outputs=[raga_logits, tonic_logits_roll])
    cce = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.0)
    rag_model.compile(loss={'raga': cce, 'tf_op_layer_tonic': 'binary_crossentropy'},
                      optimizer='adam', metrics={'raga': 'accuracy', 'tf_op_layer_tonic': 'accuracy'},
                      loss_weights={'raga': loss_weights[0], 'tf_op_layer_tonic': loss_weights[1]})
    rag_model.load_weights('model/hindustani_raga_model.hdf5', by_name=True)

    # rag_model.compile(loss={'raga': 'categorical_crossentropy'},
    #                   optimizer='adam', metrics={'raga': 'accuracy'},
    #                   loss_weights={'raga': loss_weights[0]})

    rag_model.summary()
    # rag_model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    return rag_model

def get_pitch_emb(x, n_seq, n_frames, model_capacity):
    capacity_multiplier = {
        'tiny': 4, 'small': 8, 'medium': 16, 'large': 24, 'full': 32
    }[model_capacity]
    layers = [1, 2, 3, 4, 5, 6]
    filters = [n * capacity_multiplier for n in [32, 4, 4, 4, 8, 16]]
    widths = [512, 64, 64, 64, 64, 64]
    strides = [(4, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]

    z = []
    layers_cache = []
    for i in range(n_seq):
        x_pitch = x[int(i * n_frames / n_seq):int((i + 1) * n_frames / n_seq)]

        if i == 0:
            res = Reshape(target_shape=(1024, 1, 1), name='input-reshape')
            layers_cache.append(res)
            conv_layers = []
        else:
            res = layers_cache[0]
            conv_layers = layers_cache[1]
        y = res(x_pitch)
        m = 0
        for l, f, w, s in zip(layers, filters, widths, strides):
            if i == 0:
                conv_1 = Conv2D(f, (w, 1), strides=s, padding='same',
                                activation='relu', name="conv%d" % l, trainable=False)
                bn_1 = BatchNormalization(name="conv%d-BN" % l)
                mp_1 = MaxPool2D(pool_size=(2, 1), strides=None, padding='valid',
                                 name="conv%d-maxpool" % l, trainable=False)
                do_1 = Dropout(0.25, name="conv%d-dropout" % l)
                conv_layers.append([conv_1, bn_1, mp_1, do_1])
            else:
                conv_1, bn_1, mp_1, do_1 = conv_layers[m]

            y = conv_1(y)
            y = bn_1(y)
            y = mp_1(y)
            y = do_1(y)
            m += 1

        if i == 0:
            den = Dense(360, activation='sigmoid', name="classifier", trainable=False)
            per = Permute((2, 1, 3))
            flat = Flatten(name="flatten")
            layers_cache.append(conv_layers)
            layers_cache.append(den)
            layers_cache.append(per)
            layers_cache.append(flat)
        else:
            den = layers_cache[2]
            per = layers_cache[3]
            flat = layers_cache[4]

        y = per(y)
        y = flat(y)
        y = den(y)
        z.append(y)
    y = tf.concat(z, axis=0)

    return y, den.weights[0]

def reduce_note_emb_dimensions(emb, note_dim):
    note_emb = emb
    note_emb = tf.reduce_mean(tf.reshape(note_emb, [-1, 6, 60]), axis=1)
    note_emb = tf.transpose(note_emb, name='note_emb')  # 60,note_emb
    note_emb = tf.tile(note_emb, [tf.cast(tf.math.ceil(tf.shape(note_emb)[1] / 60), tf.int32), 1])

    singular_values, u, _ = tf.linalg.svd(note_emb)
    sigma = tf.linalg.diag(singular_values)
    sigma = tf.slice(sigma, [0, 0], [tf.shape(note_emb)[-1], note_dim])
    pca = tf.matmul(u, sigma)
    note_emb = pca[:60, :]

    return note_emb

def get_top_notes(red_y):
    c_note = freq_to_cents(31.7 * 2, 25)
    c_note = np.reshape(c_note, [6, 60])
    c_note = np.sum(c_note, axis=0)

    diag_tf = tf.math.reduce_std(red_y, axis=0)
    diag_tf = Dropout(0.2)(diag_tf)
    # diag_tf = tf.reduce_mean(red_y, axis=0)
    diag_tf = AvgPool1D(pool_size=2, strides=1, padding='same')(tf.expand_dims(tf.expand_dims(diag_tf, 0), 2))[0, :, 0]
    diag_tf_p = tf.roll(diag_tf, 1, 0)
    diag_tf_n = tf.roll(diag_tf, -1, 0)
    diag_tf_1 = tf.less_equal(diag_tf_p, diag_tf)
    diag_tf_2 = tf.less_equal(diag_tf_n, diag_tf)
    diag_tf_3 = tf.logical_and(diag_tf_1, diag_tf_2)
    diag_tf_4 = tf.multiply(diag_tf, tf.cast(diag_tf_3, tf.float32))
    diag_tf_4 = tf.cast(diag_tf_4, tf.float32)
    diag_tf_3_ind = tf.where(diag_tf_3)[:, 0]
    diag_tf_3_soft = tf.keras.layers.Lambda(
        lambda x: tf.map_fn(lambda a: tf.roll(tf.constant(c_note), a, axis=-1), x, tf.float64))(diag_tf_3_ind)
    diag_tf_3_soft = tf.reduce_sum(diag_tf_3_soft, axis=0)
    diag_tf_3_soft = tf.cast(diag_tf_3_soft, tf.float32)

    diag_tf_3 = tf.cast(diag_tf_3, tf.float32)
    return diag_tf_4, diag_tf_3_soft, diag_tf_3
    # return diag_tf_3, diag_tf_3

def get_peak_notes(hist):
    diag_tf_p = tf.roll(hist, 1, 0)
    diag_tf_n = tf.roll(hist, -1, 0)
    diag_tf_1 = tf.less_equal(diag_tf_p, hist)
    diag_tf_2 = tf.less_equal(diag_tf_n, hist)
    diag_tf_3 = tf.logical_and(diag_tf_1, diag_tf_2)
    diag_tf_4 = tf.multiply(hist, tf.cast(diag_tf_3, tf.float32))

    return diag_tf_4

def get_hist_emb(red_y, cqt, note_emb, note_dim, indices, topk, is_tonic, drop_rate=0.2):
    top_notes, top_notes_soft, top_notes_hard = get_top_notes(red_y)
    emb_dot = tf.reduce_mean(tf.multiply(note_emb, tf.tile(note_emb[0:1], [60, 1])), axis=1)
    emb_dot = min_max_scale(emb_dot, is_tonic)
    # emb_dot = tf.tile(tf.expand_dims(emb_dot,0), [topk,1])[0]
    hist = tf.reduce_mean(red_y, axis=0)

    hist = min_max_scale(hist, is_tonic)
    # top_notes_soft = min_max_scale(top_notes_soft+top_notes)
    top_notes_soft = min_max_scale(top_notes_soft, is_tonic)
    cqt_mean = min_max_scale(cqt, is_tonic)
    # cqt_std = min_max_scale(tf.math.reduce_std(cqt, axis=1))
    top_notes = min_max_scale(top_notes, is_tonic)

    # hist_cc = tf.expand_dims(hist,1)
    hist_cc_all = []
    # cqt = min_max_scale(tf.reduce_mean(cqt,axis=1))

    hist_cc = tf.transpose(tf.stack([cqt_mean, hist, top_notes, top_notes_soft]))

    for i in range(topk):
        hist_cc_trans = tf.roll(hist_cc, -indices[i], axis=0)
        hist_cc_all.append(hist_cc_trans)
    # values = tf.cast(values, tf.float32)
    hist_cc_all = tf.stack(hist_cc_all)

    if not is_tonic:
        hist_cc_all = Dropout(0.4)(hist_cc_all)
    else:
        hist_cc_all = Dropout(0.3)(hist_cc_all)

    # res_net = ResNet50(include_top=False, input_shape=(60, 60, 3))

    # for layer in res_net.layers[:143]:
    #     layer.trainable = False
    # for layer in res_net.layers[143:]:
    #     layer.trainable = True
    #
    # for layer in res_net.layers:
    #     layer.trainable = False
    # for layer in res_net.layers:
    #     if "BatchNormalization" in layer.__class__.__name__:
    #         layer.trainable = True

    # model = Model(inputs=res_net.input, outputs=res_net.output)
    # model = res_net

    # hist_cc_all = tf.expand_dims(hist_cc_all,3)
    # hist_cc_all = tf.transpose(hist_cc_all, [0, 2, 1,3])
    # matmul = tf.matmul(hist_cc_all, hist_cc_all, transpose_b=True)
    # matmul = tf.transpose(matmul, [0, 2, 3, 1])
    #
    # matmul = tf.keras.applications.resnet.preprocess_input(matmul)
    # z = model(matmul)
    # z = Flatten()(z)
    # z = Dropout(drop_rate)(z)
    # z = Dense(2 * note_dim, activation='relu')(z)
    # return z, top_notes

    z = Conv1D(filters=64, kernel_size=5, strides=1, padding='same', activation='relu')(hist_cc_all)
    z = Conv1D(filters=64, kernel_size=5, strides=1, padding='same', activation='relu')(z)
    z = BatchNormalization()(z)
    z = MaxPool1D(pool_size=2)(z)
    z = Dropout(drop_rate)(z)
    z = Conv1D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu')(z)
    z = Conv1D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu')(z)
    z = BatchNormalization()(z)
    z = MaxPool1D(pool_size=2)(z)
    z = Dropout(drop_rate)(z)
    z = Conv1D(filters=192, kernel_size=3, strides=1, padding='same', activation='relu')(z)
    z = Conv1D(filters=192, kernel_size=3, strides=1, padding='same', activation='relu')(z)
    z = BatchNormalization()(z)
    z = MaxPool1D(pool_size=2)(z)
    z = Dropout(drop_rate)(z)
    z = Flatten()(z)
    z = Dense(2 * note_dim, activation='relu')(z)
    z = Dropout(drop_rate)(z)
    return z, top_notes_hard

def get_tonic_emb(red_y, cqt, note_emb, note_dim, drop_rate=0.2):
    topk=5
    indices = tf.random.uniform(shape=(topk,), minval=0, maxval=60, dtype=tf.int32)
    hist_emb, top_notes = get_hist_emb(red_y, cqt, note_emb, note_dim, indices, topk, True, drop_rate)
    # ndms = get_ndms(red_y, top_notes, None, 1, note_emb, note_dim, drop_rate)
    # tonic_emb = combine(hist_emb, ndms, note_dim, drop_rate)

    tonic_emb = hist_emb

    tonic_logits_norm = []
    tonic_logits = Dense(60, activation='sigmoid')(tonic_emb)
    for i in range(topk):
        tonic_logits_norm.append(tf.roll(tonic_logits[i],indices[i], axis=0))
    tonic_logits = tf.reduce_mean(tonic_logits_norm, axis=0, keepdims=True)

    return tonic_logits

def get_raga_emb(red_y, cqt, tonic_logits, note_emb, note_dim, drop_rate=0.2):
    topk = 5
    rs = topk//2
    re = (topk//2) + 1

    tonic_argmax = tf.argmax(tonic_logits[0])
    indices = tf.range(tonic_argmax - rs, tonic_argmax + re)
    indices = tf.math.mod(60 + indices, 60)
    values = tf.gather(tonic_logits[0], indices)

    hist_emb, top_notes = get_hist_emb(red_y, cqt, note_emb, note_dim, indices, topk, False, drop_rate)
    red_y_clustered = cluster_top_notes(red_y, top_notes)
    ndms, red_y_am = get_ndms(red_y_clustered, top_notes, indices, topk, note_emb, note_dim, drop_rate)
    raga_emb = combine(hist_emb, ndms, note_dim, drop_rate)

    raga_emb = tf.reduce_sum(tf.multiply(raga_emb, tf.expand_dims(values, 1)), axis=0, keepdims=True) / tf.reduce_sum(
        values)
    return raga_emb

def apply_note_random(red_y):
    def apply_random_roll(a):
        r = tf.random.uniform(shape=(), minval=-2, maxval=3, dtype=tf.int32)
        return tf.roll(a, r, axis=-1)

    red_y_roll = tf.keras.layers.Lambda(apply_random_roll)(red_y)
    return red_y_roll

def combine(emb1, emb2, note_dim, drop_rate):
    emb = Dense(2 * note_dim, activation='relu')(tf.concat([emb1, emb2], axis=1))
    emb = Dropout(drop_rate)(emb)
    f = tf.nn.sigmoid(Dense(2 * note_dim)(emb))
    emb = f * emb1 + (1 - f) * emb2

    return emb

def freq_to_cents(freq, std=25):
    frequency_reference = 10
    c_true = 1200 * math.log(freq / frequency_reference, 2)

    cents_mapping = np.linspace(0, 7180, 360) + 1997.3794084376191
    target = np.exp(-(cents_mapping - c_true) ** 2 / (2 * std ** 2))
    return target

def one_hot_note_emb():
    cents_mapping = tf.range(0, 60)
    target = tf.math.exp(-(cents_mapping - 29) ** 2 / (2 * 2 ** 2))
    target = tf.roll(target, -29, axis=-1)
    note_embs = []
    for i in range(60):
        note_embs.append(tf.roll(target, i, axis=-1))
    note_embs = tf.stack(note_embs)
    return note_embs

def get_ndms(red_y_am, top_notes, indices, topk, note_emb, note_dim, drop_rate=0.6):
    is_tonic = indices is None
    note_emb_mat = tf.matmul(note_emb, note_emb, transpose_b=True)
    note_emb_mat = note_emb_mat / note_dim
    note_emb_mat = tf.cast(note_emb_mat, tf.float64)

    red_y_am = get_unique_seq_1(red_y_am)
    red_y_am_first = red_y_am
    matmul_1 = get_ndms_mat(red_y_am)

    red_y_am = get_unique_seq_2(red_y_am)
    red_y_am = get_unique_seq_1(red_y_am)
    matmul_2 = get_ndms_mat(red_y_am)

    red_y_am = get_unique_seq_2(red_y_am)
    red_y_am = get_unique_seq_1(red_y_am)
    matmul_3 = get_ndms_mat(red_y_am)

    matmul_1 = min_max_scale(matmul_1, is_tonic)
    matmul_2 = min_max_scale(matmul_2, is_tonic)
    matmul_3 = min_max_scale(matmul_3, is_tonic)
    matmul = tf.stack([matmul_1, matmul_2, matmul_3, note_emb_mat], axis=2)
    ndms = []

    if indices is None:
        matmul_tmp = tf.roll(matmul, 0, axis=0)
        matmul_tmp = tf.roll(matmul_tmp, 0, axis=1)
        ndms.append(matmul_tmp)
        ndms = tf.stack(ndms)
    else:
        for i in range(topk):
            matmul_tmp = tf.roll(matmul, [-indices[i], -indices[i]], axis=[0, 1])
            ndms.append(matmul_tmp)
        ndms = tf.stack(ndms)

    ndms = Dropout(0.3)(ndms)

    z = Conv2D(filters=64, kernel_size=(5, 5), strides=(1, 1), activation='relu', padding='same')(ndms)
    z = Conv2D(filters=64, kernel_size=(5, 5), strides=(1, 1), activation='relu', padding='same')(z)
    z = BatchNormalization()(z)
    z = MaxPool2D(pool_size=(2, 2), strides=None, padding='valid')(z)
    z = Dropout(drop_rate)(z)
    z = Conv2D(filters=128, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(z)
    z = Conv2D(filters=128, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(z)
    z = BatchNormalization()(z)
    z = MaxPool2D(pool_size=(2, 2), strides=None, padding='valid')(z)
    z = Dropout(drop_rate)(z)
    z = Conv2D(filters=192, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(z)
    z = Conv2D(filters=192, kernel_size=(3, 3), strides=(1, 1), activation='relu', padding='same')(z)
    z = BatchNormalization()(z)
    z = MaxPool2D(pool_size=(2, 2), strides=None, padding='valid')(z)
    z = Flatten()(z)
    z = Dense(2 * note_dim, activation='relu')(z)
    z = Dropout(drop_rate)(z)
    return z, red_y_am_first

def cluster_top_notes(red_y, top_notes):
    c_note = freq_to_cents(31.7 * 2)
    c_note = np.reshape(c_note, [6, 60])
    c_note = np.sum(c_note, axis=0)

    top_notes_where = tf.where(top_notes)[:, 0]
    top_notes_where = tf.tile(tf.expand_dims(top_notes_where, 0), [tf.shape(red_y)[0], 1])
    red_y_arg = tf.argmax(red_y, axis=1)
    red_y_arg = tf.expand_dims(red_y_arg, 1)
    diff = tf.math.abs(red_y_arg - top_notes_where)
    diff = tf.argmin(diff, 1)
    diff_ohe = tf.one_hot(diff, tf.shape(top_notes_where)[1])
    top_notes_where = tf.multiply(tf.cast(top_notes_where, tf.float32), diff_ohe)
    top_notes_where = tf.reduce_max(top_notes_where, axis=1)
    top_notes_where = tf.cast(top_notes_where, tf.int32)
    return top_notes_where

def get_unique_seq_1(arg_y):

    arg_y = tf.concat([[0.], tf.cast(arg_y, tf.float32)], axis=-1)  # None+1

    arg_y_shifted = tf.roll(arg_y, -1, axis=-1)  # 1,None+1

    mask = tf.cast(tf.not_equal(arg_y, arg_y_shifted), tf.float32)  # 1,None+1
    mask = tf.where(mask)[:, 0]
    uni_seq_notes = tf.gather(arg_y_shifted, mask)
    uni_seq_notes = tf.cast(uni_seq_notes, tf.int32)
    return uni_seq_notes

def get_unique_seq_2(ndms):
    temp = tf.equal(ndms, tf.roll(ndms, -2, axis=-1))
    temp1 = tf.cast(tf.logical_not(tf.roll(temp, 1, axis=-1)), tf.int32)
    temp2 = tf.cast(temp, tf.int32)

    uni_seq_notes = tf.multiply(ndms, temp1) + tf.roll(tf.multiply(ndms, temp2), 1, axis=-1)
    return uni_seq_notes

def conv_tonic(pred_tonic, topk):
    pred_tonic = pred_tonic[0]
    kl = tf.keras.losses.BinaryCrossentropy()
    c_note = freq_to_cents(31.7 * 2, 25)
    c_note = np.reshape(c_note, [6, 60])
    c_note = np.sum(c_note, axis=0)
    bces = []
    for i in range(60):
        bces.append(kl(pred_tonic, np.roll(c_note, i, axis=-1)))
    bces = tf.stack(bces)
    bces_argmin = tf.argmin(bces)

    pred_tonic = tf.concat([pred_tonic, pred_tonic, pred_tonic], axis=-1)
    rs = topk // 2
    re = topk // 2 + 1
    values = pred_tonic[60 + bces_argmin - rs:60 + bces_argmin + re]

    indices = tf.range(bces_argmin - rs, bces_argmin + re)
    indices = tf.math.mod(60 + indices, 60)
    return values, indices

def get_ndms_mat(ndms):
    c_note = freq_to_cents(65.41, 25)
    c_note = np.reshape(c_note, [6, 60])
    c_note = np.sum(c_note, axis=0)
    ndms_ohe = tf.keras.layers.Lambda(
        lambda x: tf.map_fn(lambda a: tf.roll(tf.constant(c_note), a, axis=-1), x, tf.float64))(ndms)
    ndms_roll_ohe = tf.roll(ndms_ohe, -1, axis=0)
    matmul = tf.matmul(ndms_ohe, ndms_roll_ohe, transpose_a=True)

    return matmul

def get_tonic_from_cqt(cqt, tonic_emb_size=128, tonic_cnn_filters=128, drop_rate=0.1):
    # chroma, int(note_dim*3/32), int(note_dim/2), drop_rate
    chroma = tf.expand_dims(cqt, 0)
    chroma = tf.expand_dims(chroma, 3)
    y = Conv2D(tonic_cnn_filters, (5, 5), strides=1, padding='same', activation='relu', name='tonic_cnn_1')(chroma)
    y = Conv2D(tonic_cnn_filters, (5, 5), strides=1, padding='same', activation='relu', name='tonic_cnn_2')(y)
    y = Conv2D(tonic_cnn_filters, (5, 5), strides=1, padding='same', activation='relu', name='tonic_cnn_3')(y)
    y = Dropout(drop_rate)(y)
    y = Conv2D(tonic_cnn_filters, (5, 5), strides=1, padding='same', activation='relu', name='tonic_cnn_4')(y)
    y = Conv2D(tonic_cnn_filters, (5, 5), strides=1, padding='same', activation='relu', name='tonic_cnn_5')(
        y)  # (1, 60, -1, 128)
    y = Dropout(drop_rate)(y)
    y = tf.squeeze(y, 0)
    y = tf.reduce_mean(y, 1)
    y = Dense(tonic_emb_size, activation='relu', name='tonic_cnn_dense_1')(y)
    return y  # (60,32)

def get_rag_from_rnn(red_y_am, indices, topk, note_emb_add, note_dim, dropout):
    window = tf.minimum(tf.shape(red_y_am)[0], 150)
    slice_ind = tf.random.uniform(shape=(), minval=0, maxval=tf.shape(red_y_am)[0] - window + 1, dtype=tf.int32)

    red_y_am = red_y_am[slice_ind:slice_ind + window]
    red_y_am_all = []
    for i in range(topk):
        red_y_am_all.append(tf.math.mod(red_y_am + indices[i], 60))
    red_y_am_all = tf.stack(red_y_am_all)
    embs = tf.gather(note_emb_add, red_y_am_all)
    if len(embs.shape) == 2:
        embs = tf.expand_dims(embs, 0)
    rnn_1 = Bidirectional(LSTM(note_dim, return_sequences=False, recurrent_dropout=dropout, dropout=dropout))(embs)
    return Dense(2 * note_dim, activation='relu')(rnn_1)

def get_raga_from_transformer(red_y_am, note_emb_add, note_dim, dropout):
    max_len = 1000
    red_y_am = red_y_am[:, :max_len]
    pad_len = max_len - tf.shape(red_y_am)[1]
    mask_seq = tf.ones_like(red_y_am, tf.float32)

    red_y_am = tf.concat([red_y_am, tf.zeros([10, pad_len], tf.int32)], axis=1)
    mask_seq = tf.concat([mask_seq, tf.zeros([10, pad_len])], axis=1)
    mask_seq = tf.expand_dims(mask_seq, 2)

    raga_enc = Encoder(note_emb_add, enc_num=1, sequence_length=max_len, N=4, size=note_dim)
    encoding = raga_enc.encode(red_y_am, mask_seq, None, True)
    # encoding = encoder_3.encode(red_y_am, note_emb_add, None, mask_seq, 4, d_model=note_dim, dropout=0)
    encoding = tf.reduce_mean(encoding, axis=1)
    # encoding = encoding[:,-1,:]
    return encoding

def min_max_scale(y, is_tonic):
    z = y
    return (z - tf.reduce_mean(z)) / (tf.math.reduce_std(z))

def ffnn(inputs, hidden_size, drop_rate=0.4):
    x = inputs
    for hs in hidden_size:
        den = Dense(hs, activation='relu')(x)
        x = Dropout(drop_rate)(den)
    return x

def get_unique_seq(arg_y):
    arg_y = tf.concat([[0.], tf.cast(arg_y, tf.float32)], axis=-1)  # None+1

    arg_y_shifted = tf.roll(arg_y, -1, axis=-1)  # 1,None+1

    mask = tf.cast(tf.not_equal(arg_y, arg_y_shifted), tf.float32)  # 1,None+1
    mask = tf.where(mask)[:, 0]
    uni_seq_notes = tf.gather(arg_y_shifted, mask)
    uni_seq_notes = tf.cast(uni_seq_notes, tf.int32)
    return uni_seq_notes

def to_local_average_cents(salience, center=None):
    """
    find the weighted average cents near the argmax bin
    """

    if not hasattr(to_local_average_cents, 'cents_mapping'):
        # the bin number-to-cents mapping
        to_local_average_cents.cents_mapping = (
                np.linspace(0, 7180, 360) + 1997.3794084376191)

    if salience.ndim == 1:
        if center is None:
            center = int(np.argmax(salience))
        start = max(0, center - 4)
        end = min(len(salience), center + 5)
        salience = salience[start:end]
        product_sum = np.sum(
            salience * to_local_average_cents.cents_mapping[start:end])
        weight_sum = np.sum(salience)
        return product_sum / weight_sum
        # product_sum = np.sum(
        #     salience * to_local_average_cents.cents_mapping)
        # return product_sum
    if salience.ndim == 2:
        return np.array([to_local_average_cents(salience[i, :]) for i in
                         range(salience.shape[0])])

    raise Exception("label should be either 1d or 2d ndarray")

def train(task, tradition):
    if task == 'tonic':
        train_tonic(tradition)
    elif task == 'raga':
        # tonic_model_path = train_tonic(tradition)
        raga_model_path = 'model/{}_raga_model.hdf5'.format(tradition)

        config = pyhocon.ConfigFactory.parse_file("experiments.conf")[task]
        training_generator = DataGenerator(task, tradition, 'train', config, random=False)
        validation_generator = DataGenerator(task, tradition, 'validate', config, random=False)
        model = build_and_load_model(config, task)
        # model.load_weights(tonic_model_path, by_name=True)
        # model.load_weights('model/model-full.h5', by_name=True)
        # model.fit(generator)
        # model.fit(x=training_generator,
        #                     validation_data=validation_generator, verbose=2, epochs=15, shuffle=True, batch_size=1)

        checkpoint = ModelCheckpoint(raga_model_path, monitor='loss', verbose=1,
                                     save_best_only=True, mode='auto', period=1)
        # early_stopping_callback = EarlyStopping(monitor="val_loss", patience=5, mode="auto", restore_best_weights=True, min_delta=0.0000001)
        model.fit_generator(generator=training_generator,
                            validation_data=validation_generator, verbose=1, epochs=50, shuffle=True,
                            callbacks=[checkpoint])

def train_tonic(tradition):
    task = 'tonic'
    model_path = 'model/{}_tonic_model.hdf5'.format(tradition)
    # if os.path.exists(model_path):
    #     return model_path

    config = pyhocon.ConfigFactory.parse_file("experiments.conf")[task]
    training_generator = DataGenerator(task, tradition, 'train', config)
    validation_generator = DataGenerator(task, tradition, 'validate', config)
    model = build_and_load_model(config, task)
    # model.load_weights('model/model-large.h5', by_name=True)
    # model.fit(generator)
    # model.fit(x=training_generator,
    #                     validation_data=validation_generator, verbose=2, epochs=15, shuffle=True, batch_size=1)
    checkpoint = ModelCheckpoint(model_path, monitor='loss', verbose=1,
                                 save_best_only=True, mode='auto', period=1)
    model.fit_generator(generator=training_generator,
                        validation_data=validation_generator, verbose=1, epochs=15, shuffle=True,
                        callbacks=[checkpoint])

    return model_path

def test(task, tradition):
    config = pyhocon.ConfigFactory.parse_file("experiments.conf")[task]
    test_generator = DataGenerator(task, tradition, 'test', config, random=False)
    model = build_and_load_model(config, task)
    # model.fit(generator)
    # model.fit(x=training_generator,
    #                     validation_data=validation_generator, verbose=2, epochs=15, shuffle=True, batch_size=1)
    # model.load_weights('model/hindustani_tonic_model.hdf5', by_name=True)
    # model.load_weights('model/model-full.h5'.format(tradition, 'tonic'), by_name=True)
    # model.load_weights('model/{}_{}_model.hdf5'.format(tradition, 'tonic'), by_name=True)
    model.load_weights('model/hindustani_raga_model.hdf5', by_name=True)

    for i in range(5):
        model.evaluate_generator(generator=test_generator, verbose=1)
    # p = model.predict_generator(test_generator, verbose=1)
    # print(p[0, :, :, 0])
    # plt.imshow(p[1,:,:,1], cmap='hot', interpolation='nearest')
    # plt.show()

    # plt.imshow(p[1, :, :, 1], cmap='hot', interpolation='nearest')
    # plt.show()
    # plt.imshow(p[1, :, :, 2], cmap='hot', interpolation='nearest')
    # plt.show()
    # plt.imshow(p[1, :, :, 3], cmap='hot', interpolation='nearest')
    # plt.show()

    # print(p)
    # print(np.argmax(p[0]))
    # print(np.max(p, axis=1))
    # print(p[0])
    # print(np.argmax(p[0], axis=1))
    # print(np.max(p[1], axis=1))
    # print(np.argmax(p[1], axis=1))
    # for pi in p[0]:
    #     print(pi)
    # cents = to_local_average_cents(p[1])
    # frequency = 10 * 2 ** (cents / 1200)
    #
    # for f in frequency:
    #     print(f)
    # print('pred', frequency)

def test_pitch(task, tradition):
    config = pyhocon.ConfigFactory.parse_file("experiments.conf")[task]
    model = build_and_load_model(config, task)
    model.load_weights('model/model-full.h5', by_name=True)
    model_srate = config['model_srate']
    step_size = config['hop_size']
    cuttoff = config['cutoff']
    for t in ['train', 'validate', 'test']:
        data_path = config[tradition + '_' + t]
        data = pd.read_csv(data_path, sep='\t')
        data = data.reset_index()

        slice_ind = 0
        k = 0

        while k < data.shape[0]:
            path = data.loc[k, 'path']
            # path = 'E:\\E2ERaga\\data\\RagaDataset\\audio\\f5999e30-d00d-4837-b9c6-5328768ae22d.wav'
            pitch_path = path[:path.index('.wav')] + '.pitch'
            pitch_path = pitch_path.replace('audio', 'pitches')

            pitch_file = open(pitch_path, "w")
            # if os.path.exists(pitch_path):
            #     pitch_file = open(pitch_path, "a")
            # else:
            #     pitch_file = open(pitch_path, "w")
            pitches = []
            while True:
                if slice_ind == 0:
                    print(pitch_path)

                frames, slice_ind = __data_generation_pitch(path, slice_ind, model_srate, step_size, cuttoff)

                p = model.predict(np.array([frames]))
                # p = np.sum(np.reshape(p, [-1,6,60]), axis=1)
                cents = to_local_average_cents(p)
                frequency = 10 * 2 ** (cents / 1200)
                # for p1 in p:
                #     p1 = list(map(str, p1))
                #     p1 = ','.join(p1)
                #     pitches.append(p1)
                # p = ','.join(p)
                pitches.extend(frequency)
                # pitches.extend(p)
                # pitches.append(p)
                if slice_ind == 0:
                    k += 1
                    break
                # frequency = list(map(str, frequency))
            pitches = list(map(str, pitches))
            pitch_file.writelines('\n'.join(pitches))
            pitch_file.close()
            break
        break

def cache_cqt(task, tradition):
    config = pyhocon.ConfigFactory.parse_file("experiments.conf")[task]

    with h5py.File('data/RagaDataset/Hindustani/cqt_cache.hdf5', "w") as f:
        for t in ['train', 'validate', 'test']:
            data_path = config[tradition + '_' + t]
            data = pd.read_csv(data_path, sep='\t')
            data = data.reset_index()
            k = 0
            while k < data.shape[0]:
                path = data.loc[k, 'path']
                mbid = data.loc[k, 'mbid']
                cqt = get_cqt(path)
                f.create_dataset(mbid, data=cqt)
                print(mbid, k)
                k += 1

def get_cqt(path):
    # return tf.ones([1,60], np.float32)

    sr, audio = wavfile.read(path)
    if len(audio.shape) == 2:
        audio = audio.mean(1)  # make mono
    C = np.abs(librosa.cqt(audio, sr=sr, bins_per_octave=60, n_bins=60 * 7, pad_mode='wrap',
                           fmin=librosa.note_to_hz('C1')))
    #     librosa.display.specshow(C, sr=sr,x_axis='time', y_axis='cqt', cmap='coolwarm')

    # fig, ax = plt.subplots()
    c_cqt = librosa.amplitude_to_db(C, ref=np.max)
    c_cqt = np.reshape(c_cqt, [7, 60, -1])
    c_cqt = np.mean(c_cqt, axis=0)
    # c_cqt = np.mean(c_cqt, axis=1, keepdims=True)
    # img = librosa.display.specshow(c_cqt,
    #                                sr=self.model_srate, x_axis='time', y_axis='cqt_note', ax=ax, bins_per_octave=60)
    # ax.set_title('Constant-Q power spectrum')
    # fig.colorbar(img, ax=ax, format="%+2.0f dB")
    return c_cqt

def __data_generation_pitch(path, slice_ind, model_srate, step_size, cuttoff):
    # pitch_path = self.data.loc[index, 'pitch_path']
    # if self.current_data[2] == path:
    #     frames = self.current_data[0]
    #     pitches = self.current_data[1]
    #     pitches = pitches[slice_ind * int(len(frames) / n_cutoff):(slice_ind + 1) * int(len(frames) / n_cutoff)]
    #     frames = frames[slice_ind * int(len(frames) / n_cutoff):(slice_ind + 1) * int(len(frames) / n_cutoff)]
    #     return frames, pitches
    # else:
    #     sr, audio = wavfile.read(path)
    #     if len(audio.shape) == 2:
    #         audio = audio.mean(1)  # make mono
    #     audio = self.get_non_zero(audio)
    sr, audio = wavfile.read(path)
    if len(audio.shape) == 2:
        audio = audio.mean(1)  # make mono
    # audio = self.get_non_zero(audio)

    # audio = audio[:self.model_srate*15]
    # audio = self.mp3_to_wav(path)

    # print(audio[:100])
    audio = np.pad(audio, 512, mode='constant', constant_values=0)
    audio_len = len(audio)
    audio = audio[slice_ind * model_srate * cuttoff:(slice_ind + 1) * model_srate * cuttoff]
    if (slice_ind + 1) * model_srate * cuttoff >= audio_len:
        slice_ind = -1
    # audio = audio[: self.model_srate*self.cutoff]
    hop_length = int(model_srate * step_size)
    n_frames = 1 + int((len(audio) - 1024) / hop_length)
    frames = as_strided(audio, shape=(1024, n_frames),
                        strides=(audio.itemsize, hop_length * audio.itemsize))
    frames = frames.transpose().copy()
    frames -= np.mean(frames, axis=1)[:, np.newaxis]
    frames /= (np.std(frames, axis=1)[:, np.newaxis] + 1e-5)

    return frames, slice_ind + 1

def get_chroma(audio, sr):
    # logC = librosa.amplitude_to_db(np.abs(C))
    # plt.figure(figsize=(15, 5))
    # librosa.display.specshow(logC, sr=sr, x_axis='time', y_axis='cqt_note', fmin=fmin, cmap='coolwarm')

    # hop_length = 512
    # chromagram = librosa.feature.chroma_cqt(audio, sr=sr, hop_length=hop_length)
    # plt.figure(figsize=(15, 5))
    # librosa.display.specshow(chromagram, x_axis='time', y_axis='chroma', hop_length=hop_length, cmap='coolwarm')

    hop_length = 512
    chromagram = librosa.feature.chroma_cens(audio, sr=sr, hop_length=hop_length, n_chroma=60, bins_per_octave=60)
    # plt.figure(figsize=(15, 5))
    # librosa.display.specshow(chromagram, sr=sr,x_axis='time', y_axis='chroma', hop_length=hop_length, cmap='coolwarm',bins_per_octave=60)
    return chromagram

def predict(audio, sr, model_capacity='full',
            viterbi=False, center=True, step_size=10, verbose=1):
    """
    Perform pitch estimation on given audio

    Parameters
    ----------
    audio : np.ndarray [shape=(N,) or (N, C)]
        The audio samples. Multichannel audio will be downmixed.
    sr : int
        Sample rate of the audio samples. The audio will be resampled if
        the sample rate is not 16 kHz, which is expected by the model.
    model_capacity : 'tiny', 'small', 'medium', 'large', or 'full'
        String specifying the model capacity; see the docstring of
        :func:`~crepe.core.build_and_load_model`
    viterbi : bool
        Apply viterbi smoothing to the estimated pitch curve. False by default.
    center : boolean
        - If `True` (default), the signal `audio` is padded so that frame
          `D[:, t]` is centered at `audio[t * hop_length]`.
        - If `False`, then `D[:, t]` begins at `audio[t * hop_length]`
    step_size : int
        The step size in milliseconds for running pitch estimation.
    verbose : int
        Set the keras verbosity mode: 1 (default) will print out a progress bar
        during prediction, 0 will suppress all non-error printouts.

    Returns
    -------
    A 4-tuple consisting of:

        time: np.ndarray [shape=(T,)]
            The timestamps on which the pitch was estimated
        frequency: np.ndarray [shape=(T,)]
            The predicted pitch values in Hz
        confidence: np.ndarray [shape=(T,)]
            The confidence of voice activity, between 0 and 1
        activation: np.ndarray [shape=(T, 360)]
            The raw activation matrix
    """
    activation = get_activation(audio, sr, model_capacity=model_capacity,
                                center=center, step_size=step_size,
                                verbose=verbose)
    confidence = activation.max(axis=1)

    if viterbi:
        cents = to_viterbi_cents(activation)
    else:
        cents = to_local_average_cents(activation)

    frequency = 10 * 2 ** (cents / 1200)
    frequency[np.isnan(frequency)] = 0

    time = np.arange(confidence.shape[0]) * step_size / 1000.0

    # z = np.reshape(activation, [-1, 6, 60])
    # z = np.mean(z, axis=1) #(None, 60)
    # z = np.reshape(z, [-1,12,5])
    # z = np.mean(z, axis=2)  # (None, 12)
    zarg = np.argmax(activation, axis=1)
    zarg = zarg % 60
    zarg = zarg / 5
    # ((((cents - 1997.3794084376191) / 20) % 60) / 5)
    return time, frequency, zarg, confidence, activation

def process_file(file, output=None, model_capacity='full', viterbi=False,
                 center=True, save_activation=False, save_plot=False,
                 plot_voicing=False, step_size=10, verbose=True):
    """
    Use the input model to perform pitch estimation on the input file.

    Parameters
    ----------
    file : str
        Path to WAV file to be analyzed.
    output : str or None
        Path to directory for saving output files. If None, output files will
        be saved to the directory containing the input file.
    model_capacity : 'tiny', 'small', 'medium', 'large', or 'full'
        String specifying the model capacity; see the docstring of
        :func:`~crepe.core.build_and_load_model`
    viterbi : bool
        Apply viterbi smoothing to the estimated pitch curve. False by default.
    center : boolean
        - If `True` (default), the signal `audio` is padded so that frame
          `D[:, t]` is centered at `audio[t * hop_length]`.
        - If `False`, then `D[:, t]` begins at `audio[t * hop_length]`
    save_activation : bool
        Save the output activation matrix to an .npy file. False by default.
    save_plot : bool
        Save a plot of the output activation matrix to a .png file. False by
        default.
    plot_voicing : bool
        Include a visual representation of the voicing activity detection in
        the plot of the output activation matrix. False by default, only
        relevant if save_plot is True.
    step_size : int
        The step size in milliseconds for running pitch estimation.
    verbose : bool
        Print status messages and keras progress (default=True).

    Returns
    -------

    """
    try:
        sr, audio = wavfile.read(file)
    except ValueError:
        print("CREPE: Could not read %s" % file, file=sys.stderr)
        raise

    time, frequency, cents, confidence, activation = predict(
        audio, sr,
        model_capacity=model_capacity,
        viterbi=viterbi,
        center=center,
        step_size=step_size,
        verbose=1 * verbose)

    # write prediction as TSV
    f0_file = output_path(file, ".f0.csv", output)
    f0_data = np.vstack([time, frequency, cents, confidence]).transpose()
    np.savetxt(f0_file, f0_data, fmt=['%.3f', '%.3f', '%.6f', '%.6f'], delimiter=',',
               header='time,frequency,cents,confidence', comments='')
    if verbose:
        print("CREPE: Saved the estimated frequencies and confidence values "
              "at {}".format(f0_file))

    # save the salience file to a .npy file
    if save_activation:
        activation_path = output_path(file, ".activation.npy", output)
        np.save(activation_path, activation)
        if verbose:
            print("CREPE: Saved the activation matrix at {}".format(
                activation_path))

    # save the salience visualization in a PNG file
    if save_plot:
        import matplotlib.cm
        from imageio import imwrite

        plot_file = output_path(file, ".activation.png", output)
        # to draw the low pitches in the bottom
        salience = np.flip(activation, axis=1)
        inferno = matplotlib.cm.get_cmap('inferno')
        image = inferno(salience.transpose())

        if plot_voicing:
            # attach a soft and hard voicing detection result under the
            # salience plot
            image = np.pad(image, [(0, 20), (0, 0), (0, 0)], mode='constant')
            image[-20:-10, :, :] = inferno(confidence)[np.newaxis, :, :]
            image[-10:, :, :] = (
                inferno((confidence > 0.5).astype(np.float))[np.newaxis, :, :])

        imwrite(plot_file, (255 * image).astype(np.uint8))
        if verbose:
            print("CREPE: Saved the salience plot at {}".format(plot_file))


import torch
import cupy as cp
import numpy as np
import math

from .CustomKernel import CustomKernel
from ..util import get_absolute_path

class ComputeCentroidsCuda(CustomKernel):
  def __init__(
      self,
      de=16,
      dk=16,
      sm_size=48*256*4,
    ):
    super(ComputeCentroidsCuda, self).__init__()
    self.de = de
    self.dk = dk
    assert dk * (de + 1) * 4 <= sm_size
    self.tpb = 256
    self.sm_size = sm_size

    with open(get_absolute_path("kernels", "cuda", "compute_centroids.cu"), "r") as f:
      self.kernel = f.read()

    kernel = (self.kernel
      .replace("_DE_", str(de))
      .replace("_DK_", str(dk))
      .replace("_TPB_", str(self.tpb))
      .replace("_NITERS_", str(math.ceil(dk / self.tpb)))
    )

    self.fn = cp.RawKernel(
      kernel,
      'compute_centroids',
      # options=('--maxrregcount=255',),
      # backend='nvcc',
    )

    self.fn.max_dynamic_shared_size_bytes = sm_size
    # print(self.fn.attributes)

  def __call__(self, data, labels, k, centroids=None):
    """
      data: shape = [m, d_vector, n_data]   ###preferably [m, d, n_data]
      labels: shape = [m, n_data]
      centroids: shape = [m, n_clusters, d_vector]
    """
    m, d_vector, n_data = data.shape
    assert labels.shape[1] == n_data
    if centroids is not None:
      assert centroids.shape[0] == d_vector

      data = data.contiguous()
      labels = labels.contiguous()

    # data = data.transpose(1, 2).contiguous()

    centroids = torch.zeros(m, d_vector, k, device="cuda:0", dtype=torch.float32)
    threads_per_block = (self.tpb,)
    blocks_per_grid = (m, math.ceil(k / self.dk) , math.ceil(d_vector / self.de))
    # print("Blocks per grid", blocks_per_grid)

    torch.cuda.synchronize()
    self.fn(
      grid=blocks_per_grid,
      block=threads_per_block,
      shared_mem = self.dk * (self.de + 1) * 4,
      args=[
        data.data_ptr(),
        labels.data_ptr(),
        centroids.data_ptr(),
        m, # n_subquantizers / n_subvectors
        n_data,
        d_vector,
        k # n_clusters
        ],
      stream=self.stream
    )
    torch.cuda.synchronize()
    return centroids


import numpy as np
import os

from .encoder import OsNetEncoder

dirpath = os.path.abspath(os.path.dirname(__file__))

# Declare an encoder object
encoder = OsNetEncoder(
    # input_width=704,
    # input_height=480,
    # weight_filepath="weights/model_weights.pth.tar-40",
    weight_filepath=os.path.join(dirpath, "osnet_ibn_x1_0_imagenet.pth"),
    input_width=512,
    input_height=512,
    batch_size=32,
    num_classes=2022,
    patch_height=256,
    patch_width=128,
    norm_mean=[0.485, 0.456, 0.406],
    norm_std=[0.229, 0.224, 0.225],
    GPU=True)

if __name__ == '__main__':
    pred_imgs = np.random.rand(5, 512, 512, 3)
    pred_imgs = pred_imgs * 255
    pred_imgs = pred_imgs.astype(np.uint8)
    features = encoder.get_features(pred_imgs)


from statistics import mode
from keras.models import load_model
from utils.datasets import get_labels
from utils.inference import detect_faces
from utils.inference import draw_text
from utils.inference import draw_bounding_box
from utils.inference import apply_offsets
from utils.inference import load_detection_model
from utils.preprocessor import preprocess_input
import sys
import cv2,csv,time
import numpy as np
import subprocess

def func(video_path):
    # file to store metadata
    metaData = open('C:/Users/ASUS/Desktop/Face Recognition/trial1/Face Detection and Emotion Analysis/src/final1.csv','a')
    writer = csv.writer(metaData)

    # parameters for loading data and images
    detection_model_path = '../trained_models/detection_models/haarcascade_frontalface_default.xml'
    emotion_model_path = '../trained_models/emotion_models/fer2013_mini_XCEPTION.102-0.66.hdf5'
    emotion_labels = get_labels('fer2013')

    # hyper-parameters for bounding boxes shape
    frame_window = 10
    emotion_offsets = (20, 40)

    # loading models
    face_detection = load_detection_model(detection_model_path)
    emotion_classifier = load_model(emotion_model_path, compile=False)

    # getting input model shapes for inference
    emotion_target_size = emotion_classifier.input_shape[1:3]

    # starting lists for calculating modes
    emotion_window = []

    toc = time.time()
    # starting video streaming
    cv2.namedWindow('window_frame')
    #video_capture = cv2.VideoCapture(sys.argv[1])
    video_capture = cv2.VideoCapture(video_path)
    #video_capture = cv2.VideoCapture('videoplayback.mp4')

    while True:
        bgr_image = video_capture.read()[1]
        gray_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
        rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
        faces = detect_faces(face_detection, gray_image)

        frame_count = int(video_capture.get(cv2.CAP_PROP_POS_FRAMES))

        tic = time.time()

        for face_coordinates in faces:

            x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
            gray_face = gray_image[y1:y2, x1:x2]
            try:
                gray_face = cv2.resize(gray_face, (emotion_target_size))
            except:
                continue

            actor_face = cv2.resize(gray_face,(128,128))
            cv2.imwrite("E:/tensorflow-master/tensorflow/examples/image_retraining/face.jpg",actor_face)

            video_capture.set(1,int(frame_count))
            ret, frame = video_capture.read()
            cv2.imwrite("E:/Object Detection/models-master/tutorials/image/imagenet/object.jpg", gray_image)

            gray_face = preprocess_input(gray_face, True)
            gray_face = np.expand_dims(gray_face, 0)
            gray_face = np.expand_dims(gray_face, -1)
            emotion_prediction = emotion_classifier.predict(gray_face)
            emotion_probability = np.max(emotion_prediction)
            emotion_label_arg = np.argmax(emotion_prediction)
            emotion_text = emotion_labels[emotion_label_arg]
            emotion_window.append(emotion_text)

            s2_out = subprocess.check_output([sys.executable, "E:/tensorflow-master/tensorflow/examples/label_image/label_image.py", "--graph=E:/tmp/output_graph.pb", "--labels=E:/tmp/output_labels.txt", "--input_layer=Mul", "--output_layer=final_result", "--input_mean=128", "--input_std=128", "--image=E:/tensorflow-master/tensorflow/examples/image_retraining/face.jpg"])
            actor_confidence = s2_out.split()[1]
            if(float(actor_confidence)>0.5):
                actor = s2_out.split()[0]
            else:
                actor = ""

            print(s2_out)

            s3_out = subprocess.check_output([sys.executable, "E:/Object Detection/models-master/tutorials/image/imagenet/classify_image.py", "--image_file=E:/Object Detection/models-master/tutorials/image/imagenet/object.jpg"])
            object1 = s3_out.split()[0]
            print(s3_out)

            writer.writerows([[(tic-toc),frame_count,emotion_text,emotion_probability,actor,actor_confidence,face_coordinates,object1]])

            if len(emotion_window) > frame_window:
                emotion_window.pop(0)
            try:
                emotion_mode = mode(emotion_window)
            except:
                continue

            if emotion_text == 'angry':
                color = emotion_probability * np.asarray((255, 0, 0))
            elif emotion_text == 'sad':
                color = emotion_probability * np.asarray((0, 0, 255))
            elif emotion_text == 'happy':
                color = emotion_probability * np.asarray((255, 255, 0))
            elif emotion_text == 'surprise':
                color = emotion_probability * np.asarray((0, 255, 255))
            else:
                color = emotion_probability * np.asarray((0, 255, 0))

            color = color.astype(int)
            color = color.tolist()

            draw_bounding_box(face_coordinates, rgb_image, color)
            draw_text(face_coordinates, rgb_image, emotion_mode, color, 0, -20, 1, 1)
            draw_text(face_coordinates, rgb_image, actor, color, 0, -45, 1, 1)

        bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
        cv2.imshow('window_frame', bgr_image)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break


import sys, pysam, time, os, copy, argparse, random
from collections import Counter
import numpy as np
import multiprocessing as mp
from pysam import VariantFile
from intervaltree import Interval, IntervalTree

def get_cnd_pos(v_pos,cnd_pos, seq='ont'):
    if seq=='ont':
        ls=cnd_pos[abs(cnd_pos-v_pos)<50000]

        ls1_0= [p for p in ls if (p>=v_pos-2000) &  (p<v_pos)][:2]
        ls1_1= [p for p in ls if (p>=v_pos-5000) &  (p<v_pos-2000)][-3:]
        ls1_2= [p for p in ls if (p>=v_pos-10000) & (p<v_pos-5000)][-4:]
        ls1_3= [p for p in ls if (p>=v_pos-20000) & (p<v_pos-10000)][-5:]
        ls1_4= [p for p in ls if                  (p<v_pos-20000)][-6:]

        ls2_0= [p for p in ls if (p>v_pos) & (p<=v_pos+2000)][-2:]
        ls2_1= [p for p in ls if (p>v_pos+2000) & (p<=v_pos+5000)][:3]
        ls2_2= [p for p in ls if (p>v_pos+5000) & (p<=v_pos+10000)][:4]
        ls2_3= [p for p in ls if (p>v_pos+10000) & (p<=v_pos+20000)][:5]
        ls2_4= [p for p in ls if (p>v_pos+20000)][:6]

        ls_total_1=sorted(ls1_0+ls1_1+ls1_2+ls1_3+ls1_4)
        ls_total_2=sorted(ls2_0+ls2_1+ls2_2+ls2_3+ls2_4)

    elif seq=='ul_ont':
        ls=cnd_pos[abs(cnd_pos-v_pos)<100000]

        ls1_0= [p for p in ls if (p>=v_pos-2000) &  (p<v_pos)][:2]
        ls1_1= [p for p in ls if (p>=v_pos-5000) &  (p<v_pos-2000)][-2:]
        ls1_2= [p for p in ls if (p>=v_pos-10000) & (p<v_pos-5000)][-3:]
        ls1_3= [p for p in ls if (p>=v_pos-20000) & (p<v_pos-10000)][-3:]
        ls1_4= [p for p in ls if (p>=v_pos-40000) & (p<v_pos-20000)][-4:]
        ls1_5= [p for p in ls if (p>=v_pos-50000) & (p<v_pos-40000)][-3:]
        ls1_6= [p for p in ls if                  (p<v_pos-50000)][-3:]

        ls2_0= [p for p in ls if (p>v_pos) & (p<=v_pos+2000)][-2:]
        ls2_1= [p for p in ls if (p>v_pos+2000) & (p<=v_pos+5000)][:2]
        ls2_2= [p for p in ls if (p>v_pos+5000) & (p<=v_pos+10000)][:3]
        ls2_3= [p for p in ls if (p>v_pos+10000) & (p<=v_pos+20000)][:3]
        ls2_4= [p for p in ls if (p>v_pos+20000) & (p<=v_pos+40000)][:4]
        ls2_5= [p for p in ls if (p>v_pos+40000) & (p<=v_pos+50000)][:3]
        ls2_6= [p for p in ls if (p>v_pos+50000)][:3]

        ls_total_1=sorted(ls1_0+ls1_1+ls1_2+ls1_3+ls1_4+ls1_5+ls1_6)
        ls_total_2=sorted(ls2_0+ls2_1+ls2_2+ls2_3+ls2_4+ls2_5+ls2_6)

    elif seq=='ul_ont_extreme':
        ls=cnd_pos[abs(cnd_pos-v_pos)<300000]

        ls1_0= [p for p in ls if (p>=v_pos-10000) &  (p<v_pos)][:2]
        ls1_1= [p for p in ls if (p>=v_pos-20000) & (p<v_pos-10000)][-2:]
        ls1_2= [p for p in ls if (p>=v_pos-50000) & (p<v_pos-20000)][-3:]
        ls1_3= [p for p in ls if (p>=v_pos-75000) & (p<v_pos-50000)][-3:]
        ls1_4= [p for p in ls if (p>=v_pos-100000) & (p<v_pos-75000)][-4:]
        ls1_5= [p for p in ls if (p>=v_pos-200000) & (p<v_pos-100000)][-4:]
        ls1_6= [p for p in ls if                  (p<v_pos-200000)][-2:]

        ls2_0= [p for p in ls if (p>v_pos) & (p<=v_pos+10000)][-2:]
        ls2_1= [p for p in ls if (p>v_pos+10000) & (p<=v_pos+20000)][:2]
        ls2_2= [p for p in ls if (p>v_pos+20000) & (p<=v_pos+50000)][:3]
        ls2_3= [p for p in ls if (p>v_pos+50000) & (p<=v_pos+75000)][:3]
        ls2_4= [p for p in ls if (p>v_pos+75000) & (p<=v_pos+100000)][:4]
        ls2_5= [p for p in ls if (p>v_pos+100000) & (p<=v_pos+200000)][:4]
        ls2_6= [p for p in ls if (p>v_pos+200000)][:2]

        ls_total_1=sorted(ls1_0+ls1_1+ls1_2+ls1_3+ls1_4+ls1_5+ls1_6)
        ls_total_2=sorted(ls2_0+ls2_1+ls2_2+ls2_3+ls2_4+ls2_5+ls2_6)

    elif seq=='pacbio':
        ls=cnd_pos[abs(cnd_pos-v_pos)<20000]

        ls1_0= [p for p in ls if (p>=v_pos-2000) &  (p<v_pos)][:4]
        ls1_1= [p for p in ls if (p>=v_pos-5000) &  (p<v_pos-2000)][-5:]
        ls1_2= [p for p in ls if (p>=v_pos-10000) & (p<v_pos-5000)][-5:]
        ls1_3= [p for p in ls if (p>=v_pos-20000) & (p<v_pos-10000)][-6:]

        ls2_0= [p for p in ls if (p>v_pos) & (p<=v_pos+2000)][-4:]
        ls2_1= [p for p in ls if (p>v_pos+2000) & (p<=v_pos+5000)][:5]
        ls2_2= [p for p in ls if (p>v_pos+5000) & (p<=v_pos+10000)][:5]
        ls2_3= [p for p in ls if (p>v_pos+10000) & (p<=v_pos+20000)][:6]

        ls_total_1=sorted(ls1_0+ls1_1+ls1_2+ls1_3)
        ls_total_2=sorted(ls2_0+ls2_1+ls2_2+ls2_3)

    return ls_total_1, ls_total_2

def get_nbr(dct):
    base_to_num_map={'*':4,'A':0,'G':1,'T':2,'C':3,'N':4}

    chrom=dct['chrom']
    start=max(dct['start']-50000,1)
    end=dct['end']+50000

    sam_path=dct['sam_path']
    fasta_path=dct['fasta_path']
    samfile = pysam.Samfile(sam_path, "rb")
    fastafile=pysam.FastaFile(fasta_path)

    rlist=[s for s in fastafile.fetch(chrom,start-1,end-1)]
    output_seq={}

    if dct['supplementary']:
        flag=0x4|0x100|0x200|0x400
    else:
        flag=0x4|0x100|0x200|0x400|0x800

    for pcol in samfile.pileup(chrom,start-1,end-1,min_base_quality=0, flag_filter=flag,truncate=True):
            n=pcol.get_num_aligned()
            r=rlist[pcol.pos+1-start]

            if r in 'AGTC' and n>=dct['mincov']:
                seq=''.join([x[0] for x in pcol.get_query_sequences( mark_matches=False, mark_ends=False,add_indels=True)]).upper()
                alt_freq=max([x[1] for x in Counter(seq).items() if (x[0]!=r and x[0] in 'AGTC')]+[0])/n

                if dct['threshold'][0]<=alt_freq and alt_freq<dct['threshold'][1]:
                    name=pcol.get_query_names()
                    output_seq[pcol.pos+1]={n:base_to_num_map[s] for (n,s) in zip(name,seq)}
                    output_seq[pcol.pos+1]['ref']=base_to_num_map[r]
    return output_seq

def get_snp_testing_candidates(dct):
    base_to_num_map={'*':4,'A':0,'G':1,'T':2,'C':3,'N':4}

    chrom=dct['chrom']
    start=dct['start']
    end=dct['end']
    include_intervals, exclude_intervals=None, None

    if dct['include_bed']:
        tbx = pysam.TabixFile(dct['include_bed'])
        include_intervals=IntervalTree(Interval(int(row[1]), int(row[2]), "%s" % (row[1])) for row in tbx.fetch(chrom, parser=pysam.asBed()))

        def in_bed(tree,pos):
            return tree.overlaps(pos)

        include_intervals=IntervalTree(include_intervals.overlap(start,end))

        if not include_intervals:
            return [],[],[],[],[]

        else:
            start=max(start, min(x[0] for x in include_intervals))
            end=min(end, max(x[1] for x in include_intervals))

    else:
        def in_bed(tree, pos):
            return True

    if dct['exclude_bed']:
        tbx = pysam.TabixFile(dct['exclude_bed'])
        try:
            exclude_intervals=IntervalTree(Interval(int(row[1]), int(row[2]), "%s" % (row[1])) for row in tbx.fetch(chrom, parser=pysam.asBed()))

            def ex_bed(tree, pos):
                return tree.overlaps(pos)

        except ValueError:
            def ex_bed(tree, pos):
                return False
    else:
        def ex_bed(tree, pos):
            return False

    sam_path=dct['sam_path']
    fasta_path=dct['fasta_path']
    threshold=dct['threshold']

    nbr_size=20

    cnd_seq=get_nbr(dct)

    cnd_pos=np.array(list(cnd_seq.keys()))

    samfile = pysam.Samfile(sam_path, "rb")
    fastafile=pysam.FastaFile(fasta_path)

    ref_dict={j:s.upper() if s in 'AGTC' else '*' for j,s in zip(range(max(1,start-40),end+40+1),fastafile.fetch(chrom,max(1,start-40)-1,end+40)) }

    pileup_dict={}

    output={}

    if dct['supplementary']:
        flag=0x4|0x100|0x200|0x400
    else:
        flag=0x4|0x100|0x200|0x400|0x800

    for pcol in samfile.pileup(chrom,max(0,start-1),end,min_base_quality=0,\
                                           flag_filter=flag,truncate=True):

            r=ref_dict[pcol.pos+1]
            if in_bed(include_intervals, pcol.pos+1) and not ex_bed(exclude_intervals, pcol.pos+1) and r in 'AGTC':
                seq=''.join([x[0] for x in pcol.get_query_sequences( mark_matches=False, mark_ends=False,add_indels=True)]).upper()
                name=pcol.get_query_names()
                n=pcol.get_num_aligned()

                if n>=dct['mincov'] and pcol.pos+1>=start and pcol.pos+1<=end:
                    alt_freq=max([x[1] for x in Counter(seq).items() if (x[0]!=r and x[0] in 'AGTC')]+[0])/n

                    if dct['min_allele_freq']<=alt_freq:
                        pileup_dict[pcol.pos+1]={n:base_to_num_map[s] for (n,s) in zip(name,seq)}
                        output[pcol.pos+1]=(n,alt_freq)

    pileup_list=[]

    pos_list=output.keys()

    output_pos,output_ref,output_mat,output_dp,output_freq=[],[],[],[],[]
    if pos_list:

        for v_pos in pos_list:
            np.random.seed(812)
            ls_total_1, ls_total_2 = get_cnd_pos(v_pos, cnd_pos, dct['seq'])

            nbr_dict={}

            rlist=[(v_pos,base_to_num_map[ref_dict[v_pos]]) ]

            sample=list(pileup_dict[v_pos].keys())

            if len(sample) > dct['maxcov']:
                sample=random.sample(sample,min(len(sample), dct['maxcov']))

            sample=set(sample)
            nbr_dict={}

            tmp_mat=np.array([[4]*(len(ls_total_1)+1+len(ls_total_2))]*len(sample))

            for i,name in enumerate(sample):
                for j,nb_pos in enumerate(ls_total_1):
                    try:
                        tmp_mat[i][j]=cnd_seq[nb_pos][name]
                    except KeyError:
                        pass

                tmp_mat[i][len(ls_total_1)]=pileup_dict[v_pos][name]

                for j,nb_pos in enumerate(ls_total_2):

                    try:
                        tmp_mat[i][j +len(ls_total_1)+1]=cnd_seq[nb_pos][name]
                    except KeyError:
                        pass
            for nb_pos in ls_total_1+ls_total_2:
                        rlist.append((nb_pos, cnd_seq[nb_pos]['ref']))

            rlist=sorted(rlist, key=lambda x:x[0])
            total_rlist=np.array([x[1] for x in rlist])

            if len(total_rlist)<dct['min_nbr_sites']:
                continue

            mat=np.dstack([np.sum(np.eye(5)[tmp_mat[tmp_mat[:,len(ls_total_1)]==i]],axis=0) for i in range(4)]).transpose(2,0,1)[:,:,:4]

            total_ref=np.eye(5)[total_rlist.astype(int)]
            total_ref[:,4]=0
            total_ref=total_ref[np.newaxis,:]
            mat=np.dstack([mat,np.zeros([4,mat.shape[1]])+np.eye(4)[base_to_num_map[ref_dict[v_pos]]][:,np.newaxis]])
            data=np.vstack([total_ref,np.multiply(mat,1-2*total_ref)])
            data=np.hstack([np.zeros([5,nbr_size-len(ls_total_1),5]),data,np.zeros([5, nbr_size-len(ls_total_2), 5])]).astype(np.int8)

            output_pos.append(v_pos)
            output_ref.append(base_to_num_map[ref_dict[v_pos]])
            output_mat.append(data)
            output_dp.append(output[v_pos][0])
            output_freq.append(output[v_pos][1])

        if len(output_pos)>0:
            output_mat=np.array(output_mat).astype(np.float32)
            output_pos=np.array(output_pos)

            output_ref=np.eye(max(4,np.max(output_ref)+1))[np.array(output_ref)].astype(np.int8)
            output_ref=output_ref[:,:4]

            output_dp=np.array(output_dp)
            output_freq=np.array(output_freq)

    return (output_pos,output_ref,output_mat,output_dp,output_freq)


"""
This module contains the TreeGrower class.

TreeGrowee builds a regression tree fitting a Newton-Raphson step, based on
the gradients and hessians of the training data.
"""
from heapq import heappush, heappop
import numpy as np
from numba import njit, prange
from time import time

from .splitting import (SplittingContext, split_indices, find_node_split, update_prediction_values)
from .predictor import TreePredictor, PREDICTOR_RECORD_DTYPE

from pygbm.options import OptionSet

class TreeNode:
    """Tree Node class used in TreeGrower.

    This isn't used for prediction purposes, only for training (see
    TreePredictor).

    Parameters
    ----------
    depth : int
        The depth of the node, i.e. its distance from the root
    samples_indices : array of int
        The indices of the samples at the node
    sum_gradients : float
        The sum of the gradients of the samples at the node
    sum_hessians : float
        The sum of the hessians of the samples at the node
    parent : TreeNode or None, optional(default=None)
        The parent of the node. None for root.

    Attributes
    ----------
    depth : int
        The depth of the node, i.e. its distance from the root
    samples_indices : array of int
        The indices of the samples at the node
    sum_g_hf, sum_gx_hfx, sum_h, sum_hx, sum_hx2 : float
        Difference sums for samples at the node
    parent : TreeNode or None, optional(default=None)
        The parent of the node. None for root.
    split_info : SplitInfo or None
        The result of the split evaluation
    left_child : TreeNode or None
        The left child of the node. None for leaves.
    right_child : TreeNode or None
        The right child of the node. None for leaves.
    value : float or None
        The value of the leaf, as computed in finalize_leaf(). None for
        non-leaf nodes
    find_split_time : float
        The total time spent computing the histogram and finding the best
        split at the node.
    construction_speed : float
        The Number of samples at the node divided find_split_time.
    apply_split_time : float
        The total time spent actually splitting the node, e.g. splitting
        samples_indices into left and right child.
    hist_subtraction : bool
        Wheter the subtraction method was used for computing the histograms.
    """

    split_info = None
    left_child = None
    right_child = None
    value = None
    left_coefficient=None
    right_coefficient=None
    histograms = None
    sibling = None
    parent = None
    find_split_time = 0.
    construction_speed = 0.
    apply_split_time = 0.
    hist_subtraction = False

    def __init__(self, depth, sample_indices, sum_g_hf, sum_gx_hfx, sum_h, sum_hx, sum_hx2, parent=None):
        self.depth = depth
        self.sample_indices = sample_indices
        self.n_samples = sample_indices.shape[0]
        self.sum_g_hf = sum_g_hf
        self.sum_gx_hfx = sum_gx_hfx
        self.sum_h = sum_h
        self.sum_hx = sum_hx
        self.sum_hx2 = sum_hx2
        self.parent = parent

    def __repr__(self):
        # To help with debugging
        out = f"TreeNode: depth={self.depth}, "
        out += f"samples={len(self.sample_indices)}"
        if self.split_info is not None:
            out += f", feature_idx={self.split_info.feature_idx}"
            out += f", bin_idx={self.split_info.bin_idx}"
        return out

    def __lt__(self, other_node):
        """Comparison for priority queue.

        Nodes with high gain are higher priority than nodes with low gain.

        heapq.heappush only need the '<' operator.
        heapq.heappop take the smallest item first (smaller is higher
        priority).

        Parameters
        -----------
        other_node : TreeNode
            The node to compare with.
        """
        if self.split_info is None or other_node.split_info is None:
            raise ValueError("Cannot compare nodes with split_info")
        return self.split_info.gain > other_node.split_info.gain

class TreeGrower:
    """Tree grower class used to build a tree.

    The tree is fitted to predict the values of a Newton-Raphson step. The
    splits are considered in a best-first fashion, and the quality of a
    split is defined in splitting._split_gain.

    Parameters
    ----------
    X_binned : array-like of int, shape=(n_samples, n_features)
        The binned input samples. Must be Fortran-aligned.
    gradients : array-like, shape=(n_samples,)
        The gradients of each training sample. Those are the gradients of the
        loss w.r.t the predictions, evaluated at iteration ``i - 1``.
    hessians : array-like, shape=(n_samples,)
        The hessians of each training sample. Those are the hessians of the
        loss w.r.t the predictions, evaluated at iteration ``i - 1``.
    max_leaf_nodes : int or None, optional(default=None)
        The maximum number of leaves for each tree. If None, there is no
        maximum limit.
    max_depth : int or None, optional(default=None)
        The maximum depth of each tree. The depth of a tree is the number of
        nodes to go from the root to the deepest leaf.
    min_samples_leaf : int, optional(default=20)
        The minimum number of samples per leaf.
    min_gain_to_split : float, optional(default=0.)
        The minimum gain needed to split a node. Splits with lower gain will
        be ignored.
    max_bins : int, optional(default=256)
        The maximum number of bins. Used to define the shape of the
        histograms.
    n_bins_per_feature : array-like of int or int, optional(default=None)
        The actual number of bins needed for each feature, which is lower or
        equal to ``max_bins``. If it's an int, all features are considered to
        have the same number of bins. If None, all features are considered to
        have ``max_bins`` bins.
    l2_regularization : float, optional(default=0)
        The L2 regularization parameter.
    min_det_to_split : float, optional(default=1e-3)
        The minimum sum of hessians needed in each node. Splits that result in
        at least one child having a sum of hessians less than
        min_determinant_to_split are discarded.
    shrinkage : float, optional(default=1)
        The shrinkage parameter to apply to the leaves values, also known as
        learning rate.
    """
    def __init__(self, dataset, gradients, hessians, options: OptionSet):

        self.max_leaf_nodes = options['max_leaf_nodes']
        self.max_depth = options['max_depth']
        self.min_samples_leaf = options['min_samples_leaf']
        self.min_gain_to_split = options['min_gain_to_split']
        self.max_bins = options['max_bins']
        self.n_bins_per_feature = dataset.n_bins_per_feature
        self.numerical_thresholds = dataset.numerical_thresholds
        self.w_l2_reg = options['w_l2_reg']
        self.b_l2_reg = options['b_l2_reg']
        self.min_hessian_to_split = options['min_hessian_to_split']
        self.shrinkage = options['learning_rate']

        self._validate_parameters(dataset.X_binned, dataset.X, self.max_leaf_nodes, self.max_depth,
                                  self.min_samples_leaf, self.min_gain_to_split,
                                  self.w_l2_reg, self.b_l2_reg, self.min_hessian_to_split)

        if self.n_bins_per_feature is None:
            self.n_bins_per_feature = self.max_bins

        if isinstance(self.n_bins_per_feature, int):
            self.n_bins_per_feature = np.array(
                [self.n_bins_per_feature] * dataset.X_binned.shape[1],
                dtype=np.uint32)

        self.splitting_context = SplittingContext(
            dataset.X_binned, dataset.X, self.max_bins, self.n_bins_per_feature, gradients,
            hessians, self.w_l2_reg, self.b_l2_reg, self.min_hessian_to_split,
            self.min_samples_leaf, self.min_gain_to_split)
        self.X_binned = dataset.X_binned
        self.splittable_nodes = []
        self.finalized_leaves = []
        self.total_find_split_time = 0.  # time spent finding the best splits
        self.total_apply_split_time = 0.  # time spent splitting nodes
        self._initialize_root()
        self.n_nodes = 1

    def _validate_parameters(self, X_binned, X, max_leaf_nodes, max_depth,
                             min_samples_leaf, min_gain_to_split,
                             w_l2_reg, b_l2_reg, min_hessian_to_split):
        """Validate parameters passed to __init__.

        Also validate parameters passed to SplittingContext because we cannot
        raise exceptions in a jitclass.
        """
        if X_binned.dtype != np.uint8:
            raise NotImplementedError(
                "Explicit feature binning required for now")
        if not X_binned.flags.f_contiguous:
            raise ValueError(
                "X_binned should be passed as Fortran contiguous "
                "array for maximum efficiency.")
        if not X.flags.f_contiguous:
            raise ValueError(
                "X should be passed as Fortran contiguous "
                "array for maximum efficiency.")
        if max_leaf_nodes is not None and max_leaf_nodes < 1:
            raise ValueError(f'max_leaf_nodes={max_leaf_nodes} should not be'
                             f' smaller than 1')
        if max_depth is not None and max_depth < 1:
            raise ValueError(f'max_depth={max_depth} should not be'
                             f' smaller than 1')
        if min_samples_leaf < 1:
            raise ValueError(f'min_samples_leaf={min_samples_leaf} should '
                             f'not be smaller than 1')
        if min_gain_to_split < 0:
            raise ValueError(f'min_gain_to_split={min_gain_to_split} '
                             f'must be positive.')
        if w_l2_reg < 0 or b_l2_reg < 0:
            raise ValueError(f'l2_regularization=w:{w_l2_reg},b:{b_l2_reg} must be '
                             f'positive.')
        if min_hessian_to_split < 0:
            raise ValueError(f'min_hessian_to_split={min_hessian_to_split} '
                             f'must be positive.')

    def grow(self):
        """Grow the tree, from root to leaves."""
        while self.can_split_further():
            self.split_next()

    def _initialize_root(self):
        """Initialize root node and finalize it if needed."""
        n_samples = self.X_binned.shape[0]
        depth = 0
        self.root = TreeNode(
            depth=depth,
            sample_indices=self.splitting_context.partition.view(),
            sum_g_hf=self.splitting_context.gradients.sum(),
            sum_h=self.splitting_context.hessians.sum(),
            sum_gx_hfx=0,
            sum_hx=0,
            sum_hx2=0
        )
        if (self.max_leaf_nodes is not None and self.max_leaf_nodes == 1):
            self._finalize_leaf(self.root)
            return
        if self.root.n_samples < 2 * self.min_samples_leaf:
            # Do not even bother computing any splitting statistics.
            self._finalize_leaf(self.root)
            return

        self._compute_root_spittability()

    def _compute_root_spittability(self):
        """Compute histograms and best possible split of a root node
        """
        node = self.root
        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, node.sample_indices)

        toc = time()
        node.find_split_time = toc - tic
        self.total_find_split_time += node.find_split_time
        node.construction_speed = node.n_samples / node.find_split_time
        node.split_info = split_info
        node.histograms = histograms

        if node.split_info.gain <= 0:  # no valid split
#            import pdb; pdb.set_trace()
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(node)

        else:
            heappush(self.splittable_nodes, node)

    def _compute_sibling_splittability(self, left_child_node, right_child_node):
        if left_child_node.n_samples > right_child_node.n_samples:
            min_child_node, max_child_node = right_child_node, left_child_node
        else:
            min_child_node, max_child_node = left_child_node, right_child_node

        if max_child_node.n_samples < 2 * self.min_samples_leaf:
            self._finalize_leaf(max_child_node)
            self._finalize_leaf(min_child_node)
            return left_child_node, right_child_node

        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, min_child_node.sample_indices)
        toc = time()
        min_child_node.find_split_time = toc - tic
        self.total_find_split_time += min_child_node.find_split_time
        min_child_node.construction_speed = min_child_node.n_samples / min_child_node.find_split_time
        min_child_node.split_info = split_info
        min_child_node.histograms = histograms

        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, max_child_node.sample_indices)
        toc = time()
        max_child_node.find_split_time = toc - tic
        self.total_find_split_time += max_child_node.find_split_time
        max_child_node.construction_speed = max_child_node.n_samples / max_child_node.find_split_time
        max_child_node.split_info = split_info
        max_child_node.histograms = histograms

        if max_child_node.split_info.gain <= 0:  # no valid split
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(max_child_node)

        else:
            heappush(self.splittable_nodes, max_child_node)

        if min_child_node.split_info.gain <= 0:  # no valid split
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(min_child_node)

        else:
            heappush(self.splittable_nodes, min_child_node)

        return left_child_node, right_child_node

    def split_next(self):
        """Split the node with highest potential gain.

        Returns
        -------
        left : TreeNode
            The resulting left child.
        right : TreeNode
            The resulting right child.
        """
        if len(self.splittable_nodes) == 0:
            raise StopIteration("No more splittable nodes")

        # Consider the node with the highest loss reduction (a.k.a. gain)
        node = heappop(self.splittable_nodes)

        tic = time()
        (sample_indices_left, sample_indices_right) = split_indices(
            self.splitting_context, node.split_info, node.sample_indices)
        toc = time()
        node.apply_split_time = toc - tic
        self.total_apply_split_time += node.apply_split_time

        depth = node.depth + 1
        n_leaf_nodes = len(self.finalized_leaves) + len(self.splittable_nodes)
        n_leaf_nodes += 2

        left_child_node = TreeNode(depth,
                                   sample_indices_left,
                                   node.split_info.left_g_hf,
                                   node.split_info.left_gx_hfx,
                                   node.split_info.left_h,
                                   node.split_info.left_hx,
                                   node.split_info.left_hx2,
                                   parent=node)
        right_child_node = TreeNode(depth,
                                    sample_indices_right,
                                    node.split_info.right_g_hf,
                                    node.split_info.right_gx_hfx,
                                    node.split_info.right_h,
                                    node.split_info.right_hx,
                                    node.split_info.right_hx2,
                                    parent=node)
        left_child_node.sibling = right_child_node
        right_child_node.sibling = left_child_node
        node.right_child = right_child_node
        node.left_child = left_child_node
        self.n_nodes += 2

        node.left_coefficient = self._compute_linear_coefficient(left_child_node)
        node.right_coefficient = self._compute_linear_coefficient(right_child_node)

        update_prediction_values(
            self.splitting_context, sample_indices_left,
            node.left_coefficient, node.split_info.feature_idx
        )
        update_prediction_values(
            self.splitting_context, sample_indices_right,
            node.right_coefficient, node.split_info.feature_idx
        )

        if (self.max_leaf_nodes is not None
            and n_leaf_nodes == self.max_leaf_nodes):
            self._finalize_leaf(left_child_node)
            self._finalize_leaf(right_child_node)
            self._finalize_splittable_nodes()
            return left_child_node, right_child_node

        if self.max_depth is not None and depth == self.max_depth:
            self._finalize_leaf(left_child_node)
            self._finalize_leaf(right_child_node)
            return left_child_node, right_child_node

        self._compute_sibling_splittability(left_child_node, right_child_node)

        return left_child_node, right_child_node

    def can_split_further(self):
        """Return True if there are still nodes to split."""
        return len(self.splittable_nodes) >= 1

    def _compute_linear_coefficient(self, node):

        sum_hx2_reg = node.sum_hx2 + self.splitting_context.w_l2_reg
        sum_h_reg = node.sum_h + self.splitting_context.b_l2_reg
        coefficient = (node.sum_hx * node.sum_g_hf - sum_h_reg * node.sum_gx_hfx) / \
                      (sum_hx2_reg * sum_h_reg - node.sum_hx ** 2)

        return coefficient

    def _finalize_leaf(self, node):
        if node.parent is None:
            self._finalize_leaf_plain(node)
        else:
            self._finalize_leaf_linear(node)

    def _finalize_leaf_linear(self, node):
        sum_hx2_reg = node.sum_hx2 + self.splitting_context.w_l2_reg
        sum_h_reg = node.sum_h + self.splitting_context.b_l2_reg
        node.value = self.shrinkage * (node.sum_hx * node.sum_gx_hfx - sum_hx2_reg * node.sum_g_hf) / (
            sum_hx2_reg * sum_h_reg - node.sum_hx**2)
        self.finalized_leaves.append(node)

    def _finalize_leaf_plain(self, node):
        """Compute the prediction value that minimizes the objective function.

        This sets the node.value attribute (node is a leaf iff node.value is
        not None).

        See Equation 5 of:
        XGBoost: A Scalable Tree Boosting System, T. Chen, C. Guestrin, 2016
        https://arxiv.org/abs/1603.02754
        """
        node.value = -self.shrinkage * node.sum_g_hf / (
            node.sum_h + self.splitting_context.b_l2_reg)
        self.finalized_leaves.append(node)

    def _finalize_splittable_nodes(self):
        """Transform all splittable nodes into leaves.

        Used when some constraint is met e.g. maximum number of leaves or
        maximum depth."""
        while len(self.splittable_nodes) > 0:
            node = self.splittable_nodes.pop()
            self._finalize_leaf(node)

    def make_predictor(self):
        """Make a TreePredictor object out of the current tree.

        Parameters
        ----------
        numerical_thresholds : array-like of floats, optional (default=None)
            The actual thresholds values of each bin, expected to be in sorted
            increasing order. None if the training data was pre-binned.

        Returns
        -------
        A TreePredictor object.
        """
        predictor_nodes = np.zeros(self.n_nodes, dtype=PREDICTOR_RECORD_DTYPE)
        self._fill_predictor_node_array(
            predictor_nodes, self.root,
            numerical_thresholds=self.numerical_thresholds
        )
        return TreePredictor(nodes=predictor_nodes)

    def _fill_predictor_node_array(self, predictor_nodes, grower_node,
                                   numerical_thresholds=None, next_free_idx=0):
        """Helper used in make_predictor to set the TreePredictor fields."""
        node = predictor_nodes[next_free_idx]
        node['count'] = grower_node.n_samples
        node['depth'] = grower_node.depth
        if grower_node.split_info is not None:
            node['gain'] = grower_node.split_info.gain
        else:
            node['gain'] = -1

        if grower_node.value is not None:
            # Leaf node
            node['is_leaf'] = True
            node['value'] = grower_node.value
            return next_free_idx + 1
        else:
            # Decision node
            split_info = grower_node.split_info
            feature_idx, bin_idx = split_info.feature_idx, split_info.bin_idx
            node['feature_idx'] = feature_idx
            node['bin_threshold'] = bin_idx
            if numerical_thresholds is not None:
                node['threshold'] = numerical_thresholds[feature_idx][bin_idx]
            node['left_coefficient'] = grower_node.left_coefficient * self.shrinkage
            node['right_coefficient'] = grower_node.right_coefficient * self.shrinkage
            next_free_idx += 1

            node['left'] = next_free_idx
            next_free_idx = self._fill_predictor_node_array(
                predictor_nodes, grower_node.left_child,
                numerical_thresholds=numerical_thresholds,
                next_free_idx=next_free_idx)

            node['right'] = next_free_idx
            return self._fill_predictor_node_array(
                predictor_nodes, grower_node.right_child,
                numerical_thresholds=numerical_thresholds,
                next_free_idx=next_free_idx)

    def update_raw_predictions(self, raw_predictions):
        # prepare leaves_data so that _update_raw_predictions can be
        # @njitted
        leaves_data = [(l.value, l.sample_indices)
                        for l in self.finalized_leaves]
        prediction_value = self.splitting_context.prediction_value * self.shrinkage
        _update_raw_predictions(leaves_data, raw_predictions, prediction_value)

@njit(parallel=True)
def _update_raw_predictions(leaves_data, raw_predictions, prediction_value):
    """Update raw_predictions by reading the predictions of the ith tree
    directly form the leaves.

    Can only be used for predicting the training data. raw_predictions
    contains the sum of the tree values from iteration 0 to i - 1. This adds
    the predictions of the ith tree to raw_predictions.

    Parameters
    ----------
    leaves_data: list of tuples (leaf.value, leaf.sample_indices)
        The leaves data used to update raw_predictions.
    raw_predictions : array-like, shape=(n_samples,)
        The raw predictions for the training data.
    prediction_value:array-like, shape=(n_samples,)
        The accumulated linear values
    """
    for leaf_idx in prange(len(leaves_data)):
        leaf_value, sample_indices = leaves_data[leaf_idx]
        for sample_idx in sample_indices:
            raw_predictions[sample_idx] += leaf_value
    n_samples = raw_predictions.shape[0]
    for i in prange(n_samples):
        raw_predictions[i] += prediction_value[i]


from styx_msgs.msg import TrafficLight
import yaml
import os
import rospy
import cv2
import tensorflow as tf
import numpy as np

IMAGE_H = 300
IMAGE_W = 300

class TLClassifier(object):
    def __init__(self):
        config_string = rospy.get_param("/traffic_light_config")
        self.config = yaml.load(config_string)

        # Tf graph and sess
        self.model_graph = None
        self.session = None

        self.load_graph(os.path.dirname(os.path.realpath(__file__)) + '/' + self.config['model'])

        self.classes = {1: TrafficLight.RED,
                2: TrafficLight.YELLOW,
                3: TrafficLight.GREEN,
                4: TrafficLight.UNKNOWN}

    def get_classification(self, image):
        """Determines the color of the traffic light in the image
        Args:
            image (cv::Mat): image containing the traffic light
        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)
        """
        class_index, probability = self.predict(image)

        if class_index is not None:
            return class_index
        else:
            # If we detected nothing, just return none
            return TrafficLight.UNKNOWN

    def load_graph(self, model_path):
        """Loads a frozen inference graph"""
        # Create config
        config = tf.ConfigProto()
        config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

        # Create model
        self.model_graph = tf.Graph()

        # Create session
        with tf.Session(graph=self.model_graph, config=config) as sess:
            self.session = sess
            graph_def = tf.GraphDef()

            # Load model
            with tf.gfile.GFile(model_path, 'rb') as fid:
                serialized_graph = fid.read()
                graph_def.ParseFromString(serialized_graph)

                # Import graph definition
                tf.import_graph_def(graph_def, name='')

    def process_image(self, image):
        # Resize image to fit training size
        image = cv2.resize(image, (IMAGE_W, IMAGE_H))

        # Convert from BGR to RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image

    def predict(self, image):
        # The input placeholder for the image.
        # `get_tensor_by_name` returns the Tensor with the associated name in the Graph.
        image_tensor = self.model_graph.get_tensor_by_name('image_tensor:0')

        # Each box represents a part of the image where a particular object was detected.
        detection_boxes = self.model_graph.get_tensor_by_name('detection_boxes:0')

        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        detection_scores = self.model_graph.get_tensor_by_name('detection_scores:0')

        # The classification of the object (integer id).
        detection_classes = self.model_graph.get_tensor_by_name('detection_classes:0')

        image = self.process_image(image)

        (boxes, scores, classes) = self.session.run(
            [detection_boxes, detection_scores, detection_classes],
            feed_dict={image_tensor: np.expand_dims(image, axis=0)})

        # Remove unnecessary dimensions
        scores = np.squeeze(scores)
        classes = np.squeeze(classes)
        boxes = np.squeeze(boxes)

        min_score = 0.5

        # For each bounding box
        for i, box in enumerate(boxes):

            # If detection prob is higher than minimum
            if scores[i] > min_score:

                # Determin class of each detection
                detected_class = self.classes[classes[i]]
                return detected_class, scores[i]

        return None, None