# Version: 0.12.3
# This is a set of scripts that are used in processing 10x single cell data

import gzip
from scipy import io
from scipy.sparse import csc_matrix
from ast import literal_eval as make_tuple
import pandas as pd
import numpy as np
from copy import deepcopy
import os
import matplotlib.pyplot as plt

def get_version():
    print('0.12.3', 'cleaning vdj code')

def make_dir(directory):
    if not os.path.exists(directory):
        os.mkdir(directory)

def load_crv3_feature_matrix(inst_path, to_csc=True, hto_list=None,
                             drop_default_lane=True, add_lane_to_barcodes=None):

    # Read Barcodes
    ###########################
    filename = inst_path + 'barcodes.tsv.gz'
    f = gzip.open(filename, 'rt')
    lines = f.readlines()
    f.close()

    # if we are adding a lane, then we always want to drop the default cr lane
    if add_lane_to_barcodes is not None:
        drop_default_lane = True

    barcodes = []
    for inst_bc in lines:
        inst_bc = inst_bc.strip().split('\t')

        # remove dash from barcodes if necessary
        if drop_default_lane:
            if '-' in inst_bc[0]:
                inst_bc[0] = inst_bc[0].split('-')[0]
            else:
                print('did not find initial lane from cellranger')

        barcodes.append(inst_bc[0])

    if add_lane_to_barcodes is not None:
        barcodes = [x + '-' + add_lane_to_barcodes for x in barcodes]

    # Load Matrix
    #################
    mat = io.mmread(inst_path + 'matrix.mtx.gz')
    mat_csr = mat.tocsr()

    # Get Indexes of Feature Types
    ##################################
    filename = inst_path + 'features.tsv.gz'
    f = gzip.open(filename, 'rt')
    lines = f.readlines()
    f.close()

    feature_indexes = {}
    feature_lines = {}
    for index in range(len(lines)):

        inst_line = lines[index].strip().split('\t')
        inst_feat = inst_line[2].replace('Gene Expression', 'gex')#
        if hto_list is None:
            inst_feat = inst_feat.replace('Antibody Capture', 'adt').replace('Custom', 'custom')
        else:
            if inst_feat == 'Custom':
                inst_feat = ('hto' if inst_line[0] in hto_list else 'adt')

        if inst_feat not in feature_indexes:
            feature_indexes[inst_feat] = []

        feature_indexes[inst_feat].append(index)

    feature_data = {}

    for inst_feat in feature_indexes:
        feature_data[inst_feat] = {}

        feature_data[inst_feat]['barcodes'] = barcodes

        inst_indexes = feature_indexes[inst_feat]

        # Separate feature lists
        ser_lines = pd.Series(lines)
        ser_lines_found = ser_lines[inst_indexes]
        lines_found = ser_lines_found.get_values().tolist()

        # save feature lines
        feature_lines[inst_feat] = lines_found

        # save as compressed sparse column matrix (for barcode filtering)
        mat_filt = (mat_csr[inst_indexes, :].tocsc() if to_csc else mat_csr[inst_indexes, :])

        feature_data[inst_feat]['mat'] = mat_filt

    # Make unique feature names
    for inst_feat in feature_lines:
        feat_lines = feature_lines[inst_feat]
        feat_lines = [x.strip().split('\t') for x in feat_lines]

        # find non-unique initial feature names (add id later if necessary)
        ini_names = [x[1] for x in feat_lines]

        ini_name_count = pd.Series(ini_names).value_counts()
        duplicate_names = ini_name_count[ini_name_count > 1].index.tolist()

        new_names = [x[1] if x[1] not in duplicate_names else x[1] + '_' + x[0] for x in feat_lines]

        # quick hack to clean up names
        new_names = [x.replace('_TotalSeqB', '') for x in new_names]

        feature_data[inst_feat]['features'] = new_names

    return feature_data

def load_crv2_gene_matrix(inst_path):
    '''
    Loads gene expression data from 10x in sparse matrix format and returns a
    Pandas dataframe
    '''

    import pandas as pd
    from scipy import io
    from scipy import sparse
    from ast import literal_eval as make_tuple

    # matrix
    mat = io.mmread( inst_path + 'matrix.mtx').tocsc()

    # genes
    filename = inst_path + 'genes.tsv'
    f = open(filename, 'r')
    lines_genes = f.readlines()
    f.close()

    # make unique gene names
    #############################
    gene_list = [x.strip().split('\t') for x in lines_genes]

    # find non-unique initial gene names
    ini_names = [x[1] for x in gene_list]

    ini_name_count = pd.Series(ini_names).value_counts()
    duplicate_names = ini_name_count[ini_name_count > 1].index.tolist()
    genes = [x[1] if x[1] not in duplicate_names else x[1] + '_' + x[0] for x in gene_list]

    # barcodes
    filename = inst_path + 'barcodes.tsv'
    f = open(filename, 'r')
    lines = f.readlines()
    f.close()

    cell_barcodes = []
    for inst_bc in lines:
        inst_bc = inst_bc.strip().split('\t')

        # remove dash from barcodes if necessary
        if '-' in inst_bc[0]:
          inst_bc[0] = inst_bc[0].split('-')[0]

        cell_barcodes.append(inst_bc[0])

    # parse tuples if necessary
    try:
        cell_barcodes = [make_tuple(x) for x in cell_barcodes]
    except:
        pass

    try:
        genes = [make_tuple(x) for x in genes]
    except:
        pass

    # generate feature_data
    feature_data = {}
    feature_data['gex'] = {}
    feature_data['gex']['features'] = genes
    feature_data['gex']['barcodes'] = cell_barcodes
    feature_data['gex']['mat'] = mat

    return feature_data

def plot_metadata(meta_cell, metadata_type='gex-umi-sum', logy=True, logx=False, min_umi=0, max_umi=1e9, figsize=(10,5)):

    ser_meta = meta_cell[metadata_type]

    # filter
    ser_meta = ser_meta[ser_meta >= min_umi]
    ser_meta = ser_meta[ser_meta <= max_umi]
    ser_meta = ser_meta.sort_values(ascending=False)

    ser_meta.plot(logy=logy, logx=logx, figsize=figsize)

def plot_umi_levels(feature_data, feature_type='gex', logy=True, logx=False,
                    figsize=(10,5), min_umi=0, max_umi=1e9, zscore_features=False):
    '''
    This function takes a feature data format or dictionary of DataFrames and plots
    UMI levels
    '''

    if 'mat' in feature_data[feature_type]:
        mat_csc = feature_data[feature_type]['mat']

        if zscore_features:
            print('z-scoring feature_data')
            inst_df = pd.DataFrame(data=mat_csc.todense(), columns=feature_data[feature_type]['barcodes'])

            net.load_df(inst_df)
            net.normalize(axis='row', norm_type='zscore')
            inst_df = net.export_df()

            # sort
            ser_sum = inst_df.sum(axis=0).sort_values(ascending=False)

        else:
            # drop cells with fewer than threshold events
            ser_sum = mat_csc.sum(axis=0)
            arr_sum = np.asarray(ser_sum[0,:])

            # sort
            ser_sum = pd.Series(arr_sum[0], index=feature_data[feature_type]['barcodes']).sort_values(ascending=False)

        # filter
        ser_sum = ser_sum[ser_sum >= min_umi]
        ser_sum = ser_sum[ser_sum <= max_umi]

    else:
        inst_df = feature_data[feature_type]

        if zscore_features:
            print('zscore features')
            net.load_df(inst_df)
            net.normalize(axis='row', norm_type='zscore')
            inst_df = net.export_df()

        # sort
        ser_sum = inst_df.sum(axis=0).sort_values(ascending=False)

        # filter
        ser_sum = ser_sum[ser_sum >= min_umi]
        ser_sum = ser_sum[ser_sum <= max_umi]

    ser_sum.plot(logy=logy, logx=logx, figsize=figsize)
    return ser_sum

def filter_barcodes_by_umi(feature_data, feature_type, min_umi=0, max_umi=1e8,
                                      make_sparse=True, zscore_features=False):

    # feature data format
    ########################
    if 'mat' in feature_data[feature_type]:

        # sparse matrix
        ##################
        mat_csc = feature_data[feature_type]['mat']

        if zscore_features:
            print('*** warning, z-scoring not supported in feature_data format')

        # drop barcodes with fewer than threshold UMI
        ser_sum = mat_csc.sum(axis=0)
        arr_sum = np.asarray(ser_sum[0,:])
        ser_sum = pd.Series(arr_sum[0])
        ser_keep = ser_sum[ser_sum >= min_umi]
        ser_keep = ser_keep[ser_keep <= max_umi]

        # these indexes will be used to filter all features
        keep_indexes = ser_keep.index.tolist()

        # filter barcodes
        barcodes = feature_data[feature_type]['barcodes']
        ser_barcodes = pd.Series(barcodes)
        barcodes_filt = ser_barcodes[keep_indexes].get_values()

        # return Dictionary of DataFrames
        filtered_data = {}
        for inst_feat in feature_data:

            if 'meta' not in inst_feat:

                inst_mat = feature_data[inst_feat]['mat']
                mat_filt = inst_mat[:, keep_indexes]
                feature_names = feature_data[inst_feat]['features']

                inst_data = {}
                inst_data['mat'] = mat_filt
                inst_data['barcodes'] = barcodes_filt
                inst_data['features'] = feature_names

                filtered_data[inst_feat] = inst_data

    else:

        # dense matrix
        ###############
        # drop barcodes with fewer than threshold UMI
        inst_df = feature_data[feature_type]

        if zscore_features:
            print('z-scoring features')
            net.load_df(inst_df)
            net.normalize(axis='row', norm_type='zscore')
            inst_df = net.export_df()

        ser_sum = inst_df.sum(axis=0)
        ser_keep = ser_sum[ser_sum >= min_umi]
        ser_keep = ser_keep[ser_keep <= max_umi]
        keep_cols = ser_keep.index.tolist()

        # filter data
        filtered_data = {}
        for inst_feat in feature_data:

            filtered_data[inst_feat] = feature_data[inst_feat][keep_cols]

    return filtered_data

def convert_to_dense(feat_data, df=None):
    # initialize df if necessary
    if df is None:
        df = {}
    for inst_feat in feat_data:
        mat = feat_data[inst_feat]['mat']
        rows = feat_data[inst_feat]['features']
        cols = feat_data[inst_feat]['barcodes']

        dense_mat = mat.todense()
        df[inst_feat] = pd.DataFrame(dense_mat, index=rows, columns=cols)

    return df

def check_feature_data_size(feature_data):
    for inst_feat in feature_data:
        if 'meta' not in inst_feat:
            print(inst_feat)
            print(len(feature_data[inst_feat]['features']), len(feature_data[inst_feat]['barcodes']))
            print(feature_data[inst_feat]['mat'].shape, '\n')

def get_mito_genes(gene_list):
    # Removing Mitochondrial Genes
    ini_mito_list = ['MTRNR2L11', 'MTRF1', 'MTRNR2L12', 'MTRNR2L13', 'MTRF1L',
                     'MTRNR2L6', 'MTRNR2L7','MTRNR2L10', 'MTRNR2L8', 'MTRNR2L5',
                     'MTRNR2L1', 'MTRNR2L3', 'MTRNR2L4']

    list_mito_genes = list(map(lambda x:x.lower(), ini_mito_list))

    found_mito_genes = [x for x in gene_list if 'mt-' == x[:3].lower() or
                 x.split('_')[0].lower() in list_mito_genes]

    return found_mito_genes

def mito_prop_and_suspected_dead(df_gex, meta_cell, mito_thresh=0.9,
                                 plot_mito=True, s=5, alpha=0.2):

    all_genes = df_gex.index.tolist()
    mito_genes = get_mito_genes(all_genes)

    mito_sum = df_gex.loc[mito_genes].sum(axis=0)

    if 'gex-umi-sum' not in meta_cell.columns:
        print('calculating gex-umi-sum, not in meta_cell')
        gex_sum = df_gex.sum(axis=0)
    else:
        gex_sum = meta_cell['gex-umi-sum']

    mito_proportion = mito_sum/gex_sum

    list_mito_dead = []
    cells = mito_proportion.index.tolist()
    for inst_cell in cells:
        inst_mito = mito_proportion[inst_cell]
        if inst_mito >= mito_thresh:
            inst_state = 'dead-cell'
        else:
            inst_state = 'live-cell'
        list_mito_dead.append(inst_state)

    ser_dead = pd.Series(list_mito_dead, index=cells)

    meta_cell['gex-mito-proportion-umi'] = mito_proportion
    meta_cell['dead-cell-mito'] = ser_dead

    if plot_mito:
        # mito_proportion.sort_values(ascending=False).plot()

        all_cells = meta_cell.index.tolist()
        color_list = ['red' if meta_cell.loc[x,'dead-cell-mito'] == 'dead-cell' else 'blue' for x in all_cells]

        meta_cell.plot(kind='scatter',
                       x='gex-umi-sum-ash',
                       y='gex-mito-proportion-umi',
                       alpha=alpha,
                       s=s,
                       figsize=(10,10),
                       c=color_list)

    print('live-cell', ser_dead.value_counts()['live-cell'])
    print('dead-cell', ser_dead.value_counts()['dead-cell'])

    return meta_cell

def set_hto_thresh(df_hto, meta_hto, hto_name, xlim=7, thresh=1, ylim =100):

    if 'hto-threshold' not in meta_hto.columns.tolist():
        ser_thresh = pd.Series(np.nan, index=meta_hto.index)
        meta_hto['hto-threshold'] = ser_thresh

    ser_hto = df_hto.loc[hto_name]

    n, bins, patches = plt.hist(ser_hto, bins=100, range=(0, xlim))

    colors = []
    for inst_bin in bins:
        if inst_bin <= thresh:
            colors.append('red')
        else:
            colors.append('blue')

    # apply the same color for each class to match the map
    for patch, color in zip(patches, colors):
        patch.set_facecolor(color)

    meta_hto.loc[hto_name, 'hto-threshold-ash'] = thresh
    meta_hto.loc[hto_name, 'hto-threshold-umi'] = np.sinh(thresh) * 5

    plt.ylim((0,ylim))

def ini_meta_cell(df):
    list_ser = []

    # look for available data types
    found_types = list(set(['gex', 'adt', 'hto']).intersection(df.keys()))
    for inst_type in found_types:

        # calc umi sum
        inst_ser = df[inst_type].sum(axis=0)
        inst_ser.name = inst_type + '-umi-sum'

        list_ser.append(inst_ser)

    df['meta_cell'] = pd.DataFrame(data=list_ser).transpose()
    if 'gex' in df.keys():
        df_gex = deepcopy(df['gex'])
        df_gex[df_gex >= 1] = 1
        ser_gene_num = df_gex.sum(axis=0)
        df['meta_cell']['num_expressed_genes'] = ser_gene_num
    return df

def meta_cell_gex_wo_mito_ribo(df_gex_ini, meta_cell):

    df_gex = deepcopy(df_gex_ini)

    # calc umi sum
    ser_umi_sum = df_gex.sum(axis=0)

    meta_cell['gex-umi-sum-no-ribo-mito'] = ser_umi_sum

    # count number of measured genes

    df_gex[df_gex >= 1] = 1
    ser_gene_num = df_gex.sum(axis=0)

    meta_cell['num_expressed_genes_no-ribo-mito'] = ser_gene_num

    return meta_cell

def ini_meta_gene(df_gex_ini):

    df_gex = deepcopy(df_gex_ini)

    # Mean UMI
    ser_gene_mean = df_gex.mean(axis=1)
    ser_gene_mean.name = 'mean'

    # Variance UMI
    ser_gene_var = df_gex.mean(axis=1)
    ser_gene_var.name = 'variance'

    # fraction of cells measured
    df_gex[df_gex >= 1] = 1
    ser_gene_meas = df_gex.sum(axis=1)/df_gex.shape[1]
    ser_gene_meas.name = 'fraction of cells measured'

    meta_gene = pd.concat([ser_gene_mean, ser_gene_var, ser_gene_meas], axis=1)

    return meta_gene

def plot_signal_vs_noise(df, alpha=0.25, s=10, hto_range=7, inf_replace=1000):

    fig, axes = plt.subplots(nrows=1, ncols=2)

    list_first = []
    list_second = []
    list_cells = []
    for inst_cell in df.columns.tolist():
        inst_ser = df[inst_cell].sort_values(ascending=False)
        inst_first  = inst_ser.get_values()[0]
        inst_second = inst_ser.get_values()[1]

        list_first.append(inst_first)
        list_second.append(inst_second)
        list_cells.append(inst_cell)

    ser_first  = pd.Series(data=list_first,  index=list_cells, name='first highest HTO')
    ser_second = pd.Series(data=list_second, index=list_cells, name='second highest HTO')

    df_comp = pd.concat([ser_first, ser_second], axis=1).transpose()

    df_comp.transpose().plot(kind='scatter', figsize=(5,5),
                             x='first highest HTO', y='second highest HTO',
                             ylim=(0,hto_range), xlim=(0,hto_range), alpha=alpha, s=s, ax=axes[0])

    sn_ratio = np.log2(df_comp.loc['first highest HTO']/df_comp.loc['second highest HTO'])

    # replace positive infinities with set value
    sn_ratio = sn_ratio.replace(np.Inf, inf_replace)
    sn_ratio.hist(bins=100, ax=axes[1], figsize=(15,7))

    return df_comp, sn_ratio

def filter_ribo_mito_from_gex(df):

    # save avg values to meta_cell
    if 'gex-mito-avg' not in df['meta_cell']:

        df_gex = deepcopy(df['gex'])
        meta_cell = deepcopy(df['meta_cell'])

        all_genes = df_gex.index.tolist()

        ini_genes = deepcopy(all_genes)

        ribo_rpl = [x for x in all_genes if x.lower().startswith('rpl')]
        ribo_rps = [x for x in all_genes if x.lower().startswith('rps')]
        ribo_genes = ribo_rpl + ribo_rps

        # calculate average ribo gene expression
        ser_ribo = df_gex.loc[ribo_genes].mean(axis=0)

        keep_genes = [x for x in all_genes if x not in ribo_genes]

        df_gex = df_gex.loc[keep_genes]

        all_genes = df_gex.index.tolist()

        mito_genes = get_mito_genes(all_genes)

        # calculate average mito gene expression
        ser_mito = df_gex.loc[mito_genes].mean(axis=0)

        keep_genes = [x for x in all_genes if x not in mito_genes]

        # save mito and ribo genes
        mr_genes = sorted(list(set(ini_genes).difference(keep_genes)))
        df_mr = df['gex'].loc[mr_genes]

        # drop mito and ribo genes
        df_gex = df['gex'].loc[keep_genes]

        meta_cell['gex-ribo-avg'] = ser_ribo
        meta_cell['gex-mito-avg'] = ser_mito

        df['gex'] = df_gex
        df['meta_cell'] = meta_cell
        df['gex-mr'] = df_mr
    else:
        print('already filtered mito and ribo genes')

    return df

def add_cats_from_meta(barcodes, df_meta, add_cat_list):
    '''
    Add categories from df_meta.
    '''

    # get metadata of interest (add_cat_list) from barcodes of interest
    df_cats = df_meta.loc[barcodes][add_cat_list]

    # get list of cats
    list_cat_ini = [list(x) for x in df_cats.values]

    # add titles to cats
    list_cat_titles = [ list([str(x) + ': ' + str(y) for x,y in zip(add_cat_list, a)]) for a in list_cat_ini]

    # add barcodes to new columns
    new_cols = [tuple([x] + y) for x,y in zip(barcodes, list_cat_titles)]

    return new_cols

def make_cyto_export(df, num_var_genes=500, inf_replace=10):

    keep_meta_base = ['gex-umi-sum',
                      'gex-num-unique',
                      'gex-mito-proportion-umi',
                      'gex-mito-avg',
                      'gex-ribo-avg',
                      'hto-umi-sum',
                      'hto-sn',
                      'adt-umi-sum'
                      ]

    df_cyto = None

    for inst_type in ['gex', 'adt', 'hto', 'meta_cell']:
        if inst_type in df.keys():
            inst_df = deepcopy(df[inst_type])

            # filter for top var genes
            if inst_type == 'gex':
                keep_var_genes = inst_df.var(axis=1).sort_values(ascending=False).index.tolist()[:num_var_genes]

                inst_df = inst_df.loc[keep_var_genes]

            if 'meta' not in inst_type:
                inst_df.index = [inst_type.upper() + '_' + x for x in inst_df.index.tolist()]

            else:
                keep_meta = [metadata for metadata in keep_meta_base if metadata in inst_df.columns]
                inst_df = inst_df[keep_meta].transpose()
                inst_df.index = [ x.split('-')[0].upper() + '_der_' +
                                 '_'.join( x.split('-')[1:]).replace('num_unique', 'unique_gene_count')
                                 for x in inst_df.index.tolist()]

            print(inst_type, inst_df.shape)

            if df_cyto is None:
                df_cyto = inst_df
            else:
                df_cyto = df_cyto.append(inst_df)

    df_export = df_cyto.transpose()

    cells = df_export.index.tolist()
    index_cells = [str(x/100) for x in range(len(cells))]
    df_export.index = index_cells

    # Add noise
    data_columns = [x for x in df_export.columns if '_der_' not in x]

    not_derived_dataframe_shape = df_export[data_columns].shape

    # center the noise about zero
    rand_mat = np.random.rand(not_derived_dataframe_shape[0], not_derived_dataframe_shape[1]) - 0.5

    df_noise = pd.DataFrame(data=rand_mat, index=index_cells, columns=data_columns).round(2)
    df_export[data_columns] += df_noise

    ser_index = pd.Series(data=index_cells, index=cells)
    df['meta_cell']['Cytobank-Index'] = ser_index

    df_export.index.name = 'cell_index'

    # replace inf and nans
    df_export[df_export == np.inf] = inf_replace
    df_export.fillna(0, inplace=True)

    df['cyto-export'] = df_export

    return df

# # alternate lambda function
# def sum_field(dataframe, field):
#     return dataframe[field].sum(axis=0)

# list_ser_functions = {**{inst_type+'-umi-sum':(lambda y,inst_type=inst_type: sum_field(y,inst_type))\
#                          for inst_type in ['gex', 'adt', 'hto']},
#                      }

# for key,value in list_ser_functions.items():
#     list_ser.append(value(df))
# df['meta_cell'] = pd.DataFrame(data=list_ser).transpose()

def load_prod_vdj(inst_path):
    inst_df = pd.read_csv(inst_path)
    print('all contigs', inst_df.shape)
    ser_prod = inst_df['productive']

    keep_contigs = ser_prod[ser_prod == True].index.tolist()
    inst_df = inst_df.loc[keep_contigs]
    print('productive contigs', inst_df.shape)
    return inst_df

def concat_contig(ser_row):
    inst_v = str(ser_row['v_gene'])
    inst_d = str(ser_row['d_gene'])
    inst_j = str(ser_row['j_gene'])
    inst_c = str(ser_row['c_gene'])
    inst_cdr3 = str(ser_row['cdr3'])

    # do not include c gene in clonotype definition (do not include c_gene)
    inst_contig = inst_v + '_' + inst_d + '_' + inst_j + '_' + inst_cdr3

    return inst_contig

def get_unique_contigs(inst_df):
    '''
    Define contigs as the merge of v, d, j, and cdr3 genes
    Then, find all unique contigs.
    '''
    all_contigs = []
    for inst_index in inst_df.index.tolist():
        ser_row = inst_df.loc[inst_index]
        inst_contig = concat_contig(ser_row)
        all_contigs.append(inst_contig)
    unique_contigs = sorted(list(set(all_contigs)))
    return unique_contigs

def assign_ids_to_contigs(unique_contigs):
    '''
    Generate a unique contig id for all contigs
    return dictionary of contig-to-id and vice versa
    '''
    contig_id_dict = {}
    id_contig_dict = {}
    for inst_index in range(len(unique_contigs)):
        inst_id = 'contig-id-' + str(inst_index)
        inst_contig = unique_contigs[inst_index]
        contig_id_dict[inst_contig] = inst_id
        id_contig_dict[inst_id] = inst_contig

    return contig_id_dict, id_contig_dict

def get_bc_contig_combos(inst_df, contig_id_dict):
    '''
    Loop through the merged (across samples) filtered contigs
    which has one row per contig

    Define the contig (concat vdj genes) and find its unique id
    using contig_id_dict

    Assemble list of contigs associated with each barcode (dict
    with barcode keys)
    '''
    bc_contig_combos = {}
    for inst_row in inst_df.index.tolist():
        ser_row = inst_df.loc[inst_row]
        inst_bc = ser_row['barcode']
        inst_contig = concat_contig(ser_row)
        inst_id = contig_id_dict[inst_contig]

        if inst_bc not in bc_contig_combos:
            bc_contig_combos[inst_bc] = []

        bc_contig_combos[inst_bc].append(inst_id)

    return bc_contig_combos

def generate_new_clonotypes(bc_contig_combos):
    '''
    Define contig combinations as a new set of clones
    Number the new clones (rank by abundance)

    Look up contig combo for each barcode (e.g. clone)
    Look up new clone name for contig comb
    '''

    # find most abundant contig combos (clones)
    contig_id_combos = []
    for inst_bc in bc_contig_combos:
        inst_combo = '_'.join(sorted(bc_contig_combos[inst_bc]))
        contig_id_combos.append(inst_combo)
    ser_combos = pd.Series(contig_id_combos).value_counts()

    # number new clones (contig combos) based on abundance
    inst_id = 1
    combo_clone_dict = {}
    for inst_combo in ser_combos.index.tolist():
        new_clone_name = 'custom-clone-' + str(inst_id)
        combo_clone_dict[inst_combo] = new_clone_name
        inst_id = inst_id + 1

    # make dictionary of new clones for each barcode
    cell_new_clone = {}
    for inst_bc in bc_contig_combos:
        inst_combo = '_'.join(sorted(bc_contig_combos[inst_bc]))
        new_clone = combo_clone_dict[inst_combo]
        cell_new_clone[inst_bc] = new_clone

    return cell_new_clone

def add_uniform_noise(df_ini):
    df = deepcopy(df_ini)
    rows = df.index.tolist()
    cols = df.columns.tolist()

    # generate random matrix
    np.random.seed(99)
    num_rows = df.shape[0]
    num_cols = df.shape[1]
    mat = np.random.rand(num_rows, num_cols)

    # make random noise dataframe centered about zero
    df_noise = pd.DataFrame(data=mat, columns=cols, index=rows).round(2) - 0.5

    df_new = df + df_noise

    return df_new

def filter_sparse_matrix_by_list(feat, feature_type='gex', keep_rows='all', keep_cols='all'):
    '''
    This function filters sparse data by lists of rows/cols.
    filter_by_all_cols is the default because we want all data to have the same number of barcodes/cells
    '''

    feat_filt = deepcopy(feat)

    # get all cols from any feature
    tmp_feat = list(feat_filt.keys())[0]
    cols = feat_filt[tmp_feat]['barcodes']

    # Feature (row) Level Filtering
    #################################
    # apply to single feature
    if isinstance(keep_rows, list):

        # get initial feature list
        rows_orig = feat_filt[feature_type]['features']

        index_dict = dict((value, idx) for idx,value in enumerate(rows_orig))
        rows_idx = [index_dict[x] for x in keep_rows]

        # copy feature data of interest
        inst_mat = deepcopy(feat_filt[feature_type]['mat'])
        inst_mat = inst_mat[rows_idx,:]

        # filter rows for single feature
        feat_filt[feature_type]['barcodes'] = cols
        feat_filt[feature_type]['features'] = keep_rows
        feat_filt[feature_type]['mat'] = inst_mat

    # Cell (col) Level Filtering
    #################################
    # apply to all features
    if isinstance(keep_cols, list):
        index_dict = dict((value, idx) for idx,value in enumerate(cols))
        cols_idx = [index_dict[x] for x in keep_cols]

        # filter all features by columns
        for inst_feat in feat:

            # get initial feature list
            rows_orig = feat_filt[inst_feat]['features']

            inst_mat = deepcopy(feat_filt[inst_feat]['mat'])
            inst_mat = inst_mat[:,cols_idx]

            # filter single feature by columns
            feat_filt[inst_feat]['barcodes'] = keep_cols
            feat_filt[inst_feat]['features'] = rows_orig
            feat_filt[inst_feat]['mat'] = inst_mat

    return feat_filt

def preserve_genes_most_variant(input_df, genes_most_variant=500):
    gene_variance = (input_df['gex']['mat'].power(2)).mean(1) - (
        np.power(input_df['gex']['mat'].mean(1), 2))
    gene_variance_sorted = sorted([(index, variance) for index, variance in enumerate(gene_variance)],
                                  key=(lambda x: x[1]), reverse=True)
    feature_data_gene_variance_filtered = filter_sparse_matrix_by_list(input_df,
                                                                          feature_type='gex',
                                                                          keep_rows=[input_df['gex'][
                                                                                'features'][
                                                                                each_gene_variance[0]] for
                                                                                each_gene_variance in
                                                                                gene_variance_sorted[
                                                                                :genes_most_variant]])

    return feature_data_gene_variance_filtered

def filter_ribo_mito_from_list(all_genes):

    # find ribosomal genes
    ribo_rpl = [x for x in all_genes if 'RPL' in x]
    ribo_rps = [x for x in all_genes if 'RPS' in x]
    ribo_genes = ribo_rpl + ribo_rps

    # Find mitochondrial genes
    list_mito_genes = ['MTRNR2L11', 'MTRF1', 'MTRNR2L12', 'MTRNR2L13', 'MTRF1L', 'MTRNR2L6', 'MTRNR2L7',
                    'MTRNR2L10', 'MTRNR2L8', 'MTRNR2L5', 'MTRNR2L1', 'MTRNR2L3', 'MTRNR2L4']

    mito_genes = [x for x in all_genes if 'MT-' == x[:3] or
                 x.split('_')[0] in list_mito_genes]

    # filter genes
    keep_genes = [x for x in all_genes if x not in ribo_genes]
    keep_genes = [x for x in keep_genes if x not in mito_genes]

    return keep_genes

def calc_feat_sum_and_unique_count_across_cells(feat_data, inst_feat):
    barcodes = (feat_data[inst_feat]['barcodes'] if 'barcodes' in feat_data[inst_feat].keys() else feat_data[inst_feat].columns)
    mat = deepcopy(feat_data[inst_feat]['mat'])

    # sum umi of measured features
    arr_sum = np.asarray(mat.sum(axis=0))[0]
    ser_sum = pd.Series(arr_sum, index=barcodes, name=inst_feat + '-umi-sum')

    # save ash version of umi sum
    ser_sum_ash = np.arcsinh(ser_sum/5)
    ser_sum_ash.name = inst_feat + '-umi-sum-ash'

    # count number of measured features
    mat[mat > 1] = 1
    arr_count = np.asarray(mat.sum(axis=0))[0]
    ser_count = pd.Series(arr_count, index=barcodes, name=inst_feat + '-num-unique')

    inst_df = pd.concat([ser_sum, ser_sum_ash, ser_count], axis=1)

    return inst_df

def sample_meta(df_meta_ini, sample_name):
    list_index = []
    list_data = []

    df_meta = deepcopy(df_meta_ini)

    # proprtion of singlets
    #########################
    ser_cell_per = df_meta['cell-per-bead'].value_counts()

    num_singlets = ser_cell_per.loc['singlet']
    num_total = ser_cell_per.sum()

    # number of singlets
    list_index.append('number-singlets')
    list_data.append(num_singlets)

    # get singlets only
    df_meta = df_meta[df_meta['cell-per-bead'] == 'singlet']

    # proportion of dead cells
    ##############################
    ser_dead = df_meta['dead-cell-mito'].value_counts()
    prop_dead = 1 - ser_dead['live-cell'] / ser_dead.sum()

    list_index.append('proportion-dead')
    list_data.append(prop_dead)

    # assemble initial metadata series
    ser_meta_ini = pd.Series(list_data, index=list_index)

    # Calculate average metadata
    meta_list = ['gex-umi-sum', 'gex-num-unique', 'gex-mito-proportion-umi', 'gex-ribo-avg', 'gex-mito-avg']
    ser_meta_mean = df_meta[meta_list].mean()

    ser_meta = pd.concat([ser_meta_ini, ser_meta_mean])
    ser_meta.name = sample_name

    return ser_meta

def merge_lanes(lane_dirs, merge_dir, data_types=['gex', 'adt', 'hto', 'meta_cell'],
    return_df=True):

    lane_dirs = sorted(lane_dirs)

    if return_df:
        df = {}

    for inst_type in data_types:

        df_merge = None

        print('\n' + inst_type + '\n----------------')

        # collect data
        for inst_dir in lane_dirs:

            inst_lane = inst_dir.split('/')[-1]
            inst_file = inst_dir + '/' + inst_type + '.parquet'
            if os.path.exists(inst_file):

                inst_df = pd.read_parquet(inst_file)

                # if meta add lane category
                if 'meta' in inst_type:
                    ser_lane = pd.Series(data=inst_lane, index=inst_df.index.tolist())
                    inst_df['Lane_10x'] = ser_lane

                print(inst_lane, inst_df.shape)

                # merge on the fly
                if df_merge is None:
                    df_merge = deepcopy(inst_df)
                else:
                    if 'meta' in inst_type:
                        df_merge = pd.concat([df_merge, inst_df], axis=0, sort=True)
                    else:
                        df_merge = pd.concat([df_merge, inst_df], axis=1)
                print('df_merge', df_merge.shape)

        print('merged', inst_type, df_merge.shape)

        # sort columns and rows
        cols = sorted(df_merge.columns.tolist())
        rows = sorted(df_merge.index.tolist())
        df_merge = df_merge.loc[rows, cols]

        df_merge.to_parquet(merge_dir + '/' + inst_type + '.parquet')

        if return_df:
            # save to dictionary
            df[inst_type] = df_merge

    if return_df:
        return df

def load_kb_vel_feature_matrix(inst_path, inst_sample, to_csc=True):

    # Load barcodes
    #################
    filename = inst_path + inst_sample + '.barcodes.txt'
    f = open(filename, 'r')
    lines = f.readlines()
    f.close()

    barcodes = []
    for inst_bc in lines:
        inst_bc = inst_bc.strip().split('\t')

        barcodes.append(inst_bc[0])

    # Load genes
    #################
    filename = inst_path + inst_sample + '.genes.txt'
    f = open(filename, 'r')
    lines = f.readlines()
    f.close()

    genes = []
    for inst_gene in lines:
        inst_gene = inst_gene.strip().split('\t')

        genes.append(inst_gene[0])

    # Load Matrix
    #################
    mat = io.mmread(inst_path + inst_sample +'.mtx').transpose()
    mat = mat.tocsc()

    print(len(genes), len(barcodes), mat.shape)

    feature_data = {}
    feature_data['gex'] = {}
    feature_data['gex']['mat'] = mat
    feature_data['gex']['features'] = genes
    feature_data['gex']['barcodes'] = barcodes

    return feature_data

def drop_debris_gex_hto_ash(df_meta, gex_ash_thresh=None, hto_ash_thresh=None):

    # convert to non-ash values for interpretation
    gex_thresh = np.sinh(gex_ash_thresh) * 5
    hto_thresh = np.sinh(hto_ash_thresh) * 5

    # initialize everything to red (debris)
    ser_colors = pd.Series('red', index=df_meta.index.tolist())

    color_list = []
    ser_gex = df_meta['gex-umi-sum-ash']
    gex_pass = ser_gex[ser_gex >= gex_ash_thresh].index.tolist()

    ser_hto = df_meta['hto-umi-sum-ash']
    hto_pass = ser_hto[ser_hto >= hto_ash_thresh].index.tolist()

    print('gex thresh UMI-ash: ', gex_ash_thresh, ' gex thresh UMI: ', gex_thresh.round(0))
    print('hto thresh UMI-ash: ', hto_ash_thresh, ' hto thresh UMI: ', hto_thresh.round(0))

    print('gex keep: ', len(gex_pass))
    print('hto keep: ', len(hto_pass))

    keep_barcodes = list(set(gex_pass + hto_pass))
    print('keep_barcodes: ', len(keep_barcodes))

    ser_colors[keep_barcodes] = 'blue'
    color_list = ser_colors.get_values()

    df_meta.plot(kind='scatter', x='gex-umi-sum-ash', y='hto-umi-sum-ash',
                         s=1, alpha=0.5, figsize=(10,10),
                         ylim=(0,10), xlim=(0,10), c=color_list)

    return keep_barcodes

def make_ct_list(df_hto, meta_hto):

    '''assign cells to debris/singlet/multiplet based on threshold only (ash normalized)'''

    ser_list = []
    df_hto_tmp = deepcopy(np.arcsinh(df_hto/5))

    for inst_row in df_hto_tmp.index.tolist():

        # get data for a HTO
        inst_ser = deepcopy(df_hto_tmp.loc[inst_row])

        # load threshold level for this HTO
        inst_thresh = meta_hto.loc[inst_row, 'hto-threshold-ash']

        # binarize HTO values about threshold
        inst_ser[inst_ser < inst_thresh] = 0
        inst_ser[inst_ser >= inst_thresh] = 1

        # assemble list of series to make dataframe later
        ser_list.append(inst_ser)

        # find cells that are positive for this HTO
        pos_hto = inst_ser[inst_ser==1].index.tolist()

    # generate binarized dataframe
    df_binary = pd.concat(ser_list, axis=1).transpose()

    # find singlets
    ser_sum = df_binary.sum(axis=0)
    ct_list = {}
    ct_list['debris']    = ser_sum[ser_sum == 0].index.tolist()
    ct_list['singlet']   = ser_sum[ser_sum == 1].index.tolist()
    ct_list['multiplet'] = ser_sum[ser_sum > 1].index.tolist()

    return ct_list

def calc_s2n_and_s2t(df_hto, meta_hto, meta_cell, inf_replace):

    # signal-to-noise refers to the highest vs second highest HTO
    #
    # signal-to-threshold refers to the highest HTO vs the
    # manually set threshold for that HTO

    # do not use ash data, calc signal to noise on HTO UMI data
    # find the highest hto
    ser_max_hto = df_hto.idxmax(axis=0)
    meta_cell['hto-max-name'] = ser_max_hto

    list_first = []
    list_second = []
    list_thresh = []
    list_cells = []

    for inst_cell in df_hto.columns.tolist():
        inst_ser = df_hto[inst_cell].sort_values(ascending=False)
        inst_first  = inst_ser.get_values()[0]
        inst_second = inst_ser.get_values()[1]

        inst_max_hto_name = ser_max_hto[inst_cell]
        inst_max_hto_thresh = meta_hto['hto-threshold-umi'][inst_max_hto_name]

        list_first.append(inst_first)
        list_second.append(inst_second)
        list_thresh.append(inst_max_hto_thresh)
        list_cells.append(inst_cell)

    ser_first  = pd.Series(data=list_first,  index=list_cells, name='first highest HTO')
    ser_second = pd.Series(data=list_second, index=list_cells, name='second highest HTO')
    ser_thresh = pd.Series(data=list_thresh, index=list_cells, name='threshold HTO')

    # df_comp = pd.concat([ser_first, ser_second, ser_thresh], axis=1).transpose()

    meta_cell['hto-max-umi'] = ser_first

    # calc signal-to-noise
    ###########################
    # sn_ratio = df_comp.loc['first highest HTO']/df_comp.loc['second highest HTO']
    sn_ratio = ser_first/ser_second
    meta_cell['hto-sn'] = sn_ratio
    # replace infinities with large number
    sn_ratio = sn_ratio.replace(np.Inf, inf_replace)

    # calc signal-to-threshold
    ###########################
    # sn_ratio = df_comp.loc['first highest HTO']/df_comp.loc['threshold HTO']
    st_ratio = ser_first/ser_thresh
    meta_cell['hto-st'] = st_ratio
    # replace infinities with large number
    st_ratio = st_ratio.replace(np.Inf, inf_replace)

    return meta_cell

def assign_htos(df_hto, meta_hto, meta_cell, sn_thresh, inf_replace=1000, perform_sn_adjustment=True):

    # assign cells to debris/singlet/multiplet based on threshold only (ash normalized)
    ######################################################################################
    ct_list = make_ct_list(df_hto, meta_hto)

    print('De-hash distributions: HTO Threshold Only')
    print('-----------------------------------------')
    print('debris', len(ct_list['debris']))
    print('singlet', len(ct_list['singlet']))
    print('multiplet', len(ct_list['multiplet']))

    # initialize dehash-thresh: debris/singlet/multiplet based on ct_list
    if 'dehash-thresh' not in meta_cell.columns.tolist():
        ser_type = pd.Series(np.nan, index=meta_cell.index)
        meta_cell['dehash-thresh'] = ser_type

    # save dehash-thresh
    ######################
    for inst_type in ct_list:
        meta_cell.loc[ct_list[inst_type], 'dehash-thresh'] = inst_type

    # Calc signal-to-noise and signal-to-threshold ratios
    ######################################################
    meta_cell = calc_s2n_and_s2t(df_hto, meta_hto, meta_cell, inf_replace)

    # Assign Cells to Samples based on Threshold Alone
    ####################################################
    list_samples = []
    for inst_cell in meta_cell.index.tolist():

        inst_type = meta_cell.loc[inst_cell, 'dehash-thresh']

        if inst_type == 'singlet':
            inst_hto = meta_cell.loc[inst_cell, 'hto-max-name']
            inst_sample = meta_hto.loc[inst_hto]['Sample']
        else:
            inst_sample = 'N.A.'

        list_samples.append(inst_sample)

    ser_sample = pd.Series(list_samples, index=meta_cell.index.tolist())
    meta_cell['Sample-thresh'] = ser_sample

    if perform_sn_adjustment:
        # Assign Cells to Samples Based on Signal to Noise Adjustment
        cells = meta_cell.index.tolist()
        for inst_cell in cells:
            inst_type = meta_cell.loc[inst_cell, 'dehash-thresh']
            inst_sn = meta_cell.loc[inst_cell, 'hto-sn']
            inst_max_hto = meta_cell.loc[inst_cell, 'hto-max-name']

            # change singlet to multiplet if low sn
            if inst_type == 'singlet':
                if inst_sn < sn_thresh['singlets']:
                    # convert to multiplet
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'multiplet'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = 'N.A.'
                else:
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'singlet'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = meta_hto.loc[inst_max_hto]['Sample']
            elif inst_type == 'debris':
                if inst_sn >= sn_thresh['debris']:
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'singlet'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = meta_hto.loc[inst_max_hto]['Sample']
                else:
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'debris'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = 'N.A.'
            elif inst_type == 'multiplet':
                if inst_sn >= sn_thresh['multiplets']:
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'singlet'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = meta_hto.loc[inst_max_hto]['Sample']
                else:
                    meta_cell.loc[inst_cell, 'dehash-thresh-sn'] = 'multiplet'
                    meta_cell.loc[inst_cell, 'Sample-thresh-sn'] = 'N.A.'

        ser_counts = meta_cell['dehash-thresh-sn'].value_counts()

        print('De-hash distributions: HTO Threshold and SN')
        print('--------------------------------------------')
        print('debris', ser_counts['debris'])
        print('singlet', ser_counts['singlet'])
        print('multiplet', ser_counts['multiplet'])

    return meta_cell

def plot_hto_sn_vs_gex_umi(df):

    if 'hto-sn-ash' not in df['meta_cell']:
        df['meta_cell']['hto-sn-ash'] = np.arcsinh(df['meta_cell']['hto-sn']/5)

    color_dict = {
        "singlet":'blue',
        "debris":'red',
        "multiplet":'yellow',
    }

    list_dehash = list(df['meta_cell']['cell-per-bead'].get_values())
    color_list = [color_dict[x] for x in list_dehash]

    df['meta_cell'].plot(kind='scatter',
                         x='gex-umi-sum-ash',
                         y='hto-sn-ash', alpha=0.25, s=10, figsize=(10,10), c=color_list, ylim=(0,5))

def load_s3_parquet(bucket_path, filename, cols=None):

    import s3fs
    import pyarrow.parquet as pq
    fs = s3fs.S3FileSystem()

    if cols == None:
        df = pq.read_table(bucket_path + filename, use_pandas_metadata=True, filesystem=fs).to_pandas()
    else:
        df = pq.read_table(bucket_path + filename, use_pandas_metadata=True, filesystem=fs, columns=cols).to_pandas()

    return df

def add_filter(cat_title, inst_cat_search, cat_filter_list=None):
    '''
    Add a single filtering step; e.g. filter for barcodes with
    this category of this category type.

    To do: support exclusion of category by adding third element to tuple.
    '''
    if cat_filter_list == None:
        cat_filter_list = []

    if type(inst_cat_search) is list:
        inst_cat_search = tuple(inst_cat_search)

    cat_filter_list.append((cat_title, inst_cat_search))
    return cat_filter_list

def filter_meta_using_cat_filter_list(df_meta_ini, cat_filter_list):
    df_meta = deepcopy(df_meta_ini)

    for cat_tuple in cat_filter_list:
        inst_cat_type = cat_tuple[0]
        inst_cat_search = cat_tuple[1]

        if type(inst_cat_search) is not tuple:
            # find indexes of barcodes that match requested caetgory
            inst_cat = inst_cat_search
            cat_ser = df_meta[inst_cat_type]
            found_barcodes = cat_ser[cat_ser == inst_cat].index.tolist()
        else:
            # find indexes of barcodes that match requested categories
            found_barcodes = []
            for inst_cat in inst_cat_search:
                cat_ser = df_meta[inst_cat_type]
                inst_found = cat_ser[cat_ser == inst_cat].index.tolist()
                found_barcodes.extend(inst_found)

        # apply progressive filters to metadata
        df_meta = df_meta.loc[found_barcodes]

    return df_meta

def set_gex_debris_thresh(meta_cell, xlim=7, ylim=100, thresh=1):

    ser_gex_ash = meta_cell['gex-umi-sum-ash']

    n, bins, patches = plt.hist(ser_gex_ash, bins=100, range=(0, xlim))

    colors = []
    for inst_bin in bins:
        if inst_bin <= thresh:
            colors.append('red')
        else:
            colors.append('blue')

    # apply the same color for each class to match the map
    for patch, color in zip(patches, colors):
        patch.set_facecolor(color)

    plt.ylim((0,ylim))

    keep_barcodes = ser_gex_ash[ser_gex_ash >= thresh].index.tolist()

    print('gex-ash-umi thresh', thresh, '; gex-umi thresh', np.sinh(thresh) * 5)
    print('keeping', len(keep_barcodes), 'cells')

    return keep_barcodes

def sort_all_dataframes(df):
    for inst_type in df:
        print('sorting', inst_type)
        inst_df = df[inst_type]
        # sort columns and rows
        cols = sorted(inst_df.columns.tolist())
        rows = sorted(inst_df.index.tolist())
        inst_df = inst_df.loc[rows, cols]

    return df


import time
import numpy as np
import tensorflow as tf
from resnet101_deep_metric import Resnet101_Deep_Metric
from Mybase.solver import Solver

def test():

    mdl = Resnet101_Deep_Metric(cls_num=81, reg=1e-4, drp=0.5, typ=tf.float32)
    sov = Solver(mdl,
                 opm_cfg={
                     "decay_rule": "fixed",
                     #"optim_rule": "adam",
                     #"lr_base": 0.0002,
                     "optim_rule": "momentum",
                     "lr_base":  0.00005, #0.004
                     "momentum": 0.9,
                     #"optim_rule": "adadelta",
                     #"lr_base": 0.01,
                 },
                 use_gpu     = True,
                 gpu_lst     = '0',
                 bat_siz     = 1,
                 tra_num     = 8000,
                 val_num     = 40,
                 epc_num     = 10000,
                 prt_ena     = True,
                 itr_per_prt = 20,
                 tst_num     = None,
                 tst_shw     = True,
                 tst_sav     = True,
                 mdl_nam     = 'model.ckpt',
                 mdl_dir     = 'Mybase/Model',
                 log_dir     = 'Mybase/logdata',
                 dat_dir     = 'Mybase/datasets',
                 mov_ave_dca = 0.99,
                 epc_per_dca = 1)
    #print("TRAINING...")
    #sov.train()
    print("TESTING...")
    sov.test()
    sov.display_detections()
    #sov.show_loss_acc()

test()


# HQ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
# HQ X
# HQ X   quippy: Python interface to QUIP atomistic simulation library
# HQ X
# HQ X   Copyright James Kermode 2020
# HQ X
# HQ X   These portions of the source code are released under the GNU General
# HQ X   Public License, version 2, http://www.gnu.org/copyleft/gpl.html
# HQ X
# HQ X   If you would like to license the source code under different terms,
# HQ X   please contact James Kermode, james.kermode@gmail.com
# HQ X
# HQ X   When using this software, please cite the following reference:
# HQ X
# HQ X   http://www.jrkermode.co.uk/quippy
# HQ X
# HQ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

import unittest

import ase
import numpy as np
import quippytest
from quippy.potential import Potential

xml_string = """
<SW_params n_types="1">
      <per_type_data type="1" atomic_num="14" />
      <per_pair_data atnum_i="14" atnum_j="14" AA="7.049556277" BB="0.6022245584"
      p="4" q="0" a="1.80" sigma="2.0951" eps="2.1675" />
      <per_triplet_data atnum_c="14" atnum_j="14" atnum_k="14"
      lambda="21.0" gamma="1.20" eps="2.1675" />
</SW_params>
"""

class Test_NeighbourList(quippytest.QuippyTestCase):
    """
    Tests the NeighbourList with varying PBC settings as accessed by the wrapper.
    """

    def setUp(self):
        p0 = []
        for i in range(2):
            for j in range(2):
                for k in range(2):
                    p0.append([i * 2 - 1,
                               j * 2 - 1,
                               k * 2 - 1])
        self.p0 = np.array(p0)
        self.sw_pot = Potential('IP SW', param_str=xml_string)

    def test_pbc(self):
        for cell in [tuple([4] * 3),
                     tuple([10] * 3),
                     tuple([20] * 3),
                     ((4, 4, 0), (4, 0, 4), (0, 4, 4)),
                     (4, 10, 10)]:
            for pbc in [tuple([True] * 3),
                        tuple([False] * 3),
                        (True, False, False),
                        (False, True, False)]:
                init_energy = None
                for offset_dir in [(1, 1, 1),
                                   (1, 0, 0)]:
                    for lat_offset in [-5.0, -1, -0.5, 0.0, 0.5, 1.0, 5.0]:
                        at = ase.Atoms(symbols=['Si'] * len(self.p0),
                                       positions=self.p0,
                                       cell=cell, pbc=pbc)
                        at.calc = self.sw_pot
                        offset_v = np.dot(offset_dir, at.get_cell()) * lat_offset
                        at.positions += offset_v
                        if init_energy is None:
                            init_energy = at.get_potential_energy()
                        self.assertAlmostEqual(init_energy,
                                               at.get_potential_energy())

    def test_large_offset_no_pbc(self):
        cell = [10] * 3
        pbc = [False] * 3
        offset = 5e6
        at = ase.Atoms(symbols=['Si'] * len(self.p0),
                       positions=self.p0 + offset,
                       cell=cell,
                       pbc=pbc)
        at.calc = self.sw_pot
        e = at.get_potential_energy()
        self.assertAlmostEqual(e, 20.863001205176083)

if __name__ == '__main__':
    unittest.main()


import os
import sys
import time
import argparse
import ast
import logging
import numpy as np
import paddle.fluid as fluid

from model import ResNet3D
from reader import KineticsReader
from config import parse_config, merge_configs, print_configs

logging.root.handlers = []
FORMAT = '[%(levelname)s: %(filename)s: %(lineno)4d]: %(message)s'
logging.basicConfig(filename='logger.log', level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)

def parse_args():
    parser = argparse.ArgumentParser("Paddle Video train script")
    parser.add_argument(
        '--model_name',
        type=str,
        default='res3d',
        help='name of model to train.')
    parser.add_argument(
        '--config',
        type=str,
        default='configs/res3d.txt',
        help='path to config file of model')
    parser.add_argument(
        '--batch_size',
        type=int,
        default=None,
        help='training batch size. None to use config file setting.')
    parser.add_argument(
        '--learning_rate',
        type=float,
        default=None,
        help='learning rate use for training. None to use config file setting.')
    parser.add_argument(
        '--pretrain',
        type=str,
        default=None,
        help='path to pretrain weights. None to use default weights path in  ~/.paddle/weights.'
    )
    parser.add_argument(
        '--use_gpu',
        type=ast.literal_eval,
        default=False,
        help='default use gpu.')
    parser.add_argument(
        '--epoch',
        type=int,
        default=100,
        help='epoch number, 0 for read from config file')
    parser.add_argument(
        '--save_dir',
        type=str,
        default='checkpoints_models',
        help='directory name to save train snapshoot')
    args = parser.parse_args()
    return args

def train(args):
    # parse config
    place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
    with fluid.dygraph.guard(place):
        config = parse_config(args.config)
        train_config = merge_configs(config, 'train', vars(args))
        print_configs(train_config, 'Train')

        # train_model = TSN1.TSNResNet('TSN', train_config['MODEL']['num_layers'],
        #                             train_config['MODEL']['num_classes'],
        #                             train_config['MODEL']['seg_num'], 0.00002)
        train_model = ResNet3D.ResNet3D('res3d', train_config['MODEL']['num_layers'],
                                    train_config['MODEL']['num_classes'],
                                    train_config['MODEL']['seg_num'], 0.00002)
        opt = fluid.optimizer.Momentum(0.001, 0.9, parameter_list=train_model.parameters())

        if args.pretrain:
            # 加载上一次训练的模型，继续训练
            model, _ = fluid.dygraph.load_dygraph(args.save_dir + '/res3d_model')
            train_model.load_dict(model)

        # build model
        if not os.path.exists(args.save_dir):
            os.makedirs(args.save_dir)

        # get reader
        train_config.TRAIN.batch_size = train_config.TRAIN.batch_size
        train_reader = KineticsReader(args.model_name.upper(), 'train', train_config).create_reader()

        epochs = args.epoch or train_model.epoch_num()
        for i in range(epochs):
            for batch_id, data in enumerate(train_reader()):
                dy_x_data = np.array([x[0] for x in data]).astype('float32')
                y_data = np.array([[x[1]] for x in data]).astype('int64')

                img = fluid.dygraph.to_variable(dy_x_data)
                # print(img.shape)     ##################测试
                label = fluid.dygraph.to_variable(y_data)
                label.stop_gradient = True

                out, acc = train_model(img, label)

                loss = fluid.layers.cross_entropy(out, label)
                avg_loss = fluid.layers.mean(loss)

                avg_loss.backward()

                opt.minimize(avg_loss)
                train_model.clear_gradients()

                if batch_id % 12 == 0:
                    logger.info("Loss at epoch {} step {}: {}, acc: {}".format(i, batch_id, avg_loss.numpy(), acc.numpy()))
                    print("Loss at epoch {} step {}: {}, acc: {}".format(i, batch_id, avg_loss.numpy(), acc.numpy()))
                    fluid.dygraph.save_dygraph(train_model.state_dict(), args.save_dir + '/res3d_model')
        logger.info("Final loss: {}".format(avg_loss.numpy()))
        print("Final loss: {}".format(avg_loss.numpy()))

if __name__ == "__main__":
    args = parse_args()
    # check whether the installed paddle is compiled with GPU
    logger.info(args)

    train(args)


#!/usr/bin/env python

import rospy
from std_msgs.msg import Header
from human_motion.msg import Skeleto
import tf
import sys

import copy
import numpy as np

sys.path.insert(0, 'Animation')
import data_utils
import playAnimation
import viz

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

parent = np.array([0, 1, 2, 3, 4, 5, 1, 7, 8, 9,10, 1,12,13,14,15,13,
                17,18,19,20,21,20,23,13,25,26,27,28,29,28,31])-1

offset = np.array(
    [0.000000,0.000000,0.000000,
    -132.948591,0.000000,0.000000,
    0.000000,-442.894612,0.000000,
    0.000000,-454.206447,0.000000,
    0.000000,0.000000,162.767078,
    0.000000,0.000000,74.999437,
    132.948826,0.000000,0.000000,
    0.000000,-442.894413,0.000000,
    0.000000,-454.206590,0.000000,
    0.000000,0.000000,162.767426,
    0.000000,0.000000,74.999948,
    0.000000,0.100000,0.000000,
    0.000000,233.383263,0.000000,
    0.000000,257.077681,0.000000,
    0.000000,121.134938,0.000000,
    0.000000,115.002227,0.000000,
    0.000000,257.077681,0.000000,
    0.000000,151.034226,0.000000,
    0.000000,278.882773,0.000000,
    0.000000,251.733451,0.000000,
    0.000000,0.000000,0.000000,
    0.000000,0.000000,99.999627,
    0.000000,100.000188,0.000000,
    0.000000,0.000000,0.000000,
    0.000000,257.077681,0.000000,
    0.000000,151.031437,0.000000,
    0.000000,278.892924,0.000000,
    0.000000,251.728680,0.000000,
    0.000000,0.000000,0.000000,
    0.000000,0.000000,99.999888,
    0.000000,137.499922,0.000000,
    0.000000,0.000000,0.000000])
offset = offset.reshape(-1,3)

expmapInd = np.split(np.arange(4,100)-1,32)

g_tf_prefix = ''

def motion_visualize_callback(data):

    rospy.loginfo('%s: %s' % (data.header.seq, data.header.stamp))

    skeleto = data.skeleto
    angles = np.array(skeleto)

    assert len(angles) == 99

    # Structure that indicates parents for each joint
    njoints = 32

    stamp = data.header.stamp
    rate = rospy.Rate(1000)
    br = tf.TransformBroadcaster()
    while rospy.get_rostime() < stamp:
        rate.sleep()
    for i in range(njoints):
        tran = offset[i]
        rot = np.eye(3)
        if (i > 0):
            rm = np.eye(4)
            rot = data_utils.expmap2rotmat(angles[expmapInd[i]])
            rm[0:3, 0:3] = np.linalg.inv(rot)
            quat = tf.transformations.quaternion_from_matrix(rm)
            br.sendTransform(tran / 1000, quat, stamp, '%sj%s' % (g_tf_prefix, i), '%sj%s' % (g_tf_prefix, parent[i]))
        else:
            quat = tf.transformations.quaternion_from_euler(np.pi/2, 0, 0)
            br.sendTransform(tran / 1000, quat, stamp, '%sj%s' % (g_tf_prefix, i), 'skeleto')

def main():
    global g_tf_prefix
    rospy.init_node('motion_tf_broadcaster', anonymous=True)
    g_tf_prefix = rospy.get_param('tf_prefix', g_tf_prefix)
    print('tf_prefix: ', g_tf_prefix)
    rospy.Subscriber('motion', Skeleto, motion_visualize_callback, queue_size=100)
    rospy.spin()

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass


# Copyright 2016-2021 The Van Valen Lab at the California Institute of
# Technology (Caltech), with support from the Paul Allen Family Foundation,
# Google, & National Institutes of Health (NIH) under Grant U24CA224309-01.
# All rights reserved.
#
# Licensed under a modified Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.github.com/vanvalenlab/deepcell-tracking/LICENSE
#
# The Work provided may be used for non-commercial academic purposes only.
# For any other use of the Work, including commercial use, please contact:
# vanvalenlab@gmail.com
#
# Neither the name of Caltech nor the names of its contributors may be used
# to endorse or promote products derived from this software without specific
# prior written permission.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""A cell tracking class capable of extending labels across sequential frames."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import json
import logging
import os
import pathlib
import tarfile
import tempfile
import timeit

import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist
from skimage.measure import regionprops

import networkx as nx
import pandas as pd

from deepcell_tracking.utils import clean_up_annotations
from deepcell_tracking.utils import get_max_cells
from deepcell_tracking.utils import normalize_adj_matrix
from deepcell_tracking.utils import get_image_features
from deepcell_tracking.utils import save_trk

class CellTracker(object):  # pylint: disable=useless-object-inheritance
    """Solves the linear assingment problem to build a cell lineage graph.

    Args:
        movie (np.array): raw time series movie of cells.
        annotation (np.array): the labeled cell movie.
        model (keras.Model): tracking model to determine if two cells are the
            same, different, or parent/daughter.
        features (list): list of strings for the features to use.
        crop_dim (int): crop size for the appearance feature.
        death (float): paramter used to fill the death matrix in the LAP,
            (top right of the cost matrix).
        birth (float): paramter used to fill the birth matrix in the LAP,
            (bottom left of the cost matrix).
        division (float): probability threshold for assigning daughter cells.
        distance_threshold (int): maximum distance to compare cells with the model.
        track_length (int): the track length used for the model.
        neighborhood_scale_size (int): neighborhood feature size to pass to the
            model.
        neighborhood_true_size (int): original size of the neighborhood feature
            which will be scaled down to neighborhood_scale_size.
        dtype (str): data type for features, can be 'float32', 'float16', etc.
        data_format (str): determines the order of the channel axis,
            one of 'channels_first' and 'channels_last'.
    """

    def __init__(self,
                 movie,
                 annotation,
                 tracking_model,
                 neighborhood_encoder=None,
                 distance_threshold=64,
                 appearance_dim=32,
                 death=0.99,
                 birth=0.99,
                 division=0.9,
                 track_length=5,
                 embedding_axis=0,
                 dtype='float32',
                 data_format='channels_last'):

        if not len(movie.shape) == 4 or not len(annotation.shape) == 4:
            raise ValueError('Input data and labels but be rank 4 '
                             '(frames, x, y, channels).  Got {} and {}.'.format(
                                 len(movie.shape), len(annotation.shape)))

        if not movie.shape[:-1] == annotation.shape[:-1]:
            raise ValueError('Input data and labels should have the same shape'
                             ' except for the channel dimension.  Got {} and '
                             '{}'.format(movie.shape, annotation.shape))

        if data_format not in {'channels_first', 'channels_last'}:
            raise ValueError('The `data_format` argument must be one of '
                             '"channels_first", "channels_last". Received: ' +
                             str(data_format))

        self.X = copy.copy(movie)
        self.y = copy.copy(annotation)
        self.tracks = {}

        self.neighborhood_encoder = neighborhood_encoder
        self.tracking_model = tracking_model
        self.distance_threshold = distance_threshold
        self.appearance_dim = appearance_dim
        self.death = death
        self.birth = birth
        self.division = division
        self.dtype = dtype
        self.track_length = track_length
        self.embedding_axis = embedding_axis

        self.a_matrix = []
        self.c_matrix = []
        self.assignments = []

        self.data_format = data_format
        self.channel_axis = 0 if data_format == 'channels_first' else -1
        self.time_axis = 1 if data_format == 'channels_first' else 0
        self.logger = logging.getLogger(str(self.__class__.__name__))

        # Clean up annotations
        self.y = clean_up_annotations(self.y, data_format=self.data_format)

        # Accounting for 0 (background) label with 0-indexing for tracks
        self.id_to_idx = {}  # int: int mapping
        self.idx_to_id = {}  # (frame, cell_idx): cell_id mapping

        # Establish features for every instance of every cell in the movie
        adj_matrices, appearances, morphologies, centroids = self._est_feats()

        # Compute embeddings for every instance of every cell in the movie
        embeddings = self._get_neighborhood_embeddings(
            appearances=appearances,
            morphologies=morphologies,
            centroids=centroids,
            adj_matrices=adj_matrices)

        # TODO: immutable dict for safety? these values should never change.
        self.features = {
            'embedding': embeddings,
            'centroid': centroids,
        }

    def _get_frame(self, tensor, frame):
        """Helper function for fetching a frame of a tensor.

        Useful for avoiding duplication of the data_format conditional.

        Args:
            tensor (np.array): The 3D tensor to slice.
            frame (int): The frame to slice out of the tensor.

        Returns:
            np.array: the 2D slice of the 3D tensor.
        """
        if self.data_format == 'channels_first':
            return tensor[:, frame]
        return tensor[frame]

    def _get_cells_in_frame(self, frame):
        """Find the labels of cells in the given frame.

        Args:
            frame (int): Frame of interest.

        Returns:
            list: All cell labels in the frame.
        """
        cells = np.unique(self._get_frame(self.y, frame))
        cells = np.delete(cells, np.where(cells == 0))  # remove the background
        return list(cells)

    def _est_feats(self):
        """
        Extract the relevant features from the label movie
        Appearance, morphologies, centroids, and adjacency matrices
        """
        max_cells = get_max_cells(self.y)
        n_frames = self.X.shape[0]
        n_channels = self.X.shape[-1]

        appearances = np.zeros((n_frames,
                                max_cells,
                                self.appearance_dim,
                                self.appearance_dim,
                                n_channels), dtype=np.float32)

        morphologies = np.zeros((n_frames, max_cells, 3), dtype=np.float32)

        centroids = np.zeros((n_frames, max_cells, 2), dtype=np.float32)

        adj_matrix = np.zeros((n_frames, max_cells, max_cells),
                              dtype=np.float32)

        for frame in range(n_frames):

            frame_features = get_image_features(
                self.X[frame], self.y[frame],
                appearance_dim=self.appearance_dim)

            for cell_idx, cell_id in enumerate(frame_features['labels']):
                self.id_to_idx[cell_id] = cell_idx
                self.idx_to_id[(frame, cell_idx)] = cell_id

            num_tracks = len(frame_features['labels'])
            centroids[frame, :num_tracks] = frame_features['centroids']
            morphologies[frame, :num_tracks] = frame_features['morphologies']
            appearances[frame, :num_tracks] = frame_features['appearances']

            cent = centroids[frame]
            distance = cdist(cent, cent, metric='euclidean') < self.distance_threshold
            adj_matrix[frame] = distance.astype('float32')

        # Normalize adj matrix
        norm_adj_matrices = normalize_adj_matrix(adj_matrix)

        return norm_adj_matrices, appearances, morphologies, centroids

    def _get_neighborhood_embeddings(self, appearances, morphologies,
                                     centroids, adj_matrices):
        """Compute the embeddings using the neighborhood encoder"""
        # Build input dictionary for neighborhood encoder model
        inputs = {
            'encoder_app_input': appearances,
            'encoder_morph_input': morphologies,
            'encoder_centroid_input': centroids,
            'encoder_adj_input': adj_matrices,
        }

        # TODO: current model doesnt organize outputs according to ordered list
        #       patching with embedding_axis
        embeddings = self.neighborhood_encoder.predict(inputs)[self.embedding_axis]
        embeddings = np.array(embeddings)
        return embeddings

    def _validate_feature_name(self, feature_name):
        if feature_name not in self.features:
            raise ValueError('{} is an invalid feature name. '
                             'Use one of embedding or centroid'.format(
                                 feature_name))

    def _get_feature(self, frame, cell_id, feature_name='embedding'):
        """Get the feature for a cell in the frame"""
        self._validate_feature_name(feature_name)
        cell_idx = self.id_to_idx[cell_id]
        return self.features[feature_name][frame, cell_idx, :]

    def _get_frame_features(self, frame, feature_name='embedding'):
        """Get the feature for the specified cells in a frame"""
        self._validate_feature_name(feature_name)

        cells_in_frame = self._get_cells_in_frame(frame)

        frame_features = {}
        for cell_id in cells_in_frame:
            f = self._get_feature(frame, cell_id, feature_name=feature_name)
            frame_features[cell_id] = f
        return frame_features

    def _create_new_track(self, frame, old_label):
        """
        This function creates new tracks
        """
        track_id = len(self.tracks)
        new_label = track_id + 1
        embedding = self._get_feature(frame, old_label, feature_name='embedding')
        centroid = self._get_feature(frame, old_label, feature_name='centroid')

        embedding = np.expand_dims(embedding, axis=0)
        centroid = np.expand_dims(centroid, axis=0)

        self.tracks[track_id] = {
            'label': new_label,
            'frames': [frame],
            'frame_labels': [old_label],
            'daughters': [],
            'capped': False,
            'frame_div': None,
            'parent': None,
            'embedding': embedding,
            'centroid': centroid
        }

        if frame > 0 and np.any(self._get_frame(self.y, frame) == new_label):
            raise Exception('new_label already in annotated frame and frame > 0')

        if self.data_format == 'channels_first':
            self.y[:, frame][self.y[:, frame] == old_label] = new_label
        else:
            self.y[frame][self.y[frame] == old_label] = new_label

    def _initialize_tracks(self):
        """Intialize the tracks. Tracks are stored in a dictionary.
        """
        frame = 0  # initial frame
        cell_ids = self._get_cells_in_frame(frame)

        for cell_id in cell_ids:
            self._create_new_track(frame, cell_id)

        # Start a tracked label array
        # TODO: This could be a source of error!
        # we are assigning a pointer not instantiating a new object
        self.y_tracked = self.y[[frame]].astype('int32')

    def _fetch_tracked_features(self, before_frame=None, feature_name='embedding'):
        """Get feature data from each tracked frame less than before_frame.

        Args:
            before_frame (int, optional): The maximum frame to from which to
                fetch feature data.
            feature_name (str): Name of feature to fetch from tracked data.

        Returns:
            dict: dictionary of feature name to np.array of feature data.
        """
        self._validate_feature_name(feature_name)

        if before_frame is None:
            before_frame = self.X.shape[0] + 1  # all frames

        track_valid_frames = ((n, [f for f in d['frames'] if f < before_frame])
                              for n, d in self.tracks.items())
        tracks_with_frames = [(n, f) for n, f in track_valid_frames if len(f) > 0]

        tracked_features = {}
        for i, (n, valid_frames) in enumerate(tracks_with_frames):
            frame_dict = {frame: j for j, frame in enumerate(valid_frames)}
            frames = valid_frames[-self.track_length:]

            if len(frames) != self.track_length:
                # Pad the the frames with the last frame if not enough
                num_missing = self.track_length - len(frames)
                frames = frames + [frames[-1]] * num_missing

            # Get the feature data from the identified frames
            fetched = self.tracks[n][feature_name][[frame_dict[f] for f in frames]]
            tracked_features[i] = fetched

        return tracked_features

    def _build_cost_matrix(self, assignment_matrix):
        """Build the full cost matrix based on the assignment_matrix.

        Args:
            assignment_matrix (np.array): assignment_matrix.

        Returns:
            numpy.array: cost matrix.
        """
        # Initialize cost matrix
        num_tracks, num_cells = assignment_matrix.shape
        cost_matrix = np.zeros((num_tracks + num_cells,) * 2, dtype=self.dtype)

        # assignment matrix - top left
        cost_matrix[0:num_tracks, 0:num_cells] = assignment_matrix

        # birth matrix - bottom left
        birth_diagonal = np.array([self.birth] * num_cells)
        birth_matrix = np.zeros((num_cells, num_cells), dtype=self.dtype)
        birth_matrix = np.diag(birth_diagonal) + np.ones(birth_matrix.shape)
        birth_matrix = birth_matrix - np.eye(num_cells)
        cost_matrix[num_tracks:, 0:num_cells] = birth_matrix

        # death matrix - top right
        death_matrix = np.ones((num_tracks, num_tracks), dtype=self.dtype)
        death_matrix = self.death * np.eye(num_tracks) + death_matrix
        death_matrix = death_matrix - np.eye(num_tracks)
        cost_matrix[0:num_tracks, num_cells:] = death_matrix

        # mordor matrix - bottom right
        cost_matrix[num_tracks:, num_cells:] = assignment_matrix.T
        return cost_matrix

    def _get_cost_matrix(self, frame):
        """Use the model predictions to build an assignment matrix to be solved.

        Args:
            frame (int): The frame with cells to assign.

        Returns:
            tuple: the assignment matrix and the predictions used to build it.
        """
        inputs = {}
        relevant_tracks = []
        for feature_name in self.features:
            # Get the embeddings for previously tracked cells
            current_feature = self._fetch_tracked_features(
                before_frame=frame, feature_name=feature_name)

            # Get the embeddings for the current frame
            future_feature = self._get_frame_features(
                frame=frame, feature_name=feature_name)

            if not relevant_tracks:
                for track_id in current_feature:
                    relevant_tracks.append(track_id)

            # Convert from dict to arrays
            current_feature_arr = np.stack([
                current_feature[k] for k in current_feature
            ], axis=1)  # time axis already included
            future_feature_arr = np.stack([
                future_feature[k] for k in future_feature
            ], axis=0)

            # Add time dimension
            future_feature_arr = np.expand_dims(future_feature_arr, axis=0)

            # Add batch dimension
            current_feature_arr = np.expand_dims(current_feature_arr, axis=0)
            future_feature_arr = np.expand_dims(future_feature_arr, axis=0)

            # Add feature to inputs
            # model expects "current_embeddings" but feature_name is "embedding"
            # TODO: this name is hardcoded based on a model from deepcell-tf
            inputs['current_{}s'.format(feature_name)] = current_feature_arr
            inputs['future_{}s'.format(feature_name)] = future_feature_arr

        t = timeit.default_timer()

        # Perform inference
        predictions = self.tracking_model.predict(inputs)
        predictions = predictions[0, 0, ...]  # Remove the batch and time dimension
        assignment_matrix = 1 - predictions[..., 0]

        for track_id in relevant_tracks:
            if self.tracks[track_id]['capped']:
                assignment_matrix[track_id, :] = 1

        self.a_matrix.append(predictions)

        # Assemble full cost matrix
        cost_matrix = self._build_cost_matrix(assignment_matrix)
        self.logger.debug('Built cost matrix for frame %s in %s s.',
                          frame, timeit.default_timer() - t)
        self.c_matrix.append(cost_matrix)

        predictions_dict = {}
        predictions_dict['predictions'] = predictions
        predictions_dict['track_ids'] = relevant_tracks

        return cost_matrix, predictions_dict

    def _update_tracks(self, assignments, frame, predictions):
        """Update the graph based on the assignment matrix for the frame.

        Use the assignment matrix to determine if each cell in the frame.
        belongs to a track or is a daughter of a track, and updates the graph
        accordingly.

        Args:
            assignments (np.array): completed assignment matrix used to assign
                cells to existing tracks.
            frame (int): the frame of cells to assign.
            predictions (dict): dictionary of trackID-cellID combination,
                and the probability they are the same cell.
        """
        t = timeit.default_timer()
        cells_in_frame = self._get_cells_in_frame(frame)

        # Number of lables present in the current frame (needed to build cost matrix)
        y_tracked_update = np.zeros((1, self.y.shape[1], self.y.shape[2], 1), dtype='int32')

        self.assignments.append(assignments)

        for a in range(assignments.shape[0]):
            track, cell = assignments[a]

            # map the index from the LAP assignment to the cell label
            try:
                cell_id = cells_in_frame[cell]
            except IndexError:
                # cell has "died" or is a shadow assignment
                # no assignment should be made
                continue

            cell_embedding = self._get_feature(frame, cell_id, feature_name='embedding')
            cell_centroid = self._get_feature(frame, cell_id, feature_name='centroid')

            cell_embedding = np.expand_dims(cell_embedding, axis=0)
            cell_centroid = np.expand_dims(cell_centroid, axis=0)

            if track in self.tracks:  # Add cell and frame to track
                self.tracks[track]['frames'].append(frame)
                self.tracks[track]['frame_labels'].append(cell_id)
                self.tracks[track]['embedding'] = np.concatenate([
                    self.tracks[track]['embedding'],
                    cell_embedding
                ], axis=0)
                self.tracks[track]['centroid'] = np.concatenate([
                    self.tracks[track]['centroid'],
                    cell_centroid
                ], axis=0)

                # Labels and indices differ by 1
                y_tracked_update[self.y[[frame]] == cell_id] = track + 1
                self.y[frame][self.y[frame] == cell_id] = track + 1

            else:  # Create a new track if there was a birth
                self._create_new_track(frame, cell_id)
                new_track_id = max(self.tracks)
                new_label = new_track_id + 1
                self.logger.info('Created new track for cell %s.', new_label)

                # See if the new track has a parent
                parent = self._get_parent(frame, cell_id, predictions)
                if parent is not None:
                    self.logger.info('Detected division! Cell %s is daughter '
                                     'of cell %s.', new_label, parent + 1)
                    self.tracks[new_track_id]['parent'] = parent
                    self.tracks[parent]['daughters'].append(new_track_id)
                else:
                    self.tracks[new_track_id]['parent'] = None

                y_tracked_update[self.y[[frame]] == new_label] = new_track_id + 1
                self.y[frame][self.y[frame] == new_label] = new_track_id + 1

        # Check and make sure cells that divided did not get assigned to the same cell
        for track in range(len(self.tracks)):
            if not self.tracks[track]['daughters']:
                continue  # Filter out cells that have not divided

            # Cap tracks for any divided cells
            if not self.tracks[track]['capped']:
                self.tracks[track]['frame_div'] = int(frame)
                self.tracks[track]['capped'] = True

            try:
                frame_idx = self.tracks[track]['frames'].index(frame)
            except ValueError:
                continue  # Filter out tracks that are not in the frame

            # Create new track
            new_track_id = len(self.tracks)
            new_label = new_track_id + 1
            self._create_new_track(frame, self.tracks[track]['frame_labels'][-1])
            self.tracks[new_track_id]['parent'] = track

            # Remove features and frame from old track
            del self.tracks[track]['frames'][frame_idx]
            del self.tracks[track]['frame_labels'][frame_idx]
            self.tracks[track]['embedding'] = np.delete(
                self.tracks[track]['embedding'],
                frame_idx, axis=0
            )
            self.tracks[track]['centroid'] = np.delete(
                self.tracks[track]['centroid'],
                frame_idx, axis=0
            )
            self.tracks[track]['daughters'].append(new_track_id)

            # Change y_tracked_update
            old_label = self.tracks[track]['label']
            y_tracked_update[self.y[[frame]] == old_label] = new_label
            # TODO: Having y and y_tracked is redundant. only one should be changed
            self.y[frame][self.y[frame] == old_label] = new_label

        # Update the tracked label array
        self.y_tracked = np.concatenate([self.y_tracked, y_tracked_update], axis=0)
        self.logger.debug('Updated tracks for frame %s in %s s.',
                          frame, timeit.default_timer() - t)

    def _get_parent(self, frame, cell, predictions):
        """Searches the tracks for the parent of a given cell.

        Args:
            frame (int): the frame the cell appears in.
            cell (int): the label of the cell in the frame.
            predictions (dict): dictionary of trackID-cellID combination,
                and the probability they are the same cell.

        Returns:
            int: The parent cell's id or None if no parent exists.
        """
        parent_id = None
        max_prob = self.division

        track_ids = predictions['track_ids']
        predictions = predictions['predictions']

        for track_idx, track_id in enumerate(track_ids):
            # Make sure capped tracks can't be assigned parents
            if self.tracks[track_id]['capped']:
                continue

            for cell_idx in range(predictions.shape[1]):
                cell_id = self.idx_to_id[(frame, cell_idx)]

                if cell_id == cell:
                    # Do not call a newly-appeared sibling of "cell" a parent
                    if self.tracks[track_id]['frames'] == [frame]:
                        continue

                    # probability cell is part of the track
                    prob = predictions[track_idx, cell_idx, 2]

                    if prob > max_prob:
                        parent_id, max_prob = track_id, prob

        return parent_id

    def _track_frame(self, frame):
        """Inner function for tracking each frame"""
        t = timeit.default_timer()
        self.logger.info('Tracking frame %s', frame)

        cost_matrix, predictions = self._get_cost_matrix(frame)

        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        assignments = np.stack([row_ind, col_ind], axis=1)

        self._update_tracks(assignments, frame, predictions)
        self.logger.info('Tracked frame %s in %s s.',
                         frame, timeit.default_timer() - t)

    def track_cells(self):
        """Tracks all of the cells in every frame.
        """
        start = timeit.default_timer()
        self._initialize_tracks()

        for frame in range(1, self.X.shape[self.time_axis]):
            self._track_frame(frame)

        self.logger.info('Tracked all %s frames in %s s.',
                         self.X.shape[self.time_axis],
                         timeit.default_timer() - start)

    def _track_review_dict(self):
        def process(key, track_item):
            if track_item is None:
                return track_item
            if key == 'daughters':
                return list(map(lambda x: x + 1, track_item))
            elif key == 'parent':
                return track_item + 1
            else:
                return track_item

        track_keys = ['label', 'frames', 'daughters', 'capped', 'frame_div', 'parent']

        return {'tracks': {track['label']: {key: process(key, track[key]) for key in track_keys}
                           for _, track in self.tracks.items()},
                'X': self.X,
                'y': self.y,
                'y_tracked': self.y_tracked}

    def dataframe(self, **kwargs):
        """Returns a dataframe of the tracked cells with lineage.
        Uses only the cell labels not the ids.
        _track_cells must be called first!
        """
        # possible kwargs are extra_columns
        extra_columns = ['cell_type', 'set', 'part', 'montage']
        track_columns = ['label', 'daughters', 'frame_div']

        incorrect_args = set(kwargs) - set(extra_columns)
        if incorrect_args:
            raise ValueError('Invalid argument {}'.format(incorrect_args.pop()))

        # filter extra_columns by the ones we passed in
        extra_columns = [c for c in extra_columns if c in kwargs]

        # extra_columns are the same for every row, cache the values
        extra_column_vals = [kwargs[c] for c in extra_columns if c in kwargs]

        # fill the dataframe
        data = []
        for cell_id, track in self.tracks.items():
            data.append(extra_column_vals + [track[c] for c in track_columns])
        dataframe = pd.DataFrame(data, columns=extra_columns + track_columns)

        # daughters contains track_id not labels
        dataframe['daughters'] = dataframe['daughters'].apply(
            lambda d: [self.tracks[x]['label'] for x in d])

        return dataframe

    def postprocess(self, filename=None, time_excl=9):
        """Use graph postprocessing to eliminate false positive division errors
        using a graph-based detection method. False positive errors are when a
        cell is noted as a daughter of itself before the actual division occurs.
        If a filename is passed, save the state of the cell tracker to a .trk
        ('track') file. time_excl is the minimum number of frames expected to
        exist between legitimate divisions
        """

        # Load data
        track_review_dict = self._track_review_dict()

        # Prep data
        tracked = track_review_dict['y_tracked'].astype('uint16')
        lineage = track_review_dict['tracks']

        # Identify false positives (FPs)
        G = self._track_to_graph(lineage)
        FPs = self._flag_false_pos(G, time_excl)
        FPs_candidates = sorted(FPs.items(), key=lambda v: int(v[0].split('_')[1]))
        FPs_sorted = self._review_candidate_nodes(FPs_candidates)

        # If FPs exist, use the results to correct
        while len(FPs_sorted) != 0:

            lineage, tracked = self._remove_false_pos(lineage, tracked, FPs_sorted[0])
            G = self._track_to_graph(lineage)
            FPs = self._flag_false_pos(G, time_excl)
            FPs_candidates = sorted(FPs.items(), key=lambda v: int(v[0].split('_')[1]))
            FPs_sorted = self._review_candidate_nodes(FPs_candidates)

        # Make sure the assignment is correct
        track_review_dict['y_tracked'] = tracked
        track_review_dict['tracks'] = lineage

        # Save information to a track file file if requested
        if filename is not None:
            self.dump(filename, track_review_dict)

        return track_review_dict

    def dump(self, filename, track_review_dict=None):
        """Writes the state of the cell tracker to a .trk ('track') file.
        Includes raw & tracked images, and a lineage.json for parent/daughter
        information.
        """
        if not track_review_dict:
            track_review_dict = self._track_review_dict()

        filename = pathlib.Path(filename)

        if filename.suffix != '.trk':
            filename = filename.with_suffix('.trk')

        filename = str(filename)

        save_trk(filename=filename,
                 lineage=track_review_dict['tracks'],
                 raw=track_review_dict['X'],
                 tracked=track_review_dict['y_tracked'])

    def _track_to_graph(self, tracks):
        """Create a graph from the lineage information"""
        Dattr = {}
        edges = pd.DataFrame()

        for L in tracks.values():
            # Calculate node ids
            cellid = ['{}_{}'.format(L['label'], f) for f in L['frames']]
            # Add edges from cell ids
            edges = edges.append(pd.DataFrame({'source': cellid[0:-1],
                                               'target': cellid[1:]}))

            # Collect any division attributes
            if L['frame_div'] is not None:
                Dattr['{}_{}'.format(L['label'], L['frame_div'] - 1)] = {'division': True}

            # Create any daughter-parent edges
            if L['parent'] is not None:
                source = '{}_{}'.format(L['parent'], min(L['frames']) - 1)
                target = '{}_{}'.format(L['label'], min(L['frames']))
                edges = edges.append(pd.DataFrame({'source': [source],
                                                   'target': [target]}))

        G = nx.from_pandas_edgelist(edges, source='source', target='target')
        nx.set_node_attributes(G, Dattr)
        return G

    def _flag_false_pos(self, G, time_excl):
        """Examine graph for false positive nodes
        """

        # TODO: Current implementation may eliminate some divisions at the edge of the frame -
        #       Further research needed

        # Identify false positive nodes
        node_fix = []
        for g in (G.subgraph(c) for c in nx.connected_components(G)):
            div_nodes = [n for n, d in g.nodes(data=True) if d.get('division')]
            if len(div_nodes) > 1:
                for nd in div_nodes:
                    if g.degree(nd) == 2:
                        # Check how close suspected FP is to other known divisions

                        keep_div = True
                        for div_nd in div_nodes:
                            if div_nd != nd:
                                time_spacing = abs(int(nd.split('_')[1]) -
                                                   int(div_nd.split('_')[1]))
                                # If division is sufficiently far away
                                # we should exclude it from FP list
                                if time_spacing > time_excl:
                                    keep_div = False

                        if keep_div is True:
                            node_fix.append(nd)

        # Add supplementary information for each false positive
        D = {}
        for node in node_fix:
            D[node] = {
                'false positive': node,
                'neighbors': list(G.neighbors(node)),
                'connected lineages': set([int(node.split('_')[0])
                                           for node in nx.node_connected_component(G, node)])
            }

        return D

    def _review_candidate_nodes(self, FPs_candidates):
        """ review candidate false positive nodes and remove any errant degree 2 nodes.
        """
        FPs_presort = {}
        # review candidate false positive nodes and remove any errant degree 2 nodes
        for candidate_node in FPs_candidates:
            node = candidate_node[0]
            node_info = candidate_node[1]

            neighbors = []  # structure will be [(neighbor1, frame), (neighbor2,frame)]
            for neighbor in node_info['neighbors']:
                neighbor_label = int(neighbor.split('_')[0])
                neighbor_frame = int(neighbor.split('_')[1])
                neighbors.append((neighbor_label, neighbor_frame))

            # if this cell only exists in one frame (and then it divides) but its 2 neighbors
            # both exist in the same frame it will be a degree 2 node but not be a false positive
            if neighbors[0][1] != neighbors[1][1]:
                FPs_presort[node] = node_info

        FPs_sorted = sorted(FPs_presort.items(), key=lambda v: int(v[0].split('_')[1]))

        return FPs_sorted

    def _remove_false_pos(self, lineage, tracked, FP_info):
        """ Remove nodes that have been identified as false positive divisions.
        """
        node = FP_info[0]
        node_info = FP_info[1]

        fp_label = int(node.split('_')[0])
        fp_frame = int(node.split('_')[1])

        neighbors = []  # structure will be [(neighbor1, frame), (neighbor2,frame)]
        for neighbor in node_info['neighbors']:
            neighbor_label = int(neighbor.split('_')[0])
            neighbor_frame = int(neighbor.split('_')[1])
            neighbors.append((neighbor_label, neighbor_frame))

        # Verify that the FP node only 2 neighbors - 1 before it and one after it
        if len(neighbors) == 2:
            # order the neighbors such that the time (frame order) is respected
            if neighbors[0][1] > neighbors[1][1]:
                temp = neighbors[0]
                neighbors[0] = neighbors[1]
                neighbors[1] = temp

            # Decide which labels to extend and which to remove

            # Neighbor_1 has same label as fp - the actual division hasnt occurred yet
            if fp_label == neighbors[0][0]:
                # The model mistakenly identified a division before the actual division occurred
                label_to_remove = neighbors[1][0]
                label_to_extend = neighbors[0][0]

                # Give all of the errant divisions information to the correct track
                lineage[label_to_extend]['frames'].extend(lineage[label_to_remove]['frames'])
                lineage[label_to_extend]['daughters'] = lineage[label_to_remove]['daughters']
                lineage[label_to_extend]['frame_div'] = lineage[label_to_remove]['frame_div']

                # Adjust the parent information for the actual daughters
                daughter_labels = lineage[label_to_remove]['daughters']
                for daughter in daughter_labels:
                    lineage[daughter]['parent'] = lineage[label_to_remove]['parent']

                # Remove the errant node from the annotated images
                channel = 0  # These images should only have one channel
                for frame in lineage[label_to_remove]['frames']:
                    label_loc = np.where(tracked[frame, :, :, channel] == label_to_remove)
                    tracked[frame, :, :, channel][label_loc] = label_to_extend

                # Remove the errant node from the lineage
                del lineage[label_to_remove]

            # Neighbor_2 has same label as fp - the actual division ocurred &
            # the model mistakenly allowed another
            # elif fp_label == neighbors[1][0]:
                # The model mistakenly identified a division after
                # the actual division occurred
                # label_to_remove = fp_label

            # Neither neighbor has same label as fp - the actual division
            # ocurred & the model mistakenly allowed another
            else:
                # The model mistakenly identified a division after the actual division occurred
                label_to_remove = fp_label
                label_to_extend = neighbors[1][0]

                # Give all of the errant divisions information to the correct track
                lineage[label_to_extend]['frames'] = \
                    lineage[fp_label]['frames'] + lineage[label_to_extend]['frames']
                lineage[label_to_extend]['parent'] = lineage[fp_label]['parent']

                # Adjust the parent information for the actual daughter
                parent_label = lineage[fp_label]['parent']
                for d_idx, daughter in enumerate(lineage[parent_label]['daughters']):
                    if daughter == fp_label:
                        lineage[parent_label]['daughters'][d_idx] = label_to_extend

                # Remove the errant node from the annotated images
                channel = 0  # These images should only have one channel
                for frame in lineage[label_to_remove]['frames']:
                    label_loc = np.where(tracked[frame, :, :, channel] == label_to_remove)
                    tracked[frame, :, :, channel][label_loc] = label_to_extend

                # Remove the errant node
                del lineage[label_to_remove]

        else:
            self.logger.error('Error: More than 2 neighbor nodes')

        return lineage, tracked