from skimage import filters, exposure
from skimage.color import label2rgb
from skimage.morphology import watershed
from skimage.feature import peak_local_max
from scipy import ndimage as nd
from skimage import img_as_ubyte, img_as_float
from skimage.feature import hessian_matrix, hessian_matrix_eigvals, canny
import imageio as io
import cv2
import os
import numpy as np
import pandas as pd
import math
from matplotlib import pyplot as plt
import czifile

def get_file_extension(file_path):
    file_ext = os.path.splitext(file_path)

    return file_ext[1]  # because splitext returns a tuple and the extension is the second element

def find_image_channel_name(file_name):
    str_idx = file_name.find('Conf ')  # this is specific to our microscopes file name format
    channel_name = file_name[str_idx + 5 : str_idx + 8]

    return channel_name

def max_project(img):
    projection = np.max(img, axis=0)

    return projection

def clear_axis_ticks(ax):
    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])

def make_color_image(img, c):
    output = np.zeros(shape=(img.shape[0], img.shape[0], 3))

    if c == 'green':
        output[..., 0] = 0.0  # R
        output[..., 1] = img  # G
        output[..., 2] = 0.0  # B
    elif c == 'magenta':
        output[..., 0] = img  # R
        output[..., 1] = 0.0  # G
        output[..., 2] = img  # B
    elif c == 'cyan':
        output[..., 0] = 0.0  # R
        output[..., 1] = img  # G
        output[..., 2] = img  # B

    else:
        print('ERROR: Could not identify color to pseudocolor image in grapher.make_color_image')
        sys.exit(0)

    return output

def find_region_area(r):
    a = ((r[0].stop - r[0].start)) * ((r[1].stop - r[1].start))

    return a

def find_region_volume(r):
    a = (r[0].stop - r[0].start) * (r[1].stop - r[1].start) * (r[2].stop - r[2].start)

    return a

# load images for replicate
def load_images(data, input_params):

    with czifile.CziFile(data.img_path) as czi:
        img = czi.asarray()
        img = np.squeeze(img)  # this gets rid of the extra dimensions if they are 1 (e.g. if it isn't a timecourse)
        # metadata = czi.metadata  # if you want metadata. But I don't currently parse this.

    # NUCLEUS IMAGE
    if len(img.shape) < 4:
        return None

    else:
        data.nuc_img =  img[input_params.nuc_idx, :, :, :]# image is [z, x, y] array

        '''
        if data.nuc_img.dtype == 'float32':
            data.nuc_img = data.nuc_img.astype(np.uint16)  # my pipeline to align image stacks keeps values but in 32-bit format. This conflicts with later processing steps so we have to trick it.
        '''

        # PROTEIN IMAGES
        data.pro_imgs = [img[input_params.pro_idx, :, : ,:]]
        data.pro_ch_names = ['ch488']

        '''
        for idx, p in enumerate(protein_image_files):
            data.protein_image_paths.append(os.path.join(input_params.parent_path, data.folder, p))
            data.protein_channel_names.append(find_image_channel_name(p))
            data.pro_imgs.append(io.volread(data.protein_image_paths[idx]))

            if data.pro_imgs[idx].dtype == 'float32':
                data.pro_imgs[idx] = data.pro_imgs[idx].astype(np.uint16) # my pipeline to align image stacks keeps values but in 32-bit format. This conflicts with later processing steps so we have to trick it.
        '''

        return data

def find_nucleus_2D(data, input_params):

    img = max_project(data.nuc_img)
#     med_img = img_as_ubyte(img)
#     med_img = cv2.medianBlur(med_img, ksize=5)
#     med_img = img_as_float(med_img)

    med_img = filters.gaussian(img, sigma=2)
    # threshold = np.mean(med_img) + 0.05*np.std(med_img)
    threshold = 1000/65536

    # if input_params.threshold is None:
    #     threshold = filters.threshold_otsu(med_img)
#     else:
#         threshold = input_params.threshold

    # nuc_mask = canny(med_img, sigma=1)

    nuc_mask = med_img >= threshold

    nuc_mask = nd.morphology.binary_opening(nuc_mask)
    # nuc_mask = nd.morphology.binary_dilation(nuc_mask)
    nuc_mask = nd.morphology.binary_fill_holes(nuc_mask)
    #nuclear_label, num_features = nd.label(nuclear_mask)
    #nuclear_label = find_watershed_3D(nuclear_mask)

    ## WATERSHEDDING ##
    if False:
        nuc_label = find_watershed_2D(nuc_mask)
        nuc_mask = nuc_label >= 1

    nuc_label, _ = nd.label(nuc_mask)

    data.nuc_mask = nuc_mask
    data.nuc_label = nuc_label

    data.nuc_regions = nd.find_objects(nuc_label)

    return data

def find_watershed_2D(z_slice):
    labels, _ = nd.label(z_slice)
    distance = nd.distance_transform_edt(z_slice)
    local_maxi = peak_local_max(distance, min_distance=50, indices=False, labels=labels, exclude_border=True)
    markers = nd.label(local_maxi)[0]

    output = watershed(-distance, markers, mask=z_slice, watershed_line=True)

    return output

def make_nucleus_montage_2D(data, input_params):

    fig, ax = plt.subplots(1, 3)

    disp_nuc_img = max_project(data.nuc_img)
    disp_nuc_img = exposure.equalize_adapthist(disp_nuc_img)
    disp_nuc_img = make_color_image(disp_nuc_img, 'cyan')

    ax[0].imshow(disp_nuc_img)
    ax[1].imshow(data.nuc_mask, cmap='gray')

    labeled_image = label2rgb(data.nuc_label, image=disp_nuc_img,
                             alpha=0.5, bg_label=0, bg_color=[0, 0, 0])

    ax[2].imshow(labeled_image)

    for a in ax:
        clear_axis_ticks(a)

    plt.tight_layout()
    plt.show()

def find_dense_objects_3D(img, mask, input_params, data):
    # bg_mean is uint16
    #img[np.invert(mask)] = bg_mean

    img = img_as_float(img)

    ''''
    #subtract background using median filter
    med_img = img_as_ubyte(img)
    for i in range(med_img.shape[0]):
        med_img[i, :, :] = cv2.medianBlur(med_img[i,:, :], ksize=13)

    med_img = img_as_float(med_img)
    img = img - med_img
    img[np.where(img < 0)] = 0
    '''

    # just do simple gaussian filter and thresholding (z plotting)

    img = filters.gaussian(img, sigma=3)
    threshold = np.mean(img) + 2.35*np.std(img)

    object_mask = np.logical_and(mask, img >= threshold)
    # object_mask = nd.morphology.binary_opening(object_mask)
    object_mask = nd.morphology.binary_opening(object_mask, iterations=3) # we are trying to see if you can open until nothing changes
    bg_nuclear_mask = np.logical_and(mask, img < threshold)

    #parameters
    fg_dist_threshold = 0.2
    struct_size = 3
    erosion_iterations = 1
    distance = nd.distance_transform_edt(object_mask)
    z_size = img.shape[0]

    if False:
        ### WATERSHED TEST
            labels, _ = nd.label(object_mask)

            distance = nd.distance_transform_edt(object_mask)

            ##custom faster solution for getting markers
            sure_fg = distance
            sure_fg[sure_fg <= fg_dist_threshold*distance.max()] = 0.0
            sure_fg = sure_fg > 0.0

            ci_struct = make_struct_element(struct_size, shape='circle', dim=2).astype(img.dtype)
            ew_struct = make_struct_element(struct_size, shape='ellipse_wide', dim=2).astype(img.dtype)
            et_struct = make_struct_element(struct_size, shape='ellipse_tall', dim=2).astype(img.dtype)

            for z in range(z_size):
                sure_fg[z, :, :] = nd.morphology.binary_erosion(sure_fg[z, :, :], structure=ci_struct, iterations=erosion_iterations)
                # sure_fg[z, :, :] = nd.morphology.binary_erosion(sure_fg[z, :, :], structure=ew_struct, iterations=erosion_iterations)
                # sure_fg[z, :, :] = nd.morphology.binary_erosion(sure_fg[z, :, :], structure=et_struct, iterations=erosion_iterations)

            markers, num_regions = nd.label(sure_fg)
            row, col = optimum_subplots(object_mask.shape[0])
            fig, ax = plt.subplots(row, col)

            ax = ax.flatten()

            for idx, a in enumerate(ax):
                if idx < z_size:
                    labeled_image = label2rgb(markers[idx, :, :], image=exposure.equalize_adapthist(img[idx, :, :]),
                                     alpha=0.3, bg_label=0, bg_color=[0, 0, 0])
                    a.imshow(labeled_image)

                clear_axis_ticks(a)

            plt.savefig(os.path.join(input_params.output_path, data.sample_name + '_watershed_test.png'), dpi=150)
            plt.close()

    ### WATERSHED

    if True:
        ##custom faster solution for getting markers
        sure_fg = distance
        sure_fg[sure_fg <= fg_dist_threshold*distance.max()] = 0
        sure_fg = sure_fg > 0

        ci_struct = make_struct_element(struct_size, shape='circle', dim=2).astype(img.dtype)

        for z in range(z_size):
            sure_fg[z, :, :] = nd.morphology.binary_erosion(sure_fg[z, :, :], structure=ci_struct, iterations=erosion_iterations)

        markers, num_regions = nd.label(sure_fg)
        dense_object_label = watershed(-distance, markers, mask=object_mask, watershed_line=True)
        object_mask = dense_object_label > 0

        # dense_object_label, _ = nd.label(object_mask)
        dense_objects = nd.find_objects(dense_object_label)

        # new_object_mask = np.full(shape=img.shape, fill_value=False, dtype=bool)
        # filtered_objects = []

    vol_threshold = 500
    for idx, object in enumerate(dense_objects):
        if find_region_volume(object) < vol_threshold:
            object_mask[object] = 0
            dense_objects[idx] = []

    dense_objects = [o for o in dense_objects if o]

    if False:
        under_img = exposure.equalize_adapthist(img[z, :, :])
        labeled_img = label2rgb(object_mask[z,:,:], image=under_img,
                           alpha=0.5, bg_label=0, bg_color=[0,0,0])

        labeled_bg_img = label2rgb(bg_nuclear_mask[z,:,:], image=under_img,
                           alpha=0.5, bg_label=0, bg_color=[0,0,0])

        fig, ax = plt.subplots(1, 3)
        ax[0].imshow(img[z, :, :], cmap='magma')
        ax[1].imshow(labeled_img)
        ax[2].imshow(labeled_bg_img)

    '''
    # just do simple gaussian filter and thresholding (max_z plotting)
    if False:
        img = filters.gaussian(img, sigma=2)

        threshold = np.mean(img) + 2*np.std(img)

        object_mask = np.logical_and(mask, img >= threshold)
        object_mask = nd.morphology.binary_opening(object_mask)

        under_img = exposure.equalize_adapthist(max_project(img))
        labeled_img = label2rgb(max_project(object_mask), image=under_img,
                               alpha=0.5, bg_label=0, bg_color=[0,0,0])

        bg_nuclear_mask = np.logical_and(mask, img < threshold)
        labeled_bg_img = label2rgb(max_project(bg_nuclear_mask), image=under_img,
                               alpha=0.5, bg_label=0, bg_color=[0,0,0])
        if False:
            fig, ax = plt.subplots(1, 3)

            ax[0].imshow(max_project(img), cmap='magma')
            ax[1].imshow(labeled_img)
            ax[2].imshow(labeled_bg_img)

            for a in ax:
                clear_axis_ticks(a)

            plt.show()
    '''

    '''

    threshold = filters.threshold_otsu(img)
    object_mask = img >= threshold
    object_mask = nd.morphology.binary_opening(object_mask)

    labeled_image = label2rgb(max_project(object_mask), image=max_project(img),
                         alpha=0.5, bg_label=0, bg_color=[0, 0, 0])
    '''
    return object_mask, bg_nuclear_mask, dense_objects

def make_output_graphs(nuc_label, obj_mask, data, output_path):
    num_of_proteins = len(data.pro_imgs)
    fig_h = 3.3  # empirical
    fig_w = 4.23 * num_of_proteins # empirical

    fig, ax = plt.subplots(1, 2+num_of_proteins)

    # 1st image is nuclear mask
    nuc_under_img = exposure.equalize_adapthist(max_project(data.nuc_img))
    nuc_labeled_img = label2rgb(nuc_label, image=nuc_under_img,
                                alpha=0.5, bg_label=0, bg_color=[0, 0 ,0])
    ax[0].imshow(nuc_labeled_img)
    ax[0].set_title('nuclear mask', fontsize=10)

    for r_idx, region in enumerate(data.nuc_regions):
            region_area = find_region_area(region)
            if region_area >= 30000:
                region_center_r = int((region[0].stop + region[0].start)/2)
                region_center_c = int((region[1].stop + region[1].start)/2)
                nuc_id = data.nuc_label[region_center_r, region_center_c]
                ax[0].text(region_center_c, region_center_r, str(nuc_id),
                           fontsize='6', color='w', horizontalalignment='center', verticalalignment='center')

    # 2nd image is total object mask
    object_labeled_img = label2rgb(max_project(obj_mask), image=nuc_under_img,
                                  alpha=0.5, bg_label=0, bg_color=[0, 0, 0])
    ax[1].imshow(object_labeled_img)
    ax[1].set_title('object mask', fontsize=10)

    # the rest are the protein images with boxes around

    for p_idx, img in enumerate(data.pro_imgs):
        img = img_as_float(img)
        temp_under_img = exposure.equalize_adapthist(max_project(img))
        temp_labeled_img = label2rgb(max_project(obj_mask), image=temp_under_img,
                               alpha=0.25, bg_label=0, bg_color=[0,0,0])
        ax[2+p_idx].imshow(temp_labeled_img)
        ax[2+p_idx].set_title(str(data.pro_ch_names[p_idx]))

    for a in ax:
        clear_axis_ticks(a)

    plt.subplots_adjust(left=0.01, right=0.99, bottom=0.01, top=0.90, wspace=0.1, hspace=0.1)

    plt.savefig(output_path,dpi=300)
    plt.close()

def make_struct_element(edge_size, shape='circle', dim=3):

    if edge_size % 2 == 0:
        print('Error: Structuring element must have an odd edge size')
        sys.exit(0)
    else:
        if shape != 'circle':
            edge_size = edge_size + 2

        X, Y = np.ogrid[0:edge_size, 0:edge_size]
        center = math.floor(edge_size/2)
        scale_factor = 1./math.sqrt(edge_size)

        if shape == 'circle':
            element = ((X-center)**2 + (Y-center)**2 < edge_size).astype(np.uint8)
        elif shape == 'ellipse_tall':
            element = (scale_factor * (X-center)**2 + (Y-center)**2 < edge_size).astype(np.uint8)
        elif shape == 'ellipse_wide':
            element = ((X-center)**2 + scale_factor * (Y-center)**2 < edge_size).astype(np.uint8)
        else:
            print('Error: Could not recognize shape input for making structuring element')
            sys.exit(0)

        if dim == 3:
            # for now, we will just make the z 1 unit thick
            element = element[np.newaxis,:, :]

    '''
    ## STRUCT ELEMENT TEST
    print(f'Shape is {shape}')
    print(element)
    print()
    '''

    return element

def optimum_subplots(n):
    row = math.floor(math.sqrt(n))
    col = math.ceil(n/row)

    return row, col


import torch
import numpy as np
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim as optim
import utils

def _concat(xs):
  return torch.cat([x.view(-1) for x in xs])

class Augment(object):

  def __init__(self, target_net, aug_net, config):

    # loss function
    self.criterion = nn.CrossEntropyLoss().cuda()

    # optimizer
    # print('target net lr: {}'.format(config.lr))
    # print('target optimizer momentum: {}'.format(config.momentum))
    # print('target optimizer weight decay: {}'.format(config.weight_decay))
    # if config.decay_type is None:
    #   params = target_net.parameters()
    # elif config.decay_type == 'no_bn':
    #   params = utils.add_weight_decay(target_net, config.weight_decay)
    # else:
    #   raise Exception('unknown decay type: {}'.format(config.decay_type))
    self.target_net_optim = optim.SGD(target_net.parameters(), config.lr,
                                      momentum=config.momentum,
                                      weight_decay=config.weight_decay,
                                      nesterov=True)
    for group in self.target_net_optim.param_groups:
      print('target net lr: {}, weight_decay: {}, momentum: {}, nesterov: {}'
            .format(group['lr'], group['weight_decay'], group['momentum'], group['nesterov']))

    print('training epochs: {}'.format(config.epochs))
    # lr scheduler
    if config.lr_scheduler == 'cosine':
        self.lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.target_net_optim,
                                                                       T_max=float(config.epochs),
                                                                       eta_min=0.)
    else:
        raise ValueError('invalid lr_schduler: {}'.format(config.lr_scheduler))

    # self.network_momentum = args.momentum
    self.target_net = target_net
    self.aug_net = aug_net
    # self.criterion = criterion
    self.args = config
    print('adv weight for deformation: {}'.format(self.args.adv_weight_deform))
    print('reconstruction weight for deformation: {}'.format(self.args.div_weight_deform))
    print('smooth weight: {}'.format(self.args.smooth_weight))
    assert self.args.adv_weight_deform >= 0
    assert self.args.div_weight_deform >= 0
    assert self.args.smooth_weight >= 0
    print('aug_net_lr: {}'.format(config.aug_net_lr))
    print('aug net adam optimizer beta1: {}'.format(config.adam_beta1))
    print('aug_net_weight_decay: {}'.format(config.aug_net_weight_decay))
    self.aug_net_optim = torch.optim.Adam(aug_net.parameters(),
                                          lr=config.aug_net_lr, betas=(config.adam_beta1, 0.999),
                                          weight_decay=config.aug_net_weight_decay)

  # def call_aug_net(self, data, label, require_loss=False):
  #   if require_loss:
  #     if 'vae' in self.args.exp_type and self.args.vae_type == 'condition':
  #       # print('using conditional vae ...')
  #       # exit()
  #       data_aug, loss = self.aug_net(data, F.one_hot(label, 10).type(torch.cuda.FloatTensor), require_loss=True)
  #     else:
  #       data_aug, loss = self.aug_net(data, require_loss=True)
  #
  #     return data_aug, loss
  #   else:
  #     if 'vae' in self.args.exp_type and self.args.vae_type == 'condition':
  #       data_aug = self.aug_net(data, F.one_hot(label, 10).type(torch.cuda.FloatTensor))
  #     else:
  #       data_aug = self.aug_net(data)
  #
  #     return data_aug

  def _compute_unrolled_model(self, loss, eta):

    theta = _concat(self.target_net.parameters()).data
    dtheta = _concat(torch.autograd.grad(loss, self.target_net.parameters(), retain_graph=True)).data
    # dtheta = _concat([v.grad for v in self.target_net.parameters()]).data
    if self.args.weight_decay != 0:
      dtheta.add_(self.args.weight_decay, theta)
    if self.args.momentum != 0:
      try:
        moment = _concat(self.target_net_optim.state[v]['momentum_buffer'] for v in self.target_net.parameters()).mul_(self.args.momentum)
      except:
        # setting zeros is consistent with the original momentum optimizer
        moment = torch.zeros_like(theta)
    unrolled_model = self._construct_model_from_theta(theta.sub(eta, moment + dtheta))
    return unrolled_model

  # def _compute_unrolled_model(self, eta):
  #
  #
  #   theta = _concat([v.data for v in self.target_net.parameters()])
  #   try:
  #     moment = _concat(self.target_net_optim.state[v]['momentum_buffer'] for v in self.target_net.parameters()).mul_(self.args.momentum)
  #   except:
  #     moment = torch.zeros_like(theta)
  #   dtheta = _concat([v.grad.data for v in self.target_net.parameters()]) + self.args.weight_decay * theta
  #
  #   unrolled_target_net = self._construct_model_from_theta(theta.sub(eta, moment+dtheta))
  #
  #   return unrolled_target_net

  def step(self, input_train, target_train, input_valid, target_valid, unrolled):
    self.aug_net_optim.zero_grad()
    if unrolled:
        # self._backward_step_unrolled(input_train, target_train, input_valid, target_valid)
        loss_adv, loss_div = self._backward_step_unrolled(input_train, target_train, input_valid, target_valid)
    else:
        self._backward_step(input_valid, target_valid)
    self.aug_net_optim.step()
    return loss_adv, loss_div

  def _backward_step(self, input_valid, target_valid):
    output_valid = self.target_net(input_valid)
    loss = self.criterion(output_valid, target_valid)
    loss.backward()

  def _backward_step_unrolled(self, input_train, target_train, input_valid, target_valid):
    # input_train_aug, div_loss = self.aug_net(input_train, require_loss=True)
    input_train_aug, div_loss = self.aug_net(input_train, require_loss=True)
    output_train = self.target_net(input_train_aug)
    loss_train = self.criterion(output_train, target_train)
    # loss_train_2 = loss_train + div_loss * self.args.div_weight
    # print('div_loss: {}'.format(div_loss))
    # print('loss_train: {}'.format(loss_train))
    # print('loss_train_2: {}'.format(loss_train_2))
    # print('self.args.div_weight: {}'.format(self.args.div_weight))
    # exit()
    # loss_train_2.backward()
    eta = self.target_net_optim.param_groups[0]['lr']
    # eta = self.lr_scheduler.get_lr()[0]
    unrolled_target_net = self._compute_unrolled_model(loss_train, eta)

    # # check whether the unrolled_target_net is different from the original target_net
    # output2 = unrolled_target_net(input_train)
    # loss2 = self.criterion(output2, target_train)
    # print('loss1: {:4f}'.format(loss_train))
    # print('loss2: {:4f}'.format(loss2))
    # exit()
    output_valid = unrolled_target_net(input_valid)
    unrolled_loss = self.criterion(output_valid, target_valid)

    unrolled_loss.backward()
    #
    loss_train_aug = -loss_train * self.args.adv_weight + div_loss * self.args.div_weight
    dalpha = torch.autograd.grad(loss_train_aug, self.aug_net.parameters())
    # dalpha = [per_grad.data.clamp_(min=-1, max=1) for per_grad in dalpha]

    vector = [v.grad.data for v in unrolled_target_net.parameters()]
    implicit_grads = self._hessian_vector_product(vector, input_train, target_train, r=self.args.val_r)

    for g, ig in zip(dalpha, implicit_grads):
      g.data.sub_(eta, ig.data)

    for v, g in zip(self.aug_net.parameters(), dalpha):
      if v.grad is None:
        # print('grad is none. existing ...')
        # exit()
        v.grad = g.detach()
      else:
        v.grad.data.copy_(g.data)
    return -loss_train * self.args.adv_weight, div_loss * self.args.div_weight

  def _construct_model_from_theta(self, theta):
    # print('type of theta: {}'.format(type(theta)))
    theta = nn.Parameter(theta)
    target_net_new = utils.build_model(self.args)
    # .state_dict() stores all the persistent buffers (e.g. running averages), which are not included in .parameters()
    model_dict = self.target_net.state_dict()

    params, offset = {}, 0
    for k, v in self.target_net.named_parameters():
      v_length = np.prod(v.size())
      params[k] = theta[offset: offset+v_length].view(v.size())
      # print('type of params[k]: {}'.format(type(params[k])))
      offset += v_length

    assert offset == len(theta)
    model_dict.update(params)
    target_net_new.load_state_dict(model_dict)
    return target_net_new.cuda()

  def _hessian_vector_product(self, vector, input, target, r=2e-2):
    R = r / _concat(vector).norm()
    for p, v in zip(self.target_net.parameters(), vector):
      p.data.add_(R, v)

    # input_aug = self.aug_net(input)
    input_aug = self.call_aug_net(input, target)
    output_aug = self.target_net(input_aug)
    loss = self.criterion(output_aug, target)
    grads_p = torch.autograd.grad(loss, self.aug_net.parameters())

    for p, v in zip(self.target_net.parameters(), vector):
      p.data.sub_(2*R, v)

    # input_aug = self.aug_net(input)
    input_aug = self.call_aug_net(input, target)
    output_aug = self.target_net(input_aug)
    loss = self.criterion(output_aug, target)
    grads_n = torch.autograd.grad(loss, self.aug_net.parameters())

    # recover the original weights in self.target_net
    for p, v in zip(self.target_net.parameters(), vector):
      p.data.add_(R, v)

    return [(x-y).div_(2*R) for x, y in zip(grads_p, grads_n)]


import torch
from torch.utils.data import Dataset
from torchvision import transforms
import torch.nn.functional as F

import os
import random

from PIL import Image

import numpy as np

class YOLOv3Dataset(Dataset):

    def __init__(self, images_list_path="../data/images_list.txt", image_size=416, should_augment=True, transform=None):
        super(YOLOv3Dataset, self).__init__()
        with open(images_list_path, "r") as f:
            self.images_paths = f.read().splitlines()
        self.labels_paths = [path.replace("images", "labels")
                                 .replace(".png", ".txt")
                                 .replace(".jpg", ".txt")
                                 .replace(".jpeg", ".txt")
                             for path in self.images_paths]
        self.image_size = image_size
        self.should_augment = should_augment
        self.transform = transform

    def __getitem__(self, item):

        image_path = self.images_paths[item % len(self.images_paths)]
        image_tensor = transforms.ToTensor()(Image.open(image_path).convert("RGB"))

        if len(image_tensor.shape) != 3:
            image_tensor = image_tensor.unsqueeze(0)
            image_tensor = image_tensor.expand((3, image_tensor.shape[1:]))

        _, height, width = image_tensor.shape

        image_tensor, pad = YOLOv3Dataset.pad_to_square(image_tensor)
        _, padded_h, padded_w = image_tensor.shape

        label_path = self.labels_paths[item % len(self.labels_paths)].rstrip()
        if os.path.exists(label_path):
            boxes = torch.from_numpy(np.loadtxt(label_path).reshape(-1, 5))

            # Extract bounding boxes coordinates, and convert to desired format.
            # We leave boxes[:,0] as is, since that's the class annotation
            # We change the scale from 1*1 to w*d, since we want to apply to the tensor
            # Note: The bounding boxes are given in the format:
            #       class, left_x, up_y, right_x, bottom_y
            #       However, we need to return them in the format:
            #       class, center_x, center_y, width, height
            #       So this is why convert. #whyweconvert
            x1 = width * (boxes[:, 1] - boxes[:, 3] / 2) + pad[0]
            y1 = height * (boxes[:, 2] - boxes[:, 4] / 2) + pad[1]
            x2 = width * (boxes[:, 1] + boxes[:, 3] / 2) + pad[2]
            y2 = height * (boxes[:, 2] + boxes[:, 4] / 2) + pad[3]

            boxes[:, 1] = ((x1 + x2) / 2) / padded_w
            boxes[:, 2] = ((y1 + y2) / 2) / padded_h
            boxes[:, 3] *= width / padded_w
            boxes[:, 4] *= height / padded_h
            label_tensor = torch.zeros(len(boxes), 6)
            label_tensor[:, 1:] = boxes

        else:
            print("File is missing: "+label_path)
            return
        if self.should_augment and self.transform:
            image_tensor, label_tensor = self.transform(image_tensor), self.transform(label_tensor)

        return image_path, image_tensor, label_tensor

    def collate_fn(self, batch):
        """
        Merges a list of samples to form a mini-batch of Tensor(s). Used when using batched loading from a map-style dataset

        :param batch:
        :return:
        """
        paths, imgs, targets = list(zip(*batch))

        # Remove empty placeholder targets
        targets = [boxes for boxes in targets if boxes is not None]

        # Add sample index to targets
        for i, boxes in enumerate(targets):
            boxes[:, 0] = i
        targets = torch.cat(targets, 0)

        # Selects new image size every tenth batch
        if self.multiscale and self.batch_count % 10 == 0:
            self.img_size = random.choice(range(self.min_size, self.max_size + 1, 32))

        # Resize images to input shape
        imgs = torch.stack([transforms.Resize(img, self.img_size) for img in imgs])
        self.batch_count += 1

        return paths, imgs, targets

    def __len__(self):
        return len(self.images_paths)

    @staticmethod
    def pad_to_square(image_tensor):
        """
        Padding the image tensor in order to get a square
        :param image_tensor: A rectangle shape image tensor
        :return: A padded square tensor of the image
        """
        channels, height, width = image_tensor.shape
        dim_diff = np.abs(height - width)

        # Padding values for the needed dimension:
        pad1 = dim_diff // 2
        pad2 = dim_diff - pad1

        # Determine padding
        pad = (0, 0, pad1, pad2) if height <= width else (pad1, pad2, 0, 0)

        # Add padding
        padded_image_tensor = F.pad(image_tensor, pad, "constant", value=0)

        return padded_image_tensor, pad


import os
import time
import warnings

warnings.filterwarnings("ignore", message="numpy.dtype size changed")
warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
import numpy as np
import queue
import tensorflow as tf

import kaldi_io
from tf_block import batch_norm_wrapper, prelu
from ze_utils import set_cuda_visible_devices

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
VAR2STD_EPSILON = 0.00001

# noinspection PyAttributeOutsideInit
class Model(object):

    def __init__(self):
        self.graph = None

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

                    # Apply dropout
                    if i != len(kernel_sizes) - 1:
                        with tf.name_scope("dropout-%s" % i):
                            h = tf.nn.dropout(h, self.dropout_keep_prob)

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):

                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim
                    if i != len(embedding_sizes) - 1:
                        with tf.name_scope("dropout-%s" % i):
                            h = tf.nn.dropout(h, self.dropout_keep_prob)

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            self.loss = tf.reduce_mean(losses, name="loss")

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

    @staticmethod
    def save_model(sess, output_dir, logger):
        if logger is not None:
            logger.info("Start saving graph ...")
        saver = tf.train.Saver()
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        save_path = saver.save(sess, os.path.join(output_dir, 'model'))
        with open(os.path.join(output_dir, 'done'), 'wt') as fid:
            fid.write('done')
        if logger is not None:
            logger.info("Graph saved in path: %s" % save_path)

    def load_model(self, sess, input_dir, logger):
        if logger is not None:
            logger.info("Start loading graph ...")
        saver = tf.train.import_meta_graph(os.path.join(input_dir, 'model.meta'))
        saver.restore(sess, os.path.join(input_dir, 'model'))
        self.graph = sess.graph
        self.input_x = self.graph.get_tensor_by_name("input_x:0")
        self.input_y = self.graph.get_tensor_by_name("input_y:0")
        self.num_classes = self.input_y.shape[1]
        self.learning_rate = self.graph.get_tensor_by_name("learning_rate:0")
        self.dropout_keep_prob = self.graph.get_tensor_by_name("dropout_keep_prob:0")
        self.phase = self.graph.get_tensor_by_name("phase:0")
        self.loss = self.graph.get_tensor_by_name("loss:0")
        self.optimizer = self.graph.get_operation_by_name("optimizer")
        self.accuracy = self.graph.get_tensor_by_name("accuracy/accuracy:0")
        self.embedding = [None] * 2  # TODO make this more general
        self.embedding[0] = self.graph.get_tensor_by_name("embed_layer-0/scores:0")
        self.embedding[1] = self.graph.get_tensor_by_name("embed_layer-1/scores:0")
        if logger is not None:
            logger.info("Graph restored from path: %s" % input_dir)

    def create_one_hot_output_matrix(self, labels):
        minibatch_size = len(labels)
        one_hot_matrix = np.zeros((minibatch_size, self.num_classes), dtype=np.int32)
        for i, lab in enumerate(labels):
            one_hot_matrix[i, lab] = 1
        return one_hot_matrix

    def print_models_params(self, input_dir, logger=None):
        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)) as sess:
            self.load_model(sess, input_dir, logger)
            print('\n\nThe components are:\n')
            for v in self.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
                print(v.name)
            print('\n')

    def get_models_weights(self, input_dir, logger=None):
        import h5py
        h5file = os.path.join(input_dir, 'model.h5')
        if os.path.exists(h5file):
            name2weights = {}

            def add2weights(name, mat):
                if not isinstance(mat, h5py.Group):
                    # print('%s  shape: %s' % (name, str(mat.shape)))
                    name2weights[name] = mat.value

            with h5py.File(h5file, 'r') as hf:
                hf.visititems(add2weights)
            return name2weights

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)) as sess:
            self.load_model(sess, input_dir, logger)
            name2weights = {}
            for v in self.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
                name2weights[v.name] = sess.run(v)
                print('%s  shape: %s' % (v.name, str(name2weights[v.name].shape)))
            for i in range(5):
                for scope_name in ("frame_level_info_layer-%s" % i, "embed_layer-%s" % i):
                    for var_name in ("mean", "variance"):
                        name = '%s/%s:0' % (scope_name, var_name)
                        try:
                            name2weights[name] = sess.run(self.graph.get_tensor_by_name(name))
                            print('%s  shape: %s' % (name, str(name2weights[name].shape)))
                        except:
                            pass
            with h5py.File(h5file, 'w') as hf:
                for name, mat in name2weights.iteritems():
                    hf.create_dataset(name, data=mat.astype(np.float32))
            return name2weights

    def train_one_iteration(self, data_loader, args, logger):
        learning_rate = args.learning_rate
        print_interval = args.print_interval
        dropout_proportion = args.dropout_proportion
        input_dir = args.input_dir

        output_dir = args.output_dir
        random_seed = 4 * args.random_seed + args.random_seed % 3

        set_cuda_visible_devices(use_gpu=True, logger=logger)
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=False, log_device_placement=False)) as sess:
            if random_seed != 0:
                tf.set_random_seed(random_seed)
            self.load_model(sess, input_dir, logger)

            # Shuffle the data in each epoch
            minibatch_count = data_loader.count
            start_minibatch = 1
            dropout_keep_prob = 1 - dropout_proportion
            total_segments, minibatch_segments = 0, 0
            total_loss, minibatch_loss = 0, 0
            total_objective, minibatch_objective = 0, 0
            total_accuracy, minibatch_accuracy = 0, 0
            total_segments_len = 0
            total_gpu_waiting = 0.0
            total_disk_waiting = 0.0
            start_time = time.time()
            for minibatch_idx in range(minibatch_count):
                try:
                    disk_waiting = time.time()
                    batch_data, labels = data_loader.pop()
                    total_disk_waiting += time.time() - disk_waiting
                except queue.Empty:
                    logger.warning('Timeout reach when reading the minibatch index %d' % minibatch_idx)
                    continue
                if batch_data is None:
                    logger.warning('batch_data is None for the minibatch index %d' % minibatch_idx)
                    continue
                batch_labels = self.create_one_hot_output_matrix(labels)
                minibatch_segments += batch_data.shape[0]
                total_segments += batch_data.shape[0]
                total_segments_len += batch_data.shape[1]
                feed_dict = {self.input_x: batch_data, self.input_y: batch_labels,
                             self.dropout_keep_prob: dropout_keep_prob, self.learning_rate: learning_rate,
                             self.phase: True}

                gpu_waiting = time.time()
                _, loss, accuracy = sess.run([self.optimizer, self.loss, self.accuracy], feed_dict=feed_dict)
                total_gpu_waiting += time.time() - gpu_waiting

                objective = -loss
                total_loss += loss
                minibatch_loss += loss
                total_objective += objective
                minibatch_objective += objective
                total_accuracy += accuracy
                minibatch_accuracy += accuracy
                end_minibatch = minibatch_idx + 1
                if end_minibatch % print_interval == 0:
                    cnt = end_minibatch - start_minibatch + 1
                    logger.info("Average training loss for minibatches %d-%d is %.4f over %d segments. Also, the "
                                "average training accuracy for these minibatches is %.4f and the average "
                                "objective function for these minibatches is %.4f. Average DISK waiting: %.1f "
                                "secs and average GPU waiting: %.1f secs for each minibatch." %
                                (start_minibatch, end_minibatch, minibatch_loss / cnt,
                                 minibatch_segments, minibatch_accuracy / cnt, minibatch_objective / cnt,
                                 total_disk_waiting / cnt, total_gpu_waiting / cnt))
                    start_minibatch = end_minibatch + 1
                    minibatch_segments = 0
                    minibatch_loss = 0
                    minibatch_accuracy = 0
                    minibatch_objective = 0
                    total_gpu_waiting = 0.0
                    total_disk_waiting = 0.0

            logger.info("Processed %d segments of average size %d into %d minibatches. Avg minibatch size was %d." %
                        (total_segments, total_segments_len / minibatch_count, minibatch_count,
                         total_segments / minibatch_count))

            logger.info("Overall average training loss is %.4f over %d segments. Also, the overall "
                        "average training accuracy is %.4f." % (total_loss / minibatch_count,
                                                                total_segments, total_accuracy / minibatch_count))

            logger.info("Overall average objective function is %.4f over %d segments." %
                        (total_objective / minibatch_count, total_segments))

            Model.save_model(sess, output_dir, logger)

            logger.info("Elapsed time for processing whole training minibatches is %.2f minutes." %
                        ((time.time() - start_time) / 60.0))

    def eval(self, data_loader, input_dir, use_gpu, logger):

        set_cuda_visible_devices(use_gpu=use_gpu, logger=logger)
        with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)) as sess:
            self.load_model(sess, input_dir, logger)

            # Shuffle the data in each epoch
            minibatch_count = data_loader.count
            total_segments = 0
            total_loss = 0
            total_accuracy = 0
            total_segments_len = 0
            total_gpu_waiting = 0.0
            total_disk_waiting = 0.0
            start_time = time.time()
            for minibatch_idx in range(minibatch_count):
                try:
                    disk_waiting = time.time()
                    batch_data, labels = data_loader.pop()
                    total_disk_waiting += time.time() - disk_waiting
                except queue.Empty:
                    logger.warning('Timeout reach when reading minibatch index %d' % minibatch_idx)
                    continue
                if batch_data is None:
                    logger.warning('batch_data is None for minibatch index %d' % minibatch_idx)
                    continue
                batch_labels = self.create_one_hot_output_matrix(labels)
                total_segments += batch_data.shape[0]
                total_segments_len += batch_data.shape[1]
                feed_dict = {self.input_x: batch_data, self.input_y: batch_labels, self.dropout_keep_prob: 1.0,
                             self.phase: False}

                gpu_waiting = time.time()
                loss, accuracy = sess.run([self.loss, self.accuracy], feed_dict=feed_dict)
                total_gpu_waiting += time.time() - gpu_waiting
                total_loss += loss
                total_accuracy += accuracy

            logger.info("Processed %d segments of average size %d into %d minibatches. Avg minibatch size was %d." %
                        (total_segments, total_segments_len / minibatch_count, minibatch_count,
                         total_segments / minibatch_count))

            logger.info("Overall average loss is %.4f over %d segments. Also, the overall "
                        "average accuracy is %.4f." % (total_loss / minibatch_count, total_segments,
                                                       total_accuracy / minibatch_count))

            logger.info("Elapsed time for processing whole training minibatches is %.2f minutes." %
                        ((time.time() - start_time) / 60.0))

    def make_embedding(self, input_stream, output_stream, model_dir, min_chunk_size, chunk_size, use_gpu, logger):
        start_time = time.time()
        set_cuda_visible_devices(use_gpu=use_gpu, logger=logger)

        config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)
        if not use_gpu:
            config.intra_op_parallelism_threads = 2
            config.inter_op_parallelism_threads = 2

        with tf.Session(config=config) as sess:
            self.load_model(sess, model_dir, logger)

            total_segments = 0
            total_segments_len = 0
            total_gpu_waiting = 0.0
            num_fail = 0
            num_success = 0
            for key, mat in kaldi_io.read_mat_ark(input_stream):
                logger.info("Processing features with key '%s' which have shape '%s'" % (key, str(mat.shape)))
                total_segments += 1

                num_rows = mat.shape[0]
                if num_rows == 0:
                    logger.warning("Zero-length utterance: '%s'" % key)
                    num_fail += 1
                    continue

                if num_rows < min_chunk_size:
                    logger.warning("Minimum chunk size of %d is greater than the number of rows in utterance: %s" %
                                   (min_chunk_size, key))
                    num_fail += 1
                    continue
                this_chunk_size = chunk_size
                if num_rows < chunk_size:
                    logger.info("Chunk size of %d is greater than the number of rows in utterance: %s, "
                                "using chunk size of %d" % (chunk_size, key, num_rows))
                    this_chunk_size = num_rows
                elif chunk_size == -1:
                    this_chunk_size = num_rows

                num_chunks = int(np.ceil(num_rows / float(this_chunk_size)))
                # logger.info("num_chunks: %d" % num_chunks)
                xvector_avg = 0
                tot_weight = 0.0

                for chunk_idx in range(num_chunks):
                    # If we're nearing the end of the input, we may need to shift the
                    # offset back so that we can get this_chunk_size frames of input to
                    # the nnet.
                    offset = min(this_chunk_size, num_rows - chunk_idx * this_chunk_size)
                    if offset < min_chunk_size:
                        continue
                    # logger.info("offset: %d" % offset)
                    sub_mat = mat[chunk_idx * this_chunk_size: chunk_idx * this_chunk_size + offset, :]
                    data = np.reshape(sub_mat, (1, sub_mat.shape[0], sub_mat.shape[1]))
                    total_segments_len += sub_mat.shape[0]
                    feed_dict = {self.input_x: data, self.dropout_keep_prob: 1.0, self.phase: False}
                    gpu_waiting = time.time()
                    xvector = sess.run(self.embedding[0], feed_dict=feed_dict)
                    xvector = xvector[0]
                    # logger.info("xvector: %s" % str(xvector.shape))
                    total_gpu_waiting += time.time() - gpu_waiting
                    tot_weight += offset
                    xvector_avg += offset * xvector

                xvector_avg /= tot_weight
                kaldi_io.write_vec_flt(output_stream, xvector_avg, key=key)
                num_success += 1

            logger.info("Processed %d features of average size %d frames. Done %d and failed %d" %
                        (total_segments, total_segments_len / total_segments, num_success, num_fail))

            logger.info("Total time for neural network computations is %.2f minutes." %
                        (total_gpu_waiting / 60.0))

            logger.info("Elapsed time for extracting whole embeddings is %.2f minutes." %
                        ((time.time() - start_time) / 60.0))

# noinspection PyAttributeOutsideInit
class ModelWithoutDropout(Model):

    def __init__(self):
        super(ModelWithoutDropout, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            self.loss = tf.reduce_mean(losses, name="loss")

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelWithoutDropoutTdnn(Model):

    def __init__(self):
        super(ModelWithoutDropoutTdnn, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 3, 3, 1, 1]
        embedding_sizes = [512, 512]
        dilation_rates = [1, 2, 3, 1, 1]

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.convolution(h, w, dilation_rate=[dilation_rates[i]], padding="SAME",
                                             name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            self.loss = tf.reduce_mean(losses, name="loss")

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelWithoutDropoutPRelu(Model):

    def __init__(self):
        super(ModelWithoutDropoutPRelu, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = prelu(h, shared=False)
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    h = prelu(h, shared=False)
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            self.loss = tf.reduce_mean(losses, name="loss")

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelL2LossWithoutDropoutPRelu(Model):

    def __init__(self):
        super(ModelL2LossWithoutDropoutPRelu, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]
        beta = 0.0002

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            l2_loss = tf.constant(0.0)

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = prelu(h, shared=False)
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    # Apply L2 loss
                    if i == 0:
                        l2_loss += 0.1 * tf.nn.l2_loss(w)
                        l2_loss += 0.1 * tf.nn.l2_loss(b)
                    else:
                        l2_loss += tf.nn.l2_loss(w)
                        l2_loss += tf.nn.l2_loss(b)

                    h = prelu(h, shared=False)
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                # Apply L2 loss
                l2_loss += tf.nn.l2_loss(w)
                l2_loss += tf.nn.l2_loss(b)

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            # Normal loss function
            loss = tf.reduce_mean(losses, name="orig_loss")
            self.loss = tf.reduce_mean(loss + beta * l2_loss, name='loss')

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelL2LossWithoutDropoutLRelu(Model):

    def __init__(self):
        super(ModelL2LossWithoutDropoutLRelu, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]
        beta = 0.0002

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            l2_loss = tf.constant(0.0)

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    h = tf.nn.leaky_relu(h, alpha=0.2, name='lrelu')
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    # Apply L2 loss
                    if i == 0:
                        l2_loss += 0.1 * tf.nn.l2_loss(w)
                        l2_loss += 0.1 * tf.nn.l2_loss(b)
                    else:
                        l2_loss += tf.nn.l2_loss(w)
                        l2_loss += tf.nn.l2_loss(b)

                    h = tf.nn.leaky_relu(h, alpha=0.2, name='lrelu')
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                # Apply L2 loss
                l2_loss += tf.nn.l2_loss(w)
                l2_loss += tf.nn.l2_loss(b)

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            # Normal loss function
            loss = tf.reduce_mean(losses, name="orig_loss")
            self.loss = tf.reduce_mean(loss + beta * l2_loss, name='loss')

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelL2LossWithoutDropoutLReluAttention(Model):

    def __init__(self):
        super(ModelL2LossWithoutDropoutLReluAttention, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 6 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]
        beta = 0.0002

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            l2_loss = tf.constant(0.0)

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[layer_size]), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = tf.nn.leaky_relu(h, alpha=0.2, name='lrelu')
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            prev_dim /= 2

            # apply self attention
            with tf.variable_scope("attention"):
                b = tf.Variable(tf.constant(0.1, shape=[prev_dim]), name="b")
                v = tf.Variable(tf.constant(0.1, shape=[prev_dim]), name="v")
                # Note: the dimension of the w needs more experiments, here we simply use a square matrix
                w = tf.Variable(tf.truncated_normal([prev_dim, prev_dim], stddev=0.1), name="w")
                h1, h2 = tf.split(h, 2, axis=2)
                non_linearity = tf.nn.tanh(tf.nn.bias_add(tf.einsum('ijk,kl->ijl', h1, w), b), name="non_linearity")
                attention = tf.nn.softmax(tf.einsum('ijk,k->ij', non_linearity, v), name="attention")

            h_m = tf.einsum('ijk,ij->ik', h2, attention)
            h_s = tf.subtract(tf.einsum('ijk,ij->ik', tf.square(h2), attention), tf.square(h_m), name="stats_var")
            h = tf.concat([h_m, tf.sqrt(h_s + VAR2STD_EPSILON)], 1)

            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=0.1), name="w")
                    b = tf.Variable(tf.constant(0.1, shape=[out_dim]), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    # Apply L2 loss
                    if i == 0:
                        l2_loss += 0.1 * tf.nn.l2_loss(w)
                        l2_loss += 0.1 * tf.nn.l2_loss(b)
                    else:
                        l2_loss += tf.nn.l2_loss(w)
                        l2_loss += tf.nn.l2_loss(b)

                    h = tf.nn.leaky_relu(h, alpha=0.2, name='lrelu')
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                w = tf.get_variable("w", shape=[prev_dim, num_classes],
                                    initializer=tf.contrib.layers.xavier_initializer())
                b = tf.Variable(tf.constant(0.1, shape=[num_classes]), name="b")

                # Apply L2 loss
                l2_loss += tf.nn.l2_loss(w)
                l2_loss += tf.nn.l2_loss(b)

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            # Normal loss function
            loss = tf.reduce_mean(losses, name="orig_loss")
            self.loss = tf.reduce_mean(loss + beta * l2_loss, name='loss')

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")

# noinspection PyAttributeOutsideInit
class ModelL2LossWithoutDropoutReluHeInit(Model):

    def __init__(self):
        super(ModelL2LossWithoutDropoutReluHeInit, self).__init__()

    def build_model(self, num_classes, input_feature_dim, output_dir, logger=None):

        layer_sizes = [512, 512, 512, 512, 3 * 512]
        kernel_sizes = [5, 5, 7, 1, 1]
        embedding_sizes = [512, 512]
        beta = 0.0002

        if logger is not None:
            logger.info("Start building the model ...")

        tf.reset_default_graph()
        self.graph = tf.Graph()

        with self.graph.as_default():

            self.num_classes = num_classes

            # placeholder for parameter
            self.learning_rate = tf.placeholder(tf.float32, name="learning_rate")
            self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
            self.phase = tf.placeholder(tf.bool, name="phase")

            # Placeholders for regular data
            self.input_x = tf.placeholder(tf.float32, [None, None, input_feature_dim], name="input_x")
            self.input_y = tf.placeholder(tf.float32, [None, num_classes], name="input_y")

            l2_loss = tf.constant(0.0)

            h = self.input_x

            # Frame level information Layer
            prev_dim = input_feature_dim
            for i, (kernel_size, layer_size) in enumerate(zip(kernel_sizes, layer_sizes)):
                with tf.variable_scope("frame_level_info_layer-%s" % i):
                    kernel_shape = [kernel_size, prev_dim, layer_size]
                    # he_normal init
                    fan_in = kernel_size * prev_dim
                    w = tf.Variable(tf.truncated_normal(kernel_shape, stddev=tf.sqrt(2.0 / fan_in)), name="w")
                    # he_uniform init
                    limit = tf.sqrt(6.0 / fan_in)
                    b = tf.Variable(tf.random_uniform([layer_size], minval=-limit, maxval=limit), name="b")

                    conv = tf.nn.conv1d(h, w, stride=1, padding="SAME", name="conv-layer-%s" % i)
                    h = tf.nn.bias_add(conv, b)

                    # Apply nonlinearity and BN
                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = layer_size

            # Statistic pooling
            tf_mean, tf_var = tf.nn.moments(h, 1)
            h = tf.concat([tf_mean, tf.sqrt(tf_var + VAR2STD_EPSILON)], 1)
            prev_dim = prev_dim * 2

            # Embedding layers
            for i, out_dim in enumerate(embedding_sizes):
                with tf.variable_scope("embed_layer-%s" % i):
                    # he_normal init
                    w = tf.Variable(tf.truncated_normal([prev_dim, out_dim], stddev=tf.sqrt(2.0 / prev_dim)), name="w")
                    # he_uniform init
                    limit = tf.sqrt(6.0 / prev_dim)
                    b = tf.Variable(tf.random_uniform([out_dim], minval=-limit, maxval=limit), name="b")

                    h = tf.nn.xw_plus_b(h, w, b, name="scores")

                    # Apply L2 loss
                    if i == 0:
                        l2_loss += 0.1 * tf.nn.l2_loss(w)
                        l2_loss += 0.1 * tf.nn.l2_loss(b)
                    else:
                        l2_loss += tf.nn.l2_loss(w)
                        l2_loss += tf.nn.l2_loss(b)

                    h = tf.nn.relu(h, name="relu")
                    h = batch_norm_wrapper(h, decay=0.95, is_training=self.phase)

                    prev_dim = out_dim

            # Softmax
            with tf.variable_scope("output"):
                # glorot_normal: https://keras.io/initializers/
                stddev = tf.sqrt(2.0 / (prev_dim + num_classes))
                w = tf.Variable(tf.truncated_normal([prev_dim, num_classes], stddev=stddev), name="w")
                # glorot_uniform
                limit = tf.sqrt(6.0 / (prev_dim + num_classes))
                b = tf.Variable(tf.random_uniform([num_classes], minval=-limit, maxval=limit), name="b")

                # Apply L2 loss
                l2_loss += tf.nn.l2_loss(w)
                l2_loss += tf.nn.l2_loss(b)

                scores = tf.nn.xw_plus_b(h, w, b, name="scores")

                predictions = tf.argmax(scores, 1, name="predictions")

            losses = tf.nn.softmax_cross_entropy_with_logits(logits=scores, labels=self.input_y)

            # Normal loss function
            loss = tf.reduce_mean(losses, name="orig_loss")
            self.loss = tf.reduce_mean(loss + beta * l2_loss, name='loss')

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss,
                                                                                name="optimizer")

            with tf.name_scope("accuracy"):
                correct_predictions = tf.equal(predictions, tf.argmax(self.input_y, 1))
                self.accuracy = tf.reduce_mean(tf.cast(correct_predictions, "float"), name="accuracy")

        set_cuda_visible_devices(use_gpu=False, logger=logger)
        with tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
                                                                log_device_placement=False)) as sess:
            if logger is not None:
                logger.info("Start initializing the graph ...")
            sess.run(tf.global_variables_initializer())
            Model.save_model(sess, output_dir, logger)

        if logger is not None:
            logger.info("Building finished.")


import io, sys, os, csv, random, logging
from jacks.infer import LOG, inferJACKS
from jacks.jacks_io import loadSgrnaReference, readControlGeneset, createGeneSpec, createSampleSpec, getJacksParser, collateTestControlSamples, writeJacksWResults, writeJacksXResults
from jacks.preprocess import loadDataAndPreprocess
from run_JACKS_single import filterSampleSpec, filterCtrlSpec
import numpy as np

py_cmd = 'python'

def filterOutSampleSpec(sample_spec, cell_line, ctrl_spec):
    new_sample_spec = {}
    for filename in sample_spec:
        for sample_id, colname in sample_spec[filename]:
            if sample_id == cell_line: continue
            if filename not in new_sample_spec:
                new_sample_spec[filename] = []
            new_sample_spec[filename].append((sample_id, colname))
    return new_sample_spec

LOG.setLevel(logging.WARNING)
parser = getJacksParser()
parser.add_argument("--cell_line",
                type=str,
                default=None,
                help="cell line to run")
parser.add_argument("--separate",
                action='store_true',
                default=False,
                help="Run all cell lines separately")
args = parser.parse_args()

outprefix = args.outprefix
if '/' in outprefix and not os.path.exists(os.path.dirname(outprefix)): os.makedirs(os.path.dirname(outprefix))

# Load the specification of samples to include
LOG.info('Loading sample specification')
sample_spec, ctrl_spec, sample_num_reps = createSampleSpec(args.countfile, args.replicatefile, args.rep_hdr,
                                                            args.sample_hdr, args.common_ctrl_sample, args.ctrl_sample_hdr)
if args.cell_line != None:
    single_sample_spec = filterSampleSpec(sample_spec, args.cell_line, ctrl_spec)
    single_ctrl_spec = filterCtrlSpec(ctrl_spec, args.cell_line)
    ref_sample_spec = filterOutSampleSpec(sample_spec, args.cell_line, ctrl_spec)
    ref_outfile =  outprefix + '_xreference_' + args.cell_line + '.txt'

elif args.separate:
    for cell_line in ctrl_spec:
        if ctrl_spec[cell_line] == cell_line: continue
        cmd = '%s --cell_line=%s' % (' '.join(sys.argv), cell_line)
        os.system('%s %s' % (py_cmd,cmd))
    exit()

# Load the mappings from guides to genes
LOG.info('Loading gene mappings')
gene_spec = createGeneSpec(args.guidemappingfile, args.sgrna_hdr, args.gene_hdr, ignore_blank_genes=args.ignore_blank_genes)

# Load negative control guides (if any)
ctrl_geneset = readControlGeneset(args.ctrl_genes) if args.ctrl_genes is not None else set()

##REFERENCE (to collect X's)

# Load the data and preprocess
LOG.info('Reference: Loading data and pre-processing')
data, meta, sample_ids, genes, gene_index = loadDataAndPreprocess(ref_sample_spec, gene_spec,ctrl_spec=ctrl_spec, normtype=args.norm_type, ctrl_geneset=ctrl_geneset)
gene_grnas = {gene: [x for x in meta[gene_index[gene], 0]] for gene in gene_index}
testdata, ctrldata, test_sample_idxs = collateTestControlSamples(data, sample_ids, ctrl_spec)
sample_ids_without_ctrl = [sample_ids[idx] for idx in test_sample_idxs]

#Run all samples against their controls
LOG.info('Reference: Running JACKS inference')
jacks_results = inferJACKS(gene_index, testdata, ctrldata)
writeJacksXResults(ref_outfile, jacks_results, gene_grnas)

##TEST (using reference)
LOG.info('Test: Loading data and pre-processing')
data, meta, sample_ids, genes, gene_index = loadDataAndPreprocess(single_sample_spec, gene_spec,ctrl_spec=single_ctrl_spec, normtype=args.norm_type, ctrl_geneset=ctrl_geneset)
gene_grnas = {gene: [x for x in meta[gene_index[gene], 0]] for gene in gene_index}
testdata, ctrldata, test_sample_idxs = collateTestControlSamples(data, sample_ids, ctrl_spec)
sample_ids_without_ctrl = [sample_ids[idx] for idx in test_sample_idxs]

x_ref = loadSgrnaReference(ref_outfile)
# Create the X reference (in the correct order)
x_reference = {'X1': np.array([eval(x_ref[x]['X1']) for x in meta[:, 0]]),
                'X2': np.array([eval(x_ref[x]['X2']) for x in meta[:, 0]])}

#Run all samples against their controls
LOG.info('Test: Running JACKS inference using reference')
test_jacks_results = inferJACKS(gene_index, testdata, ctrldata, w_only=True, fixed_x=x_reference)

# Write out the results
LOG.info('Writing JACKS results')
writeJacksWResults(outprefix + '_' + args.cell_line, test_jacks_results, sample_ids_without_ctrl, ctrl_geneset=ctrl_geneset, write_types=[''])


#!/usr/bin/env python
# coding=utf-8

'''
Author:Tai Lei
Date:
Info:
'''

import random
import numpy as np
import rospkg
from lxml import etree
from lxml.etree import Element
from copy import deepcopy
import yaml

rospack = rospkg.RosPack()
with open(rospack.get_path("actor_services")+"/src/multiforceFactors.yaml", 'r') as stream:
    try:
        factorData = yaml.load(stream)
    except yaml.YAMLError as exc:
        print(exc)

SocialForce = factorData["SocialForceFactor"]
DesiredForce = factorData["DesiredForceFactor"]
ObstacleForce = factorData["ObstacleForceFactor"]
AnimationFactor = factorData["AnimationFactor"]
print(SocialForce)
print(DesiredForce)
print(ObstacleForce)
print(AnimationFactor)

plugin_pkg_path = rospack.get_path("actor_plugin")
plugin_path = plugin_pkg_path + "/lib/libactorplugin_ros.so"
actor_pkg_path = rospack.get_path("actor_services")

tree_ = etree.parse(actor_pkg_path+'/worlds/empty.world')
world_ = tree_.getroot().getchildren()[0]

skin_list = ["moonwalk.dae",
        "run.dae",
        "sit_down.dae",
        "sitting.dae",
        "stand_up.dae",
        "stand.dae",
        "talk_a.dae",
        "talk_b.dae",
        "walk.dae"]

# add speed and doghingdirection
actor_list = []
for item in range(6):
    actor = Element("actor", name="actor"+str(item))

    pose = Element("pose")
    #randomly generate position to pose text
    x = str((np.random.rand()-0.5)*4)
    y = str((np.random.rand()-0.5)*4)
    pose.text = x+" "+y+" "+"1.02 0 0 0"
    actor.append(pose)

    skin = Element("skin")
    skin_fn = Element("filename")
    skin_fn.text=random.choice(skin_list)
    skin_scale = Element("scale")
    skin_scale.text = "1"
    skin.append(skin_fn)
    skin.append(skin_scale)
    actor.append(skin)

    animation = Element("animation", name="walking")
    animate_fn = Element("filename")
    animate_fn.text = "walk.dae"
    interpolate_x = Element("interpolate_x")
    interpolate_x.text = "true"
    animate_scale = Element("scale")
    animate_scale.text = "1"
    animation.append(animate_fn)
    animation.append(animate_scale)
    animation.append(interpolate_x)
    actor.append(animation)

    plugin = Element("plugin", name="None", filename=plugin_path)
    speed = Element("speed")
    speed.text = str(1.3)
    socialForce = Element("socialForce")
    socialForce.text = str(SocialForce)
    desiredForce = Element("desiredForce")
    desiredForce.text = str(DesiredForce)
    obstacleForce = Element("obstacleForce")
    obstacleForce.text = str(ObstacleForce)
    dodgingRight = Element("dodgingRight")
    dodgingRight.text = str(np.random.rand() < 0.5).lower()
    target = Element("target")
    x = str((np.random.rand()-0.5)*4)
    y = str((np.random.rand()-0.5)*4)
    target.text =  x+" "+y+" "+"1.02"
    target_weight = Element("target_weight")
    target_weight.text = "1.5"
    obstacle_weight = Element("obstacle_weight")
    obstacle_weight.text = "1.5"
    animation_factor = Element("animation_factor")
    speed_ = str((np.random.rand()*3)+5)
    animation_factor.text = speed_
    ignore_obstacle = Element("ignore_obstacles")
    model_cafe = Element("model")
    model_cafe.text = "caffe"
    model_ground_plane = Element("model")
    model_ground_plane.text = "ground_plane"
    ignore_obstacle.append(model_cafe)
    ignore_obstacle.append(model_ground_plane)
    plugin.append(speed)
    plugin.append(socialForce)
    plugin.append(desiredForce)
    plugin.append(obstacleForce)
    plugin.append(dodgingRight)
    plugin.append(target)
    plugin.append(target_weight)
    plugin.append(obstacle_weight)
    plugin.append(animation_factor)
    plugin.append(ignore_obstacle)
    actor.append(plugin)

    world_.append(actor)

tree_.write(actor_pkg_path+'/worlds/ped_world.world', pretty_print=True, xml_declaration=True, encoding="utf-8")