text:
import os
import pandas as pd
import matplotlib
import click
from matplotlib import rc
from matplotlib import lines
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import gc

plt.rcParams['svg.fonttype'] = 'none'

image_path = 'reference_files/primirna_background.tiff'
im = plt.imread(image_path)

SMALL_SIZE = 18
MEDIUM_SIZE = 20
BIGGER_SIZE = 22

rc('font', size=MEDIUM_SIZE)          # controls default text sizes
rc('axes', titlesize=MEDIUM_SIZE)     # fontsize of the axes title
rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
rc('xtick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
rc('ytick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

def create_plot(data_df, output_name, mutations=0, genes=0, mirna_type='both'):

    df_plot_5, loop_value = prepare_data_5p(data_df)

    df_plot_3 = prepare_data_3p(data_df)

    max_value = max(df_plot_5['pos'].max(), df_plot_3['pos'].max())

    rc("pdf", fonttype=42)
    sns.set_style(style='white')

    palette = {'flanking-5': 'grey',
               'flanking-3': 'grey',
               'pre-seed': '#5481A6',
               'seed': 'darkblue',
               'post-seed': '#5481A6',
               'loop': 'grey',
               'silent-pre': '#5481A6',
               'silent-post': '#5481A6',
               'silent-seed': '#5481A6'
               }

    fig = plt.figure(figsize=(25, 10))
    ax = fig.add_axes([0, 0, 1, 1])

    ax.axis('off')
    plt.imshow(im)
    plt.xticks([])
    plt.yticks([])

    y_min, y_max = ax.get_ylim()
    x_min, x_max = ax.get_xlim()

    plt.text(x_max * 0.92, y_min * 1, str(mutations), horizontalalignment='center',
             verticalalignment='center', fontdict={'size': '42'})
    plt.text(x_max * 0.92, y_min * 1.09, str(genes), horizontalalignment='center',
             verticalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.89, x_max * 0.95], [y_min * 1.04, y_min * 1.04], lw=0.5, color='black')
    ax.add_line(line)
    line.set_clip_on(False)

    plt.text(x_max * 0.24, y_min * -0.09, 'flanking region', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.02, x_max * 0.46], [y_min * -0.07, y_min * -0.07], lw=0.5, color='grey', alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    plt.text(x_max * 0.65, y_min * -0.09, 'miRNA', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.47, x_max * 0.83], [y_min * -0.07, y_min * -0.07], lw=0.5, color='#5481A6',
                        alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    plt.text(x_max * 0.91, y_min * -0.09, 'loop', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.84, x_max * 0.98], [y_min * -0.07, y_min * -0.07], lw=0.5, color='grey', alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    if mirna_type != '3p':
        plt.text(x_max * 0.54, y_min * 0, 'seed', horizontalalignment='center', fontdict={'size': '42'})
        line = lines.Line2D([x_max * 0.49, x_max * 0.59], [y_min * 0.03, y_min * 0.03], lw=0.5, color='darkblue',
                            alpha=0.8)
        ax.add_line(line)
        line.set_clip_on(False)

    if mirna_type != '5p':
        line = lines.Line2D([x_max * 0.675, x_max * 0.775], [y_min * 1.125, y_min * 1.125], lw=0.5, color='darkblue',
                            alpha=0.8)
        ax.add_line(line)
        line.set_clip_on(False)
        plt.text(x_max * 0.725, y_min * 1.195, 'seed', horizontalalignment='center', fontdict={'size': '42'})

    if loop_value > 0:
        plt.text(
            x_max * 0.92, y_min * 0.6, '+ {} loop\nmutations'.format(loop_value),
            horizontalalignment='center',
            verticalalignment='center',
            fontdict={'size': '30'}
        )

    a = plt.axes([.058, .52, .82, .35])

    hue_order = ['flanking-5', 'pre-seed', 'seed', 'post-seed', 'loop']

    if mirna_type == '3p':
        hue_order = ['flanking-5', 'silent-pre', 'silent-seed', 'silent-post', 'loop']

        df_plot_5['type'] = df_plot_5['type'].apply(
            lambda x: 'silent-seed' if x == 'seed' else (
                'silent-pre' if x == 'pre-seed' else (
                    'silent-post' if x == 'post-seed' else x
                )
            )
        )
    labels = [str(x) for x in range(-25, 0)] + \
             [str(x) for x in range(1, 23)] + ['+' + str(x) for x in range(1, 4)]

    labels = ['L' if x == '+4' else x for x in labels]

    ax = sns.barplot(x="from_start", y="pos", hue="type",
                     data=df_plot_5, dodge=False,
                     hue_order=hue_order,
                     palette=palette,
                     ax=a)

    for loc in ['right', 'top', 'left', 'bottom']:
        ax.spines[loc].set_visible(False)

    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.legend(loc='upper right', ncol=2)

    plt.setp(ax.get_legend().get_texts(), fontsize='18')
    plt.setp(ax.get_legend().get_title(), fontsize='22')
    plt.grid(b=True, which='major', axis='y', color='lightgrey',
             linestyle='-', linewidth=0.75, zorder=2, alpha=0.5)

    ax.set_xticklabels(labels)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    if max_value > 5:
        plt.yticks(np.arange(0, max_value + 1, np.floor(max_value/3)))
        plot_limit = max_value + 2
    else:
        plt.yticks(np.arange(0, max_value + 1, 1))
        plot_limit = max_value + 1

    ax.set_ylim([0, plot_limit])
    ax.xaxis.tick_bottom()
    for label in ax.get_xticklabels():
        if label.get_text() not in ['-25', '-20', '-15', '-10', '-5', '1', '5', '10', '15', '20',
                                    '+1']:
            label.set_visible(False)

    new_width = 0.35
    for patch in ax.patches:
        current_width = patch.get_width()
        diff = current_width - new_width

        # we change the bar width
        patch.set_width(new_width)

        # we recenter the bar
        patch.set_x(patch.get_x() + diff * .5)

    ax.tick_params(axis='both', which='major', pad=8, width=0.5)

    plt.setp(ax.patches, linewidth=0)

    ax.get_legend().remove()

    b = plt.axes([.025, .02, .82, .35], facecolor='w')
    hue_order = ['flanking-3', 'pre-seed', 'seed', 'post-seed', 'loop']
    labels = ['+' + str(x) for x in range(1, 26)][::-1] + [str(x) for x in range(1, 23)][::-1] + \
             ['-' + str(x) for x in range(1, 4)]

    if mirna_type == '5p':
        hue_order = ['flanking-3', 'silent-pre', 'silent-seed', 'silent-post', 'loop']

        df_plot_3['type'] = df_plot_3['type'].apply(
            lambda x: 'silent-seed' if x == 'seed' else (
                'silent-pre' if x == 'pre-seed' else(
                    'silent-post' if x == 'post-seed' else x
                )
            )
        )

    ax = sns.barplot(x="from_start", y="pos", hue="type",
                     data=df_plot_3, dodge=False,
                     hue_order=hue_order,
                     palette=palette,
                     ax=b)

    for loc in ['right', 'top', 'left', 'bottom']:
        ax.spines[loc].set_visible(False)

    ax.set_xlabel('')
    ax.set_ylabel('')

    plt.grid(b=True, which='major', axis='y', color='lightgrey',
             linestyle='-', linewidth=0.75, zorder=1, alpha=0.5)

    ax.set_xticklabels(labels)

    if max_value > 5:
        plt.yticks(np.arange(0, max_value + 1, np.floor(max_value/3)))
        plot_limit = max_value + 2
    else:
        plt.yticks(np.arange(0, max_value + 1, 1))
        plot_limit = max_value + 1

    ax.get_legend().remove()
    ax.set_ylim([plot_limit, 0])

    ax.xaxis.tick_top()

    for tick in ax.xaxis.get_major_ticks():
        tick.label2.set_fontsize(30)

    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    for label in ax.get_xticklabels():
        if label.get_text() not in ['+25', '+20', '+15', '+10', '+5', '+1', '1', '5', '10', '15', '20',
                                    '-5']:
            label.set_visible(False)

    new_width = 0.35
    for patch in ax.patches:
        current_width = patch.get_width()
        diff = current_width - new_width

        # we change the bar width
        patch.set_width(new_width)

        # we recenter the bar
        patch.set_x(patch.get_x() + diff * .5)

    plt.setp(ax.patches, linewidth=0)

    plt.savefig(output_name, format='svg', dpi=300, transparent=True, bbox_inches='tight')
    plt.cla()
    plt.clf()
    plt.close("all")
    plt.close(fig)
    gc.collect()

def prepare_data_5p(df_temp):

    add_loop = df_temp[(df_temp['arm'] == 'loop') & ((df_temp['from_start'] >= 4) & (df_temp['from_end'] <= -4))]

    add_loop_value = add_loop.shape[0]

    # add_loop_df = pd.DataFrame([['loop', 'loop', 51, add_loop_value]], columns=['arm', 'type',
    #                                                                             'from_start',
    #                                                                             'pos'])

    dataframe = df_temp[(df_temp['arm'] == '5p') |
                        ((df_temp['arm'] == 'loop') & (df_temp['from_start'] < 4))].groupby(['arm', 'type',
                                                                                'from_start'],
                                                                                as_index=False
                                                                                )[['pos']]. \
        count()[[
            'arm', 'type',
            'from_start',
            'pos'
        ]]

    dataframe.loc[(dataframe['type'] == 'pre-seed') & (dataframe['arm'] == '5p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'pre-seed') & (dataframe['arm'] == '5p'), 'from_start'] + 25

    dataframe.loc[(dataframe['type'] == 'seed') & (dataframe['arm'] == '5p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'seed') & (dataframe['arm'] == '5p'), 'from_start'] + 26

    dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['arm'] == '5p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['arm'] == '5p'), 'from_start'] + 33

    dataframe.loc[(dataframe['type'] == 'loop') & (dataframe['arm'] == 'loop'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'loop') & (dataframe['arm'] == 'loop'), 'from_start'] + 47

    for x in range(1, 26):
        if dataframe[(dataframe['type'] == 'flanking-5') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['5p', 'flanking-5', x, 0]], columns=['arm', 'type',
                                                                          'from_start',
                                                                          'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(26, 27):
        if dataframe[(dataframe['type'] == 'pre-seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['5p', 'pre-seed', x, 0]], columns=['arm', 'type',
                                                                        'from_start',
                                                                        'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(27, 34):
        if dataframe[(dataframe['type'] == 'seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['5p', 'seed', x, 0]], columns=['arm', 'type',
                                                                    'from_start',
                                                                    'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(34, 48):
        if dataframe[(dataframe['type'] == 'post-seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['5p', 'post-seed', x, 0]], columns=['arm', 'type',
                                                                         'from_start',
                                                                         'pos'])
            dataframe = pd.concat([dataframe, new_row])
    dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['from_start'] > 47), 'from_start'] = 47

    for x in range(48, 53 - 2):
        if dataframe[(dataframe['type'] == 'loop') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['loop', 'loop', x, 0]], columns=['arm', 'type',
                                                                      'from_start',
                                                                      'pos'])
            dataframe = pd.concat([dataframe, new_row])

    # dataframe = pd.concat([dataframe, add_loop_df])

    dataframe = dataframe.groupby(['arm', 'type',
                                   'from_start'],
                                  as_index=False)[['pos']].sum()[['arm', 'type',
                                                                  'from_start',
                                                                  'pos']]

    return dataframe, add_loop_value

def prepare_data_3p(df_temp):
    dataframe = df_temp[(df_temp['arm'] == '3p')].groupby(['arm', 'type', 'from_start'],
                                                          as_index=False)[['pos']].count()[['arm', 'type',
                                                                                'from_start',
                                                                                'pos']]

    try:

        dataframe_loop = df_temp[(df_temp['arm'] == 'loop') &
                                 (df_temp['from_end'] > -4)].groupby(['arm', 'type', 'from_end'],
                                                                     as_index=False)[['pos']].count()[['arm', 'type',
                                                                                'from_end',
                                                                                'pos']]
        dataframe_loop['from_start'] = dataframe_loop['from_end'].apply(
            lambda start: start + 4
        )
        dataframe_loop.drop('from_end', inplace=True, axis=1)
    except KeyError:

        dataframe_loop = df_temp[(df_temp['arm'] == 'loop') &
                                 (df_temp['from_end'] > -4)].groupby(['arm', 'type', 'from_end'],
                                                                     as_index=False)[['pos']].count()[['arm', 'type',
                                                                                'from_end',
                                                                                'pos']]
        dataframe_loop['from_start'] = dataframe_loop['from_end'].apply(
            lambda start: start + 4
        )
        dataframe_loop.drop('from_end', inplace=True, axis=1)

    dataframe = pd.concat([dataframe, dataframe_loop], sort=False)

    dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['arm'] == '3p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['arm'] == '3p'), 'from_start'] + 13 - 2

    dataframe.loc[(dataframe['type'] == 'seed') & (dataframe['arm'] == '3p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'seed') & (dataframe['arm'] == '3p'), 'from_start'] + 6 - 2

    dataframe.loc[(dataframe['type'] == 'pre-seed') & (dataframe['arm'] == '3p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'pre-seed') & (dataframe['arm'] == '3p'), 'from_start'] + 5 - 2

    dataframe.loc[(dataframe['type'] == 'flanking-3') & (dataframe['arm'] == '3p'), 'from_start'] = \
        dataframe.loc[(dataframe['type'] == 'flanking-3') & (dataframe['arm'] == '3p'), 'from_start'] + 27 - 2

    for x in range(28 - 2, 53 - 2):
        if dataframe[(dataframe['type'] == 'flanking-3') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['3p', 'flanking-3', x, 0]], columns=['arm', 'type',
                                                                          'from_start',
                                                                          'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(14 - 2, 28 - 2):
        if dataframe[(dataframe['type'] == 'post-seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['3p', 'post-seed', x, 0]], columns=['arm', 'type',
                                                                         'from_start',
                                                                         'pos'])
            dataframe = pd.concat([dataframe, new_row])

    dataframe.loc[(dataframe['type'] == 'post-seed') & (dataframe['from_start'] > 27 - 2), 'from_start'] = 27 - 2

    for x in range(7 - 2, 14 - 2):
        if dataframe[(dataframe['type'] == 'seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['3p', 'seed', x, 0]], columns=['arm', 'type',
                                                                    'from_start',
                                                                    'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(6 - 2, 7 - 2):
        if dataframe[(dataframe['type'] == 'pre-seed') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['3p', 'pre-seed', x, 0]], columns=['arm', 'type',
                                                                        'from_start',
                                                                        'pos'])
            dataframe = pd.concat([dataframe, new_row])

    for x in range(1, 6 - 2):
        if dataframe[(dataframe['type'] == 'loop') & (dataframe['from_start'] == x)].shape[0] == 0:
            new_row = pd.DataFrame([['loop', 'loop', x, 0]], columns=['arm', 'type',
                                                                      'from_start',
                                                                      'pos'])
            dataframe = pd.concat([dataframe, new_row])

    dataframe['from_start'] = dataframe['from_start'].apply(lambda start: start * -1)

    dataframe = dataframe.groupby(['arm', 'type',
                                   'from_start'],
                                  as_index=False)[['pos']].sum()[['arm', 'type',
                                                                  'from_start',
                                                                  'pos']]

    return dataframe

def prepare_figure(output_folder):

    if not os.path.exists(output_folder + '/plots'):
        os.makedirs(output_folder + '/plots')

    df_temp = pd.read_csv(output_folder + '/all_mutations_with_n_hgvs.csv')
    # df_temp = df_temp[df_temp['mutation_type'] == 'subst']

    mutations = df_temp.shape[0]
    genes = df_temp['pre_name'].nunique()

    create_plot(df_temp, output_folder + '/plots/plot_miRNA.svg',
                mutations, genes, mirna_type='both')

    # 5' dominant miRNAs

    df_5_dominant = df_temp[df_temp['balance'] == '5p']

    mutations = df_5_dominant.shape[0]
    genes = df_5_dominant['pre_name'].nunique()

    create_plot(df_5_dominant, output_folder + '/plots/plot_5p_balance_miRNA.svg',
                mutations, genes, mirna_type='5p')

    # 3p dominant miRNAs

    df_3_dominant = df_temp[df_temp['balance'] == '3p']

    mutations = df_3_dominant.shape[0]
    genes = df_3_dominant['pre_name'].nunique()

    create_plot(df_3_dominant, output_folder + '/plots/plot_3p_balance_miRNA.svg',
                mutations, genes, mirna_type='3p')

    # balanced miRNAs

    df_no_dominant = df_temp[df_temp['balance'] == 'both']

    mutations = df_no_dominant.shape[0]
    genes = df_no_dominant['pre_name'].nunique()

    create_plot(df_no_dominant, output_folder + '/plots/plot_balanced_miRNA.svg',
                mutations, genes, mirna_type='both')

def create_plot_per_mirna(data_df, output_name, mutations=0, types='both', title1='title1', title='title'):
    df_plot_5, loop_value = prepare_data_5p(data_df)

    df_plot_3 = prepare_data_3p(data_df)

    max_value = max(df_plot_5['pos'].max(), df_plot_3['pos'].max())

    rc("pdf", fonttype=42)
    sns.set_style(style='white')

    palette = {'flanking-5': 'grey',
               'flanking-3': 'grey',
               'pre-seed': '#5481A6',
               'seed': 'darkblue',
               'post-seed': '#5481A6',
               'loop': 'grey',
               'silent-pre': '#5481A6',
               'silent-post': '#5481A6',
               'silent-seed': '#5481A6'
               }

    fig = plt.figure(figsize=(25, 10))
    ax = fig.add_axes([0, 0, 1, 1])

    ax.axis('off')
    plt.imshow(im)
    plt.xticks([])
    plt.yticks([])

    y_min, y_max = ax.get_ylim()
    x_min, x_max = ax.get_xlim()

    plt.text(0, y_min * -0.5, title1.replace('mir', 'miR'), horizontalalignment='left', weight='bold',
             verticalalignment='center', fontdict={'size': '38'})

    plt.text(0, y_min * -0.35, title, horizontalalignment='left',
             verticalalignment='center', fontdict={'size': '38'})

    plt.text(x_max * 0.92, y_min * 1, str(mutations), horizontalalignment='center',
             verticalalignment='center', fontdict={'size': '42'})

    plt.text(x_max * 0.24, y_min * -0.09, 'flanking region', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.02, x_max * 0.46], [y_min * -0.07, y_min * -0.07], lw=0.5, color='grey',
                        alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    plt.text(x_max * 0.65, y_min * -0.09, 'miRNA', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.47, x_max * 0.83], [y_min * -0.07, y_min * -0.07], lw=0.5, color='#5481A6',
                        alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    plt.text(x_max * 0.91, y_min * -0.09, 'loop', horizontalalignment='center', fontdict={'size': '42'})
    line = lines.Line2D([x_max * 0.84, x_max * 0.98], [y_min * -0.07, y_min * -0.07], lw=0.5, color='grey',
                        alpha=0.8)
    ax.add_line(line)
    line.set_clip_on(False)

    if types != '3p':
        plt.text(x_max * 0.54, y_min * 0, 'seed', horizontalalignment='center', fontdict={'size': '42'})
        line = lines.Line2D([x_max * 0.49, x_max * 0.59], [y_min * 0.03, y_min * 0.03], lw=0.5, color='darkblue',
                            alpha=0.8)
        ax.add_line(line)
        line.set_clip_on(False)

    if types != '5p':
        line = lines.Line2D([x_max * 0.675, x_max * 0.775], [y_min * 1.125, y_min * 1.125], lw=0.5,
                            color='darkblue',
                            alpha=0.8)
        ax.add_line(line)
        line.set_clip_on(False)
        plt.text(x_max * 0.725, y_min * 1.195, 'seed', horizontalalignment='center', fontdict={'size': '42'})

    if loop_value > 0:
        plt.text(
            x_max * 0.92, y_min * 0.6, '+ {} loop\nmutations'.format(loop_value),
            horizontalalignment='center',
            verticalalignment='center',
            fontdict={'size': '30'}
        )

    a = plt.axes([.058, .52, .82, .35])

    hue_order = ['flanking-5', 'pre-seed', 'seed', 'post-seed', 'loop']

    if types == '3p':
        hue_order = ['flanking-5', 'silent-pre', 'silent-seed', 'silent-post', 'loop']

        df_plot_5['type'] = df_plot_5['type'].apply(
            lambda x: 'silent-seed' if x == 'seed' else (
                'silent-pre' if x == 'pre-seed' else (
                    'silent-post' if x == 'post-seed' else x
                )
            )
        )
    labels = [str(x) for x in range(-25, 0)] + \
             [str(x) for x in range(1, 23)] + ['+' + str(x) for x in range(1, 4)]

    labels = ['L' if x == '+4' else x for x in labels]

    ax = sns.barplot(x="from_start", y="pos", hue="type",
                     data=df_plot_5, dodge=False,
                     hue_order=hue_order,
                     palette=palette,
                     ax=a)

    for loc in ['right', 'top', 'left', 'bottom']:
        ax.spines[loc].set_visible(False)

    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.legend(loc='upper right', ncol=2)

    plt.setp(ax.get_legend().get_texts(), fontsize='18')
    plt.setp(ax.get_legend().get_title(), fontsize='22')
    plt.grid(b=True, which='major', axis='y', color='lightgrey',
             linestyle='-', linewidth=0.75, zorder=2, alpha=0.5)

    ax.set_xticklabels(labels)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    if max_value > 5:
        plt.yticks(np.arange(0, max_value + 1, np.floor(max_value / 3)))
        plot_limit = max_value + 2
    else:
        plt.yticks(np.arange(0, max_value + 1, 1))
        plot_limit = max_value + 1

    ax.set_ylim([0, plot_limit])
    ax.xaxis.tick_bottom()
    for label in ax.get_xticklabels():
        if label.get_text() not in ['-25', '-20', '-15', '-10', '-5', '1', '5', '10', '15', '20',
                                    '+1']:
            label.set_visible(False)

    new_width = 0.35
    for patch in ax.patches:
        current_width = patch.get_width()
        diff = current_width - new_width

        # we change the bar width
        patch.set_width(new_width)

        # we recenter the bar
        patch.set_x(patch.get_x() + diff * .5)

    ax.tick_params(axis='both', which='major', pad=8, width=0.5)

    plt.setp(ax.patches, linewidth=0)

    ax.get_legend().remove()

    b = plt.axes([.025, .02, .82, .35], facecolor='w')
    hue_order = ['flanking-3', 'pre-seed', 'seed', 'post-seed', 'loop']
    labels = ['+' + str(x) for x in range(1, 26)][::-1] + [str(x) for x in range(1, 23)][::-1] + \
             ['-' + str(x) for x in range(1, 4)]

    if types == '5p':
        hue_order = ['flanking-3', 'silent-pre', 'silent-seed', 'silent-post', 'loop']

        df_plot_3['type'] = df_plot_3['type'].apply(
            lambda x: 'silent-seed' if x == 'seed' else (
                'silent-pre' if x == 'pre-seed' else (
                    'silent-post' if x == 'post-seed' else x
                )
            )
        )

    ax = sns.barplot(x="from_start", y="pos", hue="type",
                     data=df_plot_3, dodge=False,
                     hue_order=hue_order,
                     palette=palette,
                     ax=b)

    for loc in ['right', 'top', 'left', 'bottom']:
        ax.spines[loc].set_visible(False)

    ax.set_xlabel('')
    ax.set_ylabel('')

    plt.grid(b=True, which='major', axis='y', color='lightgrey',
             linestyle='-', linewidth=0.75, zorder=1, alpha=0.5)

    ax.set_xticklabels(labels)

    if max_value > 5:
        plt.yticks(np.arange(0, max_value + 1, np.floor(max_value / 3)))
        plot_limit = max_value + 2
    else:
        plt.yticks(np.arange(0, max_value + 1, 1))
        plot_limit = max_value + 1

    ax.get_legend().remove()
    ax.set_ylim([plot_limit, 0])

    ax.xaxis.tick_top()

    for tick in ax.xaxis.get_major_ticks():
        tick.label2.set_fontsize(30)

    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(30)

    for label in ax.get_xticklabels():
        if label.get_text() not in ['+25', '+20', '+15', '+10', '+5', '+1', '1', '5', '10', '15', '20',
                                    '-5']:
            label.set_visible(False)

    new_width = 0.35
    for patch in ax.patches:
        current_width = patch.get_width()
        diff = current_width - new_width

        # we change the bar width
        patch.set_width(new_width)

        # we recenter the bar
        patch.set_x(patch.get_x() + diff * .5)

    plt.setp(ax.patches, linewidth=0)

    plt.savefig(output_name, format='svg', dpi=300, transparent=True, bbox_inches='tight')
    plt.cla()
    plt.clf()
    plt.close("all")
    plt.close(fig)
    gc.collect()

def prepare_figures_per_mirna(output_folder):
    if not os.path.exists(output_folder + '/plots'):
        os.makedirs(output_folder + '/plots')
    if not os.path.exists(output_folder + '/plots/miRNAs'):
        os.makedirs(output_folder + '/plots/miRNAs')

    df_temp = pd.read_csv(output_folder + '/all_mutations_with_n_hgvs.csv')
    # df_temp = df_temp[df_temp['mutation_type'] == 'subst']

    mirnas = list(set(df_temp['pre_name'].unique()))

    for mirna in mirnas:

        title = ''

        title1 = mirna

        df_temp2 = df_temp[df_temp['pre_name'] == mirna].copy()

        type_of_miRNA = list(df_temp2['balance'].unique())[0]

        mutations = df_temp2.shape[0]

        create_plot_per_mirna(df_temp2, output_folder + '/plots/miRNAs/plot_miRNA_{}.svg'.format(mirna),
                              mutations, types=type_of_miRNA, title1=title1, title=title)

        del df_temp2
        gc.collect()

@click.command()
@click.argument('output_folder')
def main(output_folder):
    prepare_figure(output_folder)
    prepare_figures_per_mirna(output_folder)

if __name__ == "__main__":
    main()

max_stars_repo_name: martynaut/mirnome-mutations

max_issues_repo_name: martynaut/mirnome-mutations

max_forks_repo_name: martynaut/mirnome-mutations


text:
import tensorflow as tf
import numpy as np
import pylab as plt
import os
import astropy.time as at
import astropy.coordinates as ac
import astropy.units as au
from bayes_filter.datapack import _load_array_file, DataPack
from bayes_filter.feeds import Feed
from collections import namedtuple
from bayes_filter.frames import ENU
from scipy.spatial.distance import pdist
import tensorflow_probability as tfp
from bayes_filter.misc import flatten_batch_dims

C = 299792458.

# class SampleTimeFeed(Feed):
#     def __init__(self, time0, step=1, end=None):
#         self._step = step
#         self._end = end
#         self.step = tf.convert_to_tensor(step, tf.int32)
#         self.index_feed = tf.data.Dataset.from_generator(self.index_generator,
#                                                          (tf.int32, tf.int32),
#                                                          (tf.TensorShape([]), tf.TensorShape([])))
#         self.feed = self.index_feed
#
#     def index_generator(self):
#         if self._end is None:
#             i = 0
#             while True:
#                 i += self._step
#                 yield i - self._step, i
#         else:
#             i = 0
#             while i < self._end:
#                 i += self._step
#                 yield i - self._step, min(i, self._end)

def make_obs_setup(Nd, time_total=10, fov=2.5, interval=1, Nf=2, seed = 0):
    np.random.seed(seed)
    antenna_labels, antennas = _load_array_file(DataPack.lofar_array)
    antennas = ac.SkyCoord(x=antennas[:,0] * au.m, y=antennas[:,1] * au.m, z=antennas[:,2] * au.m, frame='itrs')
    array_center = np.mean(antennas.cartesian.xyz, axis=1)
    array_center = ac.SkyCoord(x=array_center[0], y=array_center[1], z=array_center[2], frame='itrs')
    time0 = at.Time("2019-01-01T00:00:00.000", format='isot')
    altaz = ac.AltAz(location=array_center.earth_location, obstime=time0)
    up = ac.SkyCoord(alt=90. * au.deg, az=0. * au.deg, frame=altaz).transform_to('icrs')
    directions = np.stack([np.random.normal(up.ra.rad, np.pi / 180. * fov, size=[Nd]),
                           np.random.normal(up.dec.rad, np.pi / 180. * fov, size=[Nd])], axis=1)
    directions = ac.SkyCoord(directions[:,0]*au.rad, directions[:,1]*au.rad, frame='icrs')
    patch_names = np.array(['patch_{:04d}'.format(i) for i in range(len(directions))])

    Nt = int(time_total/interval)

    times = at.Time(time0.gps + np.linspace(0, time_total, Nt), format='gps')#.mjd[:, None] * 86400.  # mjs
    freqs = np.linspace(120, 160, Nf) * 1e6

    fluxes = np.random.uniform(low=0.1, high=1., size=Nd)

    return dict(times=times, freqs=freqs, fluxes=fluxes, antennas=antennas, directions=directions)

def simulate_vis(obs_setup, vis_noise=None, vis_frac_noise=0.5, seed=0):
    np.random.seed(seed)
    times = obs_setup['times']
    directions = obs_setup['directions']
    Nd = len(directions)
    antennas = obs_setup['antennas']
    Na = len(antennas)
    freqs = obs_setup['freqs']
    Nf = len(freqs)
    fluxes = obs_setup['fluxes']
    array_center = np.mean(antennas.cartesian.xyz, axis=1)
    array_center = ac.SkyCoord(x=array_center[0], y=array_center[1], z=array_center[2], frame='itrs')
    visibilities = []
    obs_times = []

    factor = -2.*np.pi*1j*freqs/C# Nf

    geometric = []
    for time in times:
        enu = ENU(location=array_center.earth_location, obstime=time)
        k = directions.transform_to(enu).cartesian.xyz.value #3, Nd
        x = antennas.transform_to(enu).cartesian.xyz.to(au.m).value #3, Na
        kx = np.einsum('si,sj->ij', k, x) #Nd, Na
        dkx = kx[:,:,None] - kx[:,None,:]#Nd, Na, Na
        geometric.append(np.exp(factor[None, None, None, :]*dkx[:, :, :, None]))

    dt = times[1].gps - times[0].gps

    geometric = np.stack(geometric,axis=-1)
    vis = np.sum(fluxes[:,None,None,None,None]*geometric,axis=0)*dt
    if vis_noise is not None:
        vis += vis_noise*np.random.normal(size=vis.shape) + vis_noise*1j*np.random.normal(size=vis.shape)
    else:
        vis += vis*vis_frac_noise * (np.random.normal(size=vis.shape) + 1j*np.random.normal(size=vis.shape))
    return dict(sim_visibilities=vis, geometric=geometric, times=times, uncert=vis*vis_frac_noise)

class Target:
    def __init__(self, obs_vis, uncert, times, geometric):
        self.obs_vis = tf.reshape(tf.convert_to_tensor(obs_vis, tf.complex128), [-1])#Na Na Nf Nt
        self.uncert = tf.reshape(tf.convert_to_tensor(uncert, tf.float64), [-1])#Na Na Nf Nt

        self.obs_vis_real = tf.real(self.obs_vis)
        self.obs_vis_imag = tf.imag(self.obs_vis)

        self.times = tf.convert_to_tensor(times) #Nt
        self.dt = tf.cast(self.times[1] - self.times[0], tf.float64)
        self.geometric = tf.convert_to_tensor(geometric, tf.complex128) #Nd, Na, Na, Nf, Nt
        self.geometric = tf.reshape(self.geometric, tf.concat([tf.shape(self.geometric)[0:1], [-1]],axis=0)) #Nd, Na Na Nf Nt
        self.geometric_real = tf.real(self.geometric)
        self.geometric_imag = tf.imag(self.geometric)

        self.fluxes_bijector = tfp.bijectors.Exp()
        self.sigma_bijector = tfp.bijectors.Exp()

    def forward(self,fluxes):

        geometric_real = self.geometric_real#tf.gather(self.geometric_real, selection, axis=1)
        geometric_imag = self.geometric_imag#tf.gather(self.geometric_imag, selection, axis=1)

        vis_real = tf.reduce_sum(fluxes[:, :, None] * geometric_real[None, :, :],
                                 axis=1) * self.dt
        vis_imag = tf.reduce_sum(fluxes[:, :, None] * geometric_imag[None, :, :],
                                 axis=1) * self.dt # Na Na Nf Nt*
        return vis_real, vis_imag

    def log_prob(self, log_fluxes):
        fluxes = self.fluxes_bijector.forward(log_fluxes)# num_chains, Nd
        # sigma = self.sigma_bijector.forward(log_sigma)# num_chains, 1

        # n = tf.cast(tf.shape(self.obs_vis_real)[0], dtype=sigma.dtype)
        # p_erdos = 2*tf.math.log(n) / n
        # selection = tf.where(
        #     tf.equal(tfp.distributions.Binomial(total_count=tf.constant(1., dtype=sigma.dtype),
        #                                         probs=1.-p_erdos).sample(tf.cast(n, tf.int32)),
        #        selection = tf.where(
        #     tf.equal(tfp.distributions.Binomial(total_count=tf.constant(1., dtype=sigma.dtype),
        #                                         probs=1.-p_erdos).sample(tf.cast(n, tf.int32)),
        #              tf.constant(1., dtype=sigma.dtype)))[:,0]

        model_vis = self.forward(fluxes)#self.forward(fluxes, selection)

        obs_vis_real = self.obs_vis_real#tf.gather(self.obs_vis_real, selection, axis=0)
        obs_vis_imag = self.obs_vis_imag#tf.gather(self.obs_vis_imag, selection, axis=0)

        vis_sigma = self.uncert#sigma# * tf.sqrt(tf.square(model_vis[0]) + tf.square(model_vis[1]))

        likelihood = tf.reduce_sum(tfp.distributions.Normal(loc=model_vis[0], scale=vis_sigma).log_prob(obs_vis_real[None,:]) \
                     + tfp.distributions.Normal(loc=model_vis[1], scale=vis_sigma).log_prob(obs_vis_imag[None,:]),axis=1)

        # prior = tfp.distributions.Normal(loc=tf.constant(0.1,sigma.dtype), scale=tf.constant(0.1, sigma.dtype)).log_prob(sigma[:,0])

        return likelihood# + prior

def run(output_folder='bayes_vis_output_1percent_10'):
    output_folder = os.path.abspath(output_folder)
    num_chains = 4
    Nd = 10
    obs_setup = make_obs_setup(Nd=Nd, time_total=2, fov=2.5, interval=1, Nf=2)
    sim = simulate_vis(obs_setup,vis_frac_noise=0.01)
    with tf.Session(graph=tf.Graph()) as sess:
        target = Target(sim['sim_visibilities'],sim['uncert'], sim['times'].gps, sim['geometric'])
        ###
        hmc = tfp.mcmc.SimpleStepSizeAdaptation(tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=target.log_prob,
            num_leapfrog_steps=2,
            step_size=[1.]),
            num_adaptation_steps=800,
            target_accept_prob=tf.constant(0.6,tf.float64),
            adaptation_rate=0.1)

        # Run the chain (with burn-in maybe).
        # last state as initial point (mean of each chain)

        def trace_fn(_, pkr):
            # print(pkr)
            return (pkr.inner_results.log_accept_ratio,
                    pkr.inner_results.accepted_results.step_size)

        init_state = [#tf.cast(target.sigma_bijector.inverse(tf.random.uniform([num_chains, 1], 0.01,  0.1)), dtype=tf.float64),
                      tf.cast(target.fluxes_bijector.inverse(tf.random.uniform([num_chains, Nd], 0.1,  1.)), dtype=tf.float64)]

        # u = target.log_prob(*init_state)
        # print(tf.gradients(u,init_state))
        samples, (log_accept_ratio, stepsizes) = tfp.mcmc.sample_chain(
            num_results=10000,
            num_burnin_steps=1000,
            trace_fn=trace_fn,
            return_final_kernel_results=False,
            current_state=init_state,
            kernel=hmc,
            parallel_iterations=10)

        rhat = tfp.mcmc.potential_scale_reduction(samples)

        next_stepsizes = [#tf.sqrt(tf.reduce_mean(tfp.stats.variance(samples[0]))),
                          tf.sqrt(tf.reduce_mean(tfp.stats.variance(samples[0])))]

        final_state = dict(#sigma=flatten_batch_dims(target.sigma_bijector.forward(samples[0]),-1),
                       fluxes=flatten_batch_dims(target.fluxes_bijector.forward(samples[0]),-1),
                           stepsizes=stepsizes,
                           next_stepsizes=next_stepsizes,
                           rhat=rhat)

        res = sess.run(final_state)
        print(res['next_stepsizes'])

        os.makedirs(output_folder, exist_ok=True)
        idx = np.arange(Nd)
        plt.boxplot(res['fluxes'], positions=idx, sym="")
        plt.scatter(idx,obs_setup['fluxes'])
        plt.savefig(os.path.join(output_folder,'fluxes.png'))
        plt.close('all')

        plt.boxplot(res['fluxes']-obs_setup['fluxes'], positions=idx, sym="")
        plt.savefig(os.path.join(output_folder, 'res_fluxes.png'))
        plt.close('all')

        plt.hist(np.median(res['fluxes'] - obs_setup['fluxes'], axis=0), bins = max(10, int(np.sqrt(Nd))))
        plt.savefig(os.path.join(output_folder, 'hist_res_fluxes.png'))
        plt.close('all')

        # plt.hist(res['sigma'],bins=int(np.sqrt(res['sigma'].size)))
        # plt.savefig(os.path.join(output_folder,'sigma.png'))
        # plt.close('all')

        plt.plot(res['stepsizes'])
        plt.savefig(os.path.join(output_folder,'stepsizes.png'))
        plt.close('all')

        # plt.bar([0],res['rhat'][0])
        # plt.savefig(os.path.join(output_folder,'rhat_sigma.png'))
        # plt.close('all')

        plt.bar(idx, res['rhat'][0])
        plt.savefig(os.path.join(output_folder, 'rhat_fluxes.png'))
        plt.close('all')

if __name__=='__main__':
    run()

ext: py

reason: import numpy,from scipy,import astropy


text:
import sys
from typing import List

import streamlit as st
from sympy import diff, lambdify, parse_expr
from sympy.parsing.sympy_parser import (
    convert_xor,
    function_exponentiation,
    implicit_application,
    implicit_multiplication,
    split_symbols,
    standard_transformations,
)

sys.path.append('')
sys.path.append('../../..')

from src.common.config import OUTPUT_PRECISION
from src.common.methods.root_separation.segment_tabulation import Tabulator
from src.common.model.line_segment import LineSegment
from src.tasks.task1.main import METHOD_NAME_TO_ROOT_FINDER

TRANSFORMATIONS = standard_transformations + (
    split_symbols,
    implicit_application,
    implicit_multiplication,
    function_exponentiation,
    convert_xor,
)

def root_separation(*, number_of_parts: int, expression, line_segment: LineSegment) -> List[LineSegment]:
    tabulator = Tabulator(number_of_parts)

    f = lambdify('x', expression)
    segments = tabulator.separate(f=f, line_segment=line_segment)

    for index, segment in enumerate(segments):
        st.markdown(f'{index + 1}. ${segment}$')

    return segments

def root_find(*, segments: List[LineSegment], method_name: str, expression, accuracy: float):
    root_finder = METHOD_NAME_TO_ROOT_FINDER[method_name]

    for number, segment in enumerate(segments, start=1):
        f = lambdify('x', expression)
        df = lambdify('x', diff(expression))

        root_finder.find(derivatives=[f, df], line_segment=segment, accuracy=accuracy)
        stats = root_finder.stats

        st.subheader(f'Отрезок №{number}')

        st.markdown(
            f"""
            - Отрезок: ${stats.line_segment}$

            - Начальное приближение: ${stats.round_initial_approximation()}$

            - Количество шагов: ${stats.number_of_steps}$

            - $x_m = {round(stats.approximate_solution, OUTPUT_PRECISION)}$

            - $\\Delta = {round(stats.error, OUTPUT_PRECISION)}$

            - $|f(x_m)| = {round(stats.residual, OUTPUT_PRECISION)}$

            """,
            unsafe_allow_html=True,
        )

def app():
    with st.sidebar:
        st.header('Параметры')

        text_expression = st.text_input('Выражение:', value='x - 10 * sin(x)')

        left_column, right_column = st.columns(2)

        with left_column:
            a = st.number_input('A:', value=-5)

        with right_column:
            b = st.number_input('B:', value=3)

        accuracy = st.number_input('Точность:', value=10 ** -6, format='%e')

        number_of_parts = st.number_input('Количество частей:', value=1000)

    expression = parse_expr(text_expression)
    line_segment = LineSegment(a, b)

    st.title('Численные методы решения нелинейных уравнений')

    method_name = st.selectbox('Метод:', options=METHOD_NAME_TO_ROOT_FINDER.keys())

    st.subheader('Отделение корней')
    segments = root_separation(number_of_parts=number_of_parts, expression=expression, line_segment=line_segment)

    root_find(segments=segments, method_name=method_name, accuracy=accuracy, expression=expression)

if __name__ == '__main__':
    app()

reason: from sympy


text:
import numpy as np

from graph.gates import Sigmoid, Relu
from model.model import ClassificationModel
from model.layers import Layer
from model.loss_functions import MSE
from model.optimizers import Adam

model = ClassificationModel([
    Layer(28*28, 10, Relu)
], MSE, Adam, lr=0.0001)

import mnist

x_train, y_train = mnist.train_images(), mnist.train_labels()

# Flatten images
x_train = np.array([x.flatten() for x in x_train])

# Normalize images
x_train = x_train / 255.0

# Convert labels to vectors
y_train = np.array([np.array([int(i == y) for i in range(10)]) for y in y_train])

sample_size = 200

x_train = x_train[:sample_size]
y_train = y_train[:sample_size]

model.fit(x_train, y_train, epochs=15)

max_stars_repo_name: maximlopin/computational-graph

max_issues_repo_name: maximlopin/computational-graph

max_forks_repo_name: maximlopin/computational-graph

reason: import numpy


text:
#coding: utf-8

import numpy as np

import utils

####################################################
## Learners                                       ##
####################################################

# Learner structure
class Learner(object):
	def __init__(self, K):
		self.K = K
		## Arm playing mixed strategy
		self.p = np.zeros(self.K)
	def act(self):
		raise NotImplemented
	def incur(self, w):
		raise NotImplemented

# Follow the Leader learner
class FTL(Learner):
	def __init__(self, K):
		super(FTL, self).__init__(K)
	def act(self):
		p = np.asarray(self.p == np.min(self.p), dtype=float)
		p /= np.sum(p)
		return p
	def incur(self, w):
		self.p += w

# AdaHedge learner (the one used in LinGame)
class AdaHedge(Learner):
	def __init__(self, K):
		self.delta = 0.01
		super(AdaHedge, self).__init__(K)
	def act(self):
		eta = np.log(self.K)/float(self.delta)
		p = np.exp(-eta*(self.p-np.min(self.p)))
		p /= np.sum(p)
		return p
	def incur(self, w):
		p = self.act()
		eta = np.log(self.K)/float(self.delta)
		self.p += w
		m = np.min(w)-1./eta*np.log(p.T.dot(np.exp(-eta*(w-np.min(w)))))
		assert m != float("inf") and p.T.dot(w) >= m-1e-7
		self.delta += p.T.dot(w)-m

# Lazy Mirror Descent/Fixed share learner
class FixedShare(Learner):
	def __init__(self, K, S=1):
		self.K = K
		self.p = 1./float(self.K)*np.ones(self.K)
		self.t = 0
		self.S = S
	def act(self):
		return self.p
	def incur(self, w):
		self.t += 1
		eta = np.sqrt(np.log(self.K)/float(self.t))/float(self.S)
		self.p *= np.exp(-eta*(w-np.min(w)))
		self.p /= np.sum(self.p)
		gamma = 1./float(4*np.sqrt(self.t))
		self.p = (1-gamma)*self.p+gamma/float(self.K)

max_stars_repo_name: clreda/linear-top-m

max_issues_repo_name: clreda/linear-top-m

max_forks_repo_name: clreda/linear-top-m


text:
import unittest

import numpy
import pytest

import cupy
from cupy import testing

@testing.parameterize(*testing.product({
    'xp': [numpy, cupy],
    'shape': [(3, 2), (), (3, 0, 2)],
}))
@testing.gpu
class TestShapedRandom(unittest.TestCase):

    @testing.for_all_dtypes()
    def test_shape_and_dtype(self, dtype):
        a = testing.shaped_random(self.shape, self.xp, dtype)
        assert isinstance(a, self.xp.ndarray)
        assert a.shape == self.shape
        assert a.dtype == dtype

    @testing.for_all_dtypes(no_bool=True, no_complex=True)
    def test_value_range(self, dtype):
        a = testing.shaped_random(self.shape, self.xp, dtype)
        assert self.xp.all(0 <= a)
        assert self.xp.all(a < 10)

    @testing.for_complex_dtypes()
    def test_complex(self, dtype):
        a = testing.shaped_random(self.shape, self.xp, dtype)
        assert self.xp.all(0 <= a.real)
        assert self.xp.all(a.real < 10)
        assert self.xp.all(0 <= a.imag)
        assert self.xp.all(a.imag < 10)
        if 0 not in self.shape:
            assert self.xp.any(a.imag)

@testing.parameterize(*testing.product({
    'xp': [numpy, cupy],
}))
@testing.gpu
class TestShapedRandomBool(unittest.TestCase):

    def test_bool(self):
        a = testing.shaped_random(10000, self.xp, numpy.bool_)
        assert 4000 < self.xp.sum(a) < 6000

@testing.parameterize(*testing.product({
    'dtype': [
        numpy.float16, numpy.float32, numpy.float64,
        numpy.complex64, numpy.complex128,
    ],
    'xp': [numpy, cupy],
    'x_s_shapes': [
        ((0, 0), (0,)),
        ((2, 2), (2,)),
        ((2, 3), (2,)),
        ((3, 2), (2,)),
        # broadcast
        ((2, 2), ()),
    ],
}))
class TestGenerateMatrix(unittest.TestCase):

    def test_generate_matrix(self):
        dtype = self.dtype
        x_shape, s_shape = self.x_s_shapes
        sv = self.xp.random.uniform(
            0.5, 1.5, s_shape).astype(dtype().real.dtype)
        x = testing.generate_matrix(
            x_shape, xp=self.xp, dtype=dtype, singular_values=sv)
        assert x.shape == x_shape

        if 0 in x_shape:
            return

        s = self.xp.linalg.svd(
            x.astype(numpy.complex128), full_matrices=False, compute_uv=False,
        )
        sv = self.xp.broadcast_to(sv, s.shape)
        sv_sorted = self.xp.sort(sv, axis=-1)[..., ::-1]

        rtol = 1e-3 if dtype == numpy.float16 else 1e-7
        self.xp.testing.assert_allclose(s, sv_sorted, rtol=rtol)

class TestGenerateMatrixInvalid(unittest.TestCase):

    def test_no_singular_values(self):
        with self.assertRaises(TypeError):
            testing.generate_matrix((2, 2))

    def test_invalid_shape(self):
        with self.assertRaises(ValueError):
            testing.generate_matrix((2,), singular_values=1)

    def test_invalid_dtype_singular_values(self):
        with self.assertRaises(TypeError):
            testing.generate_matrix((2, 2), singular_values=1 + 0j)

    def test_invalid_dtype(self):
        with self.assertRaises(TypeError):
            testing.generate_matrix(
                (2, 2), dtype=numpy.int32, singular_values=1)

    def test_negative_singular_values(self):
        with self.assertRaises(ValueError):
            testing.generate_matrix((2, 2), singular_values=[1, -1])

    def test_shape_mismatch(self):
        with self.assertRaises(ValueError):
            testing.generate_matrix(
                (2, 2), singular_values=numpy.ones(3))

    def test_shape_mismatch_2(self):
        with self.assertRaises(ValueError):
            testing.generate_matrix(
                (0, 2, 2), singular_values=numpy.ones(3))

class TestAssertFunctionIsCalled(unittest.TestCase):

    def test_patch_ndarray(self):
        orig = cupy.ndarray
        with testing.AssertFunctionIsCalled('cupy.ndarray'):
            a = cupy.ndarray((2, 3), numpy.float32)
        assert cupy.ndarray is orig
        assert not isinstance(a, cupy.ndarray)

    def test_spy_ndarray(self):
        orig = cupy.ndarray
        with testing.AssertFunctionIsCalled(
                'cupy.ndarray', wraps=cupy.ndarray):
            a = cupy.ndarray((2, 3), numpy.float32)
        assert cupy.ndarray is orig
        assert isinstance(a, cupy.ndarray)

    def test_fail_not_called(self):
        orig = cupy.ndarray
        with pytest.raises(AssertionError):
            with testing.AssertFunctionIsCalled('cupy.ndarray'):
                pass
        assert cupy.ndarray is orig

    def test_fail_called_twice(self):
        orig = cupy.ndarray
        with pytest.raises(AssertionError):
            with testing.AssertFunctionIsCalled('cupy.ndarray'):
                cupy.ndarray((2, 3), numpy.float32)
                cupy.ndarray((2, 3), numpy.float32)
        assert cupy.ndarray is orig

    def test_times_called(self):
        orig = cupy.ndarray
        with testing.AssertFunctionIsCalled('cupy.ndarray', times_called=2):
            cupy.ndarray((2, 3), numpy.float32)
            cupy.ndarray((2, 3), numpy.float32)
        assert cupy.ndarray is orig

    def test_inner_error(self):
        orig = cupy.ndarray
        with pytest.raises(numpy.AxisError):
            with testing.AssertFunctionIsCalled('cupy.ndarray'):
                cupy.ndarray((2, 3), numpy.float32)
                raise numpy.AxisError('foo')
        assert cupy.ndarray is orig

reason: import numpy,import cupy,from cupy


text:
import sys
import pytest
import numpy as np
sys.path.append('/pymoc/src/pymoc/utils')
from gridit import gridit

class TestGridit(object):
  def test_gridit(self):
    f = lambda x, y: x ^ 2 + y
    nx = 3
    ny = 2
    test = np.asarray([[2, 3], [3, 2], [0, 1]])
    sol = gridit(range(nx), range(ny), f)

    for i in range(nx):
      for j in range(ny):
        assert test[i, j] == sol[i, j]


text:
import numpy as np
import math
def mi(pos, neg):
    P = np.array([[pos / 2, (1 - pos) / 2],
                  [neg / 2, (1 - neg) / 2]])
    Px = np.sum(P, axis=1)
    Py = np.sum(P, axis=0)
    res = 0.0
    for i in range(2):
        for j in range(2):
            if P[i,j] > 0 and Px[i] > 0 and Py[j] > 0:
                res += P[i,j] * math.log(P[i, j] / Px[i] / Py[j])
                return res

fns = ['bert_amazon_polarity.csv', 'bert_yelp_funny_balanced.csv'] # , 'yelp_polarity.csv', 'book_spoiler_balanced.csv']
for fn in fns:
    results = []
    mapping = {}
    for line in open('../search/log/' + fn):
        if len(line.strip().split(',')) != 4:
            continue
        token, _, pos, neg = line.strip().split(',')
        mapping[token] = (float(pos), float(neg))
        results.append((mi(float(pos), float(neg)), token))
        # break
        results.sort(reverse=True)
        print(fn)
        for score, token in results[:30]:
            print(score, token, mapping[token])
            print('-------------------------')

max_stars_repo_name: rit-git/tagging

max_issues_repo_name: rit-git/tagging

max_forks_repo_name: rit-git/tagging


text:
"""
Implementation of Spin permutations.
"""

# Author: Oualid Benkarim <oualid.benkarim@mcgill.ca>
# License: BSD 3 clause

import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist
from scipy.spatial import cKDTree

from sklearn.utils import check_random_state
from sklearn.base import BaseEstimator

from ..vtk_interface.wrappers import BSPolyData
from ..mesh import mesh_elements as me

def _generate_spins(points_lh, points_rh=None, unique=False, n_rep=100,
                    random_state=None, surface_algorithm='FreeSurfer'):
    """ Generate rotational spins based on points that lie on a sphere.

    Parameters
    ----------
    points_lh : BSPolyData or ndarray, shape = (n_lh, 3)
        Array of points in a sphere, where `n_lh` is the number of points.
    points_rh : BSPolyData or ndarray, shape = (n_rh, 3), optional
        Array of points in a sphere, where `n_rh` is the number of points. If
        provided, rotations are derived from the rotations computed for
        `points_lh` by reflecting the rotation matrix across the Y-Z plane.
        Default is None.
    unique : bool, optional
        Whether to enforce a one-to-one correspondence between original points
        and rotated ones. If true, the Hungarian algorithm is used.
        Default is False.
    n_rep : int, optional
        Number of random rotations. Default is 100.
    surface_algorithm : {'FreeSurfer', 'CIVET'}
        For 'CIVET', no flip is required to generate the spins for the right
        hemisphere. Only used when ``points_rh is not None``.
        Default is 'FreeSurfer'.
    random_state : int or None, optional
        Random state. Default is None.

    Returns
    -------
    result : dict[str, ndarray]
        Spin indices for left points (and also right, if provided).

    References
    ----------
    * Alexander-Bloch A, Shou H, Liu S, Satterthwaite TD, Glahn DC,
      Shinohara RT, Vandekar SN and Raznahan A (2018). On testing for spatial
      correspondence between maps of human brain structure and function.
      NeuroImage, 178:540-51.
    * Blaser R and Fryzlewicz P (2016). Random Rotation Ensembles.
      Journal of Machine Learning Research, 17(4): 1–26.
    * https://netneurotools.readthedocs.io

    """

    # Handle if user provides spheres
    if not isinstance(points_lh, np.ndarray):
        points_lh = me.get_points(points_lh)

    if points_rh is not None:
        if not isinstance(points_rh, np.ndarray):
            points_rh = me.get_points(points_rh)

    pts = {'lh': points_lh}
    if points_rh is not None:
        pts['rh'] = points_rh

        # for reflecting across Y-Z plane
        reflect = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])

    idx = {k: np.arange(p.shape[0]) for k, p in pts.items()}
    spin = {k: np.empty((n_rep, p.shape[0]), dtype=int)
            for k, p in pts.items()}
    if not unique:
        tree = {k: cKDTree(p, leafsize=20) for k, p in pts.items()}

    rs = check_random_state(random_state)

    rot = {}
    for i in range(n_rep):

        # generate rotation for left
        rot['lh'], temp = np.linalg.qr(rs.normal(size=(3, 3)))
        rot['lh'] *= np.sign(np.diag(temp))
        rot['lh'][:, 0] *= np.sign(np.linalg.det(rot['lh']))

        # reflect the left rotation across Y-Z plane
        if 'rh' in pts:
            if surface_algorithm.lower() == 'freesurfer':
                rot['rh'] = reflect @ rot['lh'] @ reflect
            else:
                rot['rh'] = rot['lh']

        for k, p in pts.items():
            if unique:
                dist = cdist(p, p @ rot[k])
                row, col = linear_sum_assignment(dist)
                spin[k][i, idx[k]] = idx[k][col]
            else:
                _, spin[k][i] = tree[k].query(p @ rot[k], k=1, n_jobs=1)

    return spin

def spin_permutations(spheres, data, unique=False, n_rep=100,
                      random_state=None, surface_algorithm='FreeSurfer'):
    """ Generate null data using spin permutations.

    Parameters
    ----------
    spheres : dict[str, ndarray or BSPolyData], BSPolyData or ndarray
        Dictionary of points in a sphere, for left ('lh' key) and
        right ('rh' key) hemispheres. The right hemisphere is optional. If
        provided, rotations are derived from the rotations computed for
        `points_lh` by reflecting the rotation matrix across the Y-Z plane.
    data : dict[str, ndarray] or ndarray
        Dictionary of data to randomize. Array of variables arranged in
        columns for each hemisphere.
    unique : bool, optional
        Whether to enforce a one-to-one correspondence between original points
        and rotated ones. If true, the Hungarian algorithm is used.
        Default is False.
    n_rep : int, optional
        Number of random rotations. Default is 100.
    surface_algorithm : {'FreeSurfer', 'CIVET'}
        For 'CIVET', no flip is required to generate the spins for the right
        hemisphere. Only used when ``points_rh is not None``.
        Default is 'FreeSurfer'.
    random_state : int or None, optional
        Random state. Default is None.

    Returns
    -------
    rand_lh : ndarray, shape = (n_rep, n_lh, n_feat)
        Permutations of data in left hemisphere.
    rand_rh : ndarray, shape = (n_rep, n_rh, n_feat)
        Permutations of data in right hemisphere. Only if right data and
        sphere are provided.

    See Also
    --------
    :class:`.SpinPermutations`

    References
    ----------
    * Alexander-Bloch A, Shou H, Liu S, Satterthwaite TD, Glahn DC,
      Shinohara RT, Vandekar SN and Raznahan A (2018). On testing for spatial
      correspondence between maps of human brain structure and function.
      NeuroImage, 178:540-51.
    * Blaser R and Fryzlewicz P (2016). Random Rotation Ensembles.
      Journal of Machine Learning Research, 17(4): 1–26.
    * https://netneurotools.readthedocs.io

    """

    if isinstance(data, np.ndarray):
        data = {'lh': data}

    if isinstance(spheres, BSPolyData) or isinstance(spheres, np.ndarray):
        spheres = {'lh': spheres}

    if data.keys() != spheres.keys():
        raise ValueError("Keys for data and spheres do not coincide.")

    if len(data) > 2:
        raise ValueError("Unknown keys. Possible keys: {'lh', 'rh'}.")

    if len(data) == 1 and 'lh' not in data.keys():
        raise ValueError("Key must be 'lh'.")
    elif len(data) == 2 and ('lh' not in data or 'rh' not in data):
        raise ValueError("Unknown keys. Possible keys: {'lh', 'rh'}.")

    points_lh = spheres['lh']
    points_rh = spheres.pop('rh', None)

    spin_idx = _generate_spins(points_lh, points_rh=points_rh, unique=unique,
                               n_rep=n_rep, random_state=random_state,
                               surface_algorithm=surface_algorithm)

    spin_lh = spin_idx['lh']
    spin_rh = spin_idx.pop('rh', None)

    x_lh = data['lh']
    rand_lh = x_lh[spin_lh]
    if spin_rh is None:
        return rand_lh

    x_rh = data.pop('rh', None)
    rand_rh = None
    if x_rh is not None:
        rand_rh = x_rh[spin_rh]

    return rand_lh, rand_rh

class SpinPermutations(BaseEstimator):
    """ Spin permutations.

    Parameters
    ----------
    unique : bool, optional
        Whether to enforce a one-to-one correspondence between original points
        and rotated ones. If True, the Hungarian algorithm is used.
        Default is False.
    n_rep : int, optional
        Number of randomizations. Default is 100.
    random_state : int or None, optional
        Random state. Default is None.
    surface_algorithm : {'FreeSurfer', 'CIVET'}
        For 'CIVET', no flip is required to generate the spins for the right
        hemisphere. Only used when ``points_rh is not None``.
        Default is 'FreeSurfer'.

    Attributes
    ----------
    spin_lh_ : ndarray, shape (n_rep, n_lh)
        Spin indices for points in left hemisphere.
    spin_rh_ : ndarray, shape (n_rep, n_rh)
        Spin indices for points in right hemisphere. Only if user provides
        right hemisphere points. None, otherwise.

    See Also
    --------
    :func:`.spin_permutations`
    :class:`.MoranRandomization`

    Notes
    -----
    Right hemisphere permutations are generated by reflecting the rotation
    matrix used for the left hemisphere.

    """

    def __init__(self, unique=False, n_rep=100, random_state=None,
                 surface_algorithm='FreeSurfer'):
        self.unique = unique
        self.n_rep = n_rep
        self.random_state = random_state
        self.surface_algorithm = surface_algorithm

    def fit(self, points_lh, points_rh=None):
        """ Compute spin indices by random rotation.

        Parameters
        ----------
        points_lh : BSPolyData or ndarray, shape = (n_lh, 3)
            Sphere for the left hemisphere. If ndarray, each row must
            represent a vertex in the sphere.
        points_rh : BSPolyData or ndarray, shape = (n_rh, 3), optional
            Sphere for the right hemisphere. If ndarray, row must
            represent a vertex in the sphere. Default is None.

        Returns
        -------
        self : object
            Returns self.

        """

        spin_idx = _generate_spins(points_lh, points_rh=points_rh,
                                   unique=self.unique, n_rep=self.n_rep,
                                   random_state=self.random_state,
                                   surface_algorithm=self.surface_algorithm)

        self.spin_lh_ = spin_idx['lh']
        self.spin_rh_ = spin_idx.pop('rh', None)
        return self

    def randomize(self, x_lh, x_rh=None):
        """ Generate random samples from `x_lh` and `x_rh`.

        Parameters
        ----------
        x_lh : ndarray, shape = (n_lh,) or (n_lh, n_feat)
            Array of variables arranged in columns, where `n_feat` is the
            number of variables.
        x_rh : ndarray, shape = (n_rh,) or (n_rh, n_feat), optional
            Array of variables arranged in columns for the right hemisphere.
            Default is None.

        Returns
        -------
        rand_lh : ndarray, shape = (n_rep, n_lh, n_feat)
            Permutations of `x_rh`. If ``n_feat == 1``, shape = (n_rep, n_lh).
        rand_lh : ndarray, shape = (n_rep, n_rh, n_feat)
            Permutations of `x_rh`. If ``n_feat == 1``, shape = (n_rep, n_rh).
            None if `x_rh` is None. Only if `spin_rh_` is not None.

        """

        rand_lh = x_lh[self.spin_lh_]
        if self.spin_rh_ is None:
            return rand_lh

        rand_rh = None
        if x_rh is not None:
            rand_rh = x_rh[self.spin_rh_]

        return rand_lh, rand_rh

max_stars_repo_name: vinferrer/BrainSpace

max_issues_repo_name: vinferrer/BrainSpace

max_forks_repo_name: vinferrer/BrainSpace

reason: import numpy,from scipy


text:
import numpy as np

def get_feature_subset(x, index):
    """Get the feature subset A_{s(index)} for a given example and a
    permutation s.

    Parameters
    -------
    x : array-like of shape (1, n_features)
        Example, where n_features is the number of features.

    index : int
        Index for the feature subset.

    Returns
    -------
    result : frozenset
        Feature subset.
    """

    result = set()

    m = len(x)
    sorted_x = np.sort(x)

    permutation = _get_permutation_for_x(x, sorted_x)

    # compute feature subset {\sigma(i),...,\sigma(m)}
    for i in range(index, m + 1):
        result.add(permutation[i])

    return frozenset(result)

def _get_permutation_for_x(x, sorted_x):
    """Get permutation for an example.

    Get a dictionary describing a permutation s of x with
    x_s(1) <= ... <= x_s(m). The dictionary maps the values
    1,...,m to the corresponding position of the permutation,
    where m is the number of features.

    Parameters
    -------
    x : array-like of shape (1, n_features)
        Example, where n_features is the number of features.

    sorted_x : array-like of shape (1, n_features)
        Sorted example, where n_features is the number of features.

    Returns
    -------
    permutation : dict
        Dictionary describing the permutation.
    """

    permutation = dict()

    x_copy = x.copy()

    for i in range(len(x)):
        value = sorted_x[i]

        # TODO which index to choose if there is a tie between values
        index = np.where(x_copy == value)[0][0]

        permutation[i + 1] = index + 1
        x_copy[index] = -9999  # TODO choose dummy value

    return permutation


text:
import logging
import os
import sys
import time

import numpy as np
import torch

sys.path.append(os.path.join(os.path.dirname(__file__), "../.."))

# "get_" 는 util에 정의된 함수, "cls_" 는 classification module에 정의된 함수

from util import get_dataloader, get_optimizer, get_scheduler

# Logger
logger = logging.getLogger("Logger")

class task_module:
    def __init__(self, save_name, writer):

        # Basic Setting
        self.save_name = save_name
        self.save_path = writer.file_writer.get_logdir() if writer is not None else None
        self.writer = writer
        self.best_score = 0

        # Setup Model
        # Don't need to model.to(device)
        self.network = torch.nn.Module()

        # Setup Optimizer and Scheduler
        self.optimizer = torch.nn.Module()
        self.scheduler = torch.nn.Module()

    def _save_state(self, i_tr, score, best=False, additional={}):
        self.network.eval()

        state = {
            "iter": i_tr + 1,
            "network_state": self.network.state_dict(),
            "optimizer_state": self.optimizer.state_dict(),
            "scheduler_state": self.scheduler.state_dict(),
            "score": score,
        }

        state.update(additional)

        name = "best" if best else i_tr + 1
        save_path = os.path.join(
            self.save_path,
            f"{self.save_name}_{name}_model.pkl",
        )
        torch.save(state, save_path)

        logger.info(f"[SAVE] {self.save_name}_{name}_model")

    def _load_state(self, cfg_load):

        if cfg_load is not None:

            assert cfg_load["file"], "No State File Given"
            assert os.path.isfile(
                cfg_load["file"]
            ), f'No State File Found At { cfg_load["file"]}'

            state = torch.load(cfg_load["file"])

            if cfg_load.get("backbone_only", False):
                network_state_dict = self.network.state_dict()
                loaded_state_dict = self._convert_state(state["network_state"])
                # construct backbone only state_dict
                loaded_state_dict_backbone = {}
                for key in loaded_state_dict:
                    if "backbone" in key:
                        loaded_state_dict_backbone[key] = loaded_state_dict[key]
                # update network_state_dict with backbone only state_dict
                network_state_dict.update(loaded_state_dict_backbone)
                self.network.load_state_dict(network_state_dict)
            else:
                network_state = self._convert_state(
                    state.get("network_state", state.get("model_state"))
                )
                self.network.load_state_dict(self._convert_state(network_state))

            # loading specific layers sample
            # self.network.load_state_dict(self._convert_state(state["network_state"]))
            logger.info(
                f'[{"NETWORK".center(9)}] Trained Model Loaded from: {cfg_load["file"]}'
            )

            if cfg_load.get("optimizer", False):
                self.optimizer.load_state_dict(state["optimizer_state"])
                self.scheduler.load_state_dict(state["scheduler_state"])
                logger.info("Optimizer loaded")

            if cfg_load.get("iteration", False):
                self.start_iter = state["iter"] - 1
                logger.info(f"Start_iter loaded: {self.start_iter + 1}")

    def _convert_state(self, state):
        # erase '.module' for multi-gpu case
        state_modified = {}

        for k, v in state.items():
            split_k = k.split(".")
            name = ""

            for k_item in split_k:
                if k_item != "module" and k_item != "network":
                    name += k_item + "."

            state_modified[name[:-1]] = v

        return state_modified

    def _check_best_score(self, i_tr, score, key, additional={}):

        assert type(score) == dict, "Score should be passed to _check_bst_score as Dictionary"

        for k, v in score.items():
            logger.info(f"{k.rjust(18)}: {v}")
            if self.writer is not None and np.isscalar(v):
                self.writer.add_scalar(f"Metric/{k}", v, i_tr + 1)

        current_score = score[key]
        if current_score >= self.best_score:
            self.best_score = current_score

            if self.writer is not None:
                self._save_state(i_tr, current_score, best=True, additional=additional)

    def train_step(self, image, label, i):
        raise Exception("You called dummy 'train_step' function! Write your own 'train_step'")

    def validation_step(self, image, label):
        raise Exception(
            "You called dummy 'validation_step' function! Write your own 'validation_step'"
        )

    def evaluation_step(self, image, label):
        raise Exception(
            "You called dummy 'evaluation_step' function! Write your own 'evaluation_step'"
        )

    def on_save_interval(self, i):
        logger.info("[NOTHING SAVED] you didn't code 'on_save_interval' function")

    def on_validation_end(self, i):
        logger.info("[NOTHING HAPPENED] you didn't code 'on_validation_end' function")

if __name__ == "__main__":
    print(1)


text:
# coding:utf-8
"""
"""
import os
import sys
import csv
import fire
import time
import json
import tarfile
import torch
import torch.nn as nn
from torchvision import models
from torchvision.models.vgg import VGG
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder

import pretrainedmodels

from glob import glob
from PIL import Image
import numpy as np

data_dir = 'D:/workspace/udacity/MLND/capstone/distracted_driver_detection/data/'

ON_CLIENT = True
DEVICE_ID = 0
EPOCHS = 10
DROPOUT = 0.5
NUM_CLASSES = 10
LEARN_RATE = 0.001
CROP_SIZE = (320, 480)
IMG_SIZE = (224, 224)
BATCH_SIZE = 1
IMG_MEAN = (0.3184719383716583, 0.3813590109348297, 0.37875279784202576)
IMG_VAR = (0.3184719383716583, 0.3813589811325073, 0.37875282764434814)
IMG_STD = (0.5643, 0.6175, 0.6154)
MASKER_KEEP_LAYER = None
FAKE_TRAIN = False
GENERATE_MASK_IMAGE = False

CUDA = torch.cuda.is_available()
# DEVICE = torch.device("cuda" if CUDA else "cpu")

DEVICE_MASKER = None
DEVICE_CLASSIFIER = None
DEVICE_DEBUG = None
DEVICE_TEST = None

def device_setting(on_client):
    if on_client is True:
        DEVICE_MASKER = torch.device("cuda:0")
        DEVICE_CLASSIFIER = torch.device("cpu")
        DEVICE_DEBUG = torch.device("cuda:0")
        DEVICE_TEST = torch.device("cpu")
    else:
        DEVICE_MASKER = torch.device("cuda:1")
        DEVICE_CLASSIFIER = torch.device("cuda:2")
        DEVICE_DEBUG = torch.device("cuda:3")
        DEVICE_TEST = torch.device("cpu")

ddd_mask_images = [torch.zeros(IMG_SIZE)] * 10

ranges = {
    'vgg11': ((0, 3), (3, 6),  (6, 11),  (11, 16), (16, 21)),
    'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
    'vgg16': ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)),
    'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37))
}

cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

class VGGNet(VGG):
    def __init__(self, pretrained=True, model='vgg16'):
        super().__init__(make_layers(cfg[model]))
        self.ranges = ranges[model]

        if pretrained:
            exec("self.load_state_dict(models.%s(pretrained=True).state_dict())" % model)

    def forward(self, x):
        output = {}
        # get the output of each maxpooling layer (5 maxpool in VGG net)
        for idx, (begin, end) in enumerate(self.ranges):
            # self.ranges = ((0, 5), (5, 10), (10, 17), (17, 24), (24, 31)) (vgg16 examples)
            for layer in range(begin, end):
                x = self.features[layer](x)
            output["x%d" % (idx+1)] = x

        return output

class MaskNet(nn.Module):
    def __init__(self, model=None, dropout=DROPOUT, num_classes=NUM_CLASSES, masker_keep_layer=None):
        super(MaskNet, self).__init__()

        if model is None:
            self.model = VGGNet()
        else:
            self.model = model
        # self.model.to(DEVICE_MASKER)

        self.dropout = dropout
        self.num_classes = num_classes
        self.masker_keep_layer = masker_keep_layer

        for name, param in self.model.named_parameters():
            param.requires_grad = self.keep_layer(name)

        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(
            512, 512, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn1 = nn.BatchNorm2d(512)
        self.deconv2 = nn.ConvTranspose2d(
            512, 256, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn2 = nn.BatchNorm2d(256)
        self.deconv3 = nn.ConvTranspose2d(
            256, 128, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.deconv4 = nn.ConvTranspose2d(
            128, 64, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn4 = nn.BatchNorm2d(64)
        self.deconv5 = nn.ConvTranspose2d(
            64, 32, kernel_size=3, stride=2, padding=1, dilation=1, output_padding=1)
        self.bn5 = nn.BatchNorm2d(32)

        #self.masker = nn.Conv2d(32, 3, kernel_size=1)
        self.masker = nn.Sequential(
            nn.Conv2d(32, 16, kernel_size=1),
            nn.Conv2d(16, 8, kernel_size=1),
            nn.Conv2d(8, 3, kernel_size=1)
        )

        self.excessive = nn.Conv2d(3, 1, kernel_size=1)

        self.classifier = nn.Sequential(
            nn.Linear(50176, 4096),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(4096, self.num_classes)
        )

    def forward(self, x):
        output = self.model(x)
        x5 = output['x5']
        x4 = output['x4']
        x3 = output['x3']
        x2 = output['x2']
        x1 = output['x1']

        mask = self.bn1(self.relu(self.deconv1(x5)))
        mask = mask + x4
        mask = self.bn2(self.relu(self.deconv2(mask)))
        mask = mask + x3
        mask = self.bn3(self.relu(self.deconv3(mask)))
        mask = mask + x2
        mask = self.bn4(self.relu(self.deconv4(mask)))
        mask = mask + x1
        mask = self.bn5(self.relu(self.deconv5(mask)))
        self.mask = self.masker(mask)

        # use filter to mask
        #self.tensor_filter(self.mask)

        #mask_cpu = self.mask.cpu()
        #mask = transforms.functional.to_pil_image(mask_cpu)
        # img.show()

        self.feature_image = self.excessive(self.mask)
        x = self.feature_image.view(self.feature_image.size(0), -1)
        pred = self.classifier(x)

        return pred

    def tensor_filter(self, tensors_items):
        for item in tensors_items:
            for tensor in item:
                mean = tensor.mean()
                tensor[tensor<mean] = 0

    def keep_layer(self, name):
        if self.masker_keep_layer is None:
            return False

        if len(self.masker_keep_layer) == 0:
            return False

        for kl in self.masker_keep_layer:
            if kl in name:
                return True
        return False

class ClassifierNet(nn.Module):
    def __init__(self, model=None, dropout=DROPOUT, num_classes=NUM_CLASSES):
        super(ClassifierNet, self).__init__()

        if model is None:
            self.model = pretrainedmodels.__dict__[
                'xception'](pretrained='imagenet')
        else:
            self.model = model
        # self.model.to(DEVICE_CLASSIFIER)

        self.dropout = dropout
        self.num_classes = num_classes

        self.classifier_transorms = transforms.Compose([
            transforms.Lambda(lambda tensors:
                              [transforms.Resize((299, 299))(transforms.ToPILImage()(tensor)) for tensor in tensors]),
            transforms.Lambda(lambda tensors:
                              torch.stack([transforms.ToTensor()(tensor) for tensor in tensors]))
        ])

        # extract freature
        # self.feature = nn.Sequential(*list(self.model.children())[:-1])

        for param in self.model.parameters():
            param.requires_grad = False

        in_dim = self.model.last_linear.in_features

        """
        finetune = nn.Sequential(
            nn.Dropout(self.dropout),
            nn.Linear(in_dim, 4096),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(4096, 1000),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(1000, self.num_classes)
        )
        """
        finetune = nn.Sequential(
            nn.Dropout(self.dropout),
            nn.Linear(in_dim, 1000),
            nn.ReLU(True),
            nn.Dropout(self.dropout),
            nn.Linear(1000, self.num_classes)
        )

        self.model.last_linear = finetune

    def forward(self, x):
        x = x.cpu()
        x = self.classifier_transorms(x)
        x = torch.as_tensor(x, device=DEVICE_CLASSIFIER)
        x = self.model(x)
        return x

# ========================================================================================================

def make_transorms():
    train_transorms = transforms.Compose([
        transforms.CenterCrop(CROP_SIZE),
        #transforms.RandomAffine(degrees=10, translate=(0.1, 0.3)),
        #transforms.RandomHorizontalFlip(),
        #transforms.RandomResizedCrop(IMG_SIZE),
        transforms.Resize(IMG_SIZE),
        # transforms.FiveCrop(CROP_SIZE),
        # transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
        transforms.ToTensor(),
        transforms.Normalize(IMG_MEAN, IMG_STD)
    ])

    valid_transorms = transforms.Compose([
        transforms.CenterCrop(CROP_SIZE),
        transforms.Resize(IMG_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(IMG_MEAN, IMG_STD)
    ])

    return train_transorms, valid_transorms

def get_datasets(data_dir, train_transorms, valid_transorms=None, batch_size=BATCH_SIZE):
    valid_gen = None
    train_gen = ImageFolder(os.path.join(
        data_dir, 'train/'), transform=train_transorms)
    if valid_transorms is not None:
        valid_gen = ImageFolder(os.path.join(
            data_dir, 'valid/'), transform=valid_transorms)

    train_dataiter = DataLoader(train_gen, batch_size=batch_size, shuffle=True)
    return train_dataiter, valid_gen

def make_model(masker_path=None, classifier_path=None,
               masker_dropout=DROPOUT, classifier_dropout=DROPOUT,
               masker_keep_layer=MASKER_KEEP_LAYER, on_client=ON_CLIENT):
    masker_model = MaskNet(
        dropout=masker_dropout, num_classes=NUM_CLASSES, masker_keep_layer=masker_keep_layer)
    classifier_model = ClassifierNet(
        dropout=classifier_dropout, num_classes=NUM_CLASSES)

    if masker_path is not None:
        if on_client:
            masker_model.load_state_dict(torch.load(
                masker_path, map_location=torch.device('cpu')))
        else:
            masker_model.load_state_dict(torch.load(masker_path))

    if classifier_path is not None:
        if on_client:
            classifier_model.load_state_dict(torch.load(
                classifier_path, map_location=torch.device('cpu')))
        else:
            classifier_model.load_state_dict(torch.load(classifier_path))

    return masker_model, classifier_model

def valid(masker_model, classifier_model, masker_criterion, classifier_criterion, datasets):
    masker_model.eval()
    classifier_model.eval()

    loss = 0
    correct = 0

    mask_loss = 0
    mask_correct = 0

    with torch.no_grad():
        for i, data in enumerate(datasets):
            image = data[0]
            label = torch.tensor(data[1], dtype=torch.long)

            image = torch.unsqueeze(image.to(DEVICE_MASKER), 0)
            label = torch.unsqueeze(label.to(DEVICE_CLASSIFIER), 0)
            mask_label = torch.as_tensor(label, device=DEVICE_MASKER)

            masker_pred = masker_model(image)
            mask_loss += masker_criterion(masker_pred, mask_label).item()
            masker_pred = masker_pred.max(1, keepdim=True)[1]
            mask_correct += masker_pred.eq(
                mask_label.view_as(masker_pred)).sum().item()

            masker_output = masker_model.mask
            classifier_input = torch.as_tensor(
                masker_output, device=DEVICE_CLASSIFIER)
            classifier_pred = classifier_model(classifier_input)

            loss += classifier_criterion(classifier_pred, label).item()
            pred = classifier_pred.max(1, keepdim=True)[1]
            correct += pred.eq(label.view_as(pred)).sum().item()
    loss /= len(datasets)
    mask_loss /= len(datasets)
    return loss, correct, mask_loss, mask_correct

def train_nn(masker_model, classifier_model, masker_criterion, classifier_criterion,
             masker_optimizer, classifier_optimizer, train_dataiter, fake_masker=False,
             generate_mask_image=GENERATE_MASK_IMAGE):
    masker_model.train()
    classifier_model.train()

    for batch_index, train_data in enumerate(train_dataiter):
        train_images_masker = train_data[0]
        train_labels_masker = train_data[1]
        train_labels_classifier = train_data[1]

        # copy CPU data to GPU
        train_images_masker = train_images_masker.to(DEVICE_MASKER)
        train_labels_masker = train_labels_masker.to(DEVICE_MASKER)

        train_labels_classifier = train_labels_classifier.to(DEVICE_CLASSIFIER)

        masker_optimizer.zero_grad()
        classifier_optimizer.zero_grad()

        masker_pred = masker_model(train_images_masker)
        masker_output = masker_model.mask

        if generate_mask_image:
            masker_feature_image = masker_model.feature_image
            for i, feature in enumerate(masker_feature_image):
                label = train_data[1][i]
                feature_cpu = feature.cpu()
                assert ddd_mask_images[label].size() == feature_cpu.size()
                ddd_mask_images[label] += feature_cpu

        classifier_input = torch.as_tensor(
            masker_output, device=DEVICE_CLASSIFIER)
        classifier_pred = classifier_model(classifier_input)

        classifier_loss = classifier_criterion(
            classifier_pred, train_labels_classifier)

        # use classifier's pred to train masknet
        if fake_masker is True:
            fake_masker_pred = torch.as_tensor(
                classifier_pred, device=DEVICE_MASKER)
            masker_loss = masker_criterion(
                fake_masker_pred, train_labels_masker)
        else:
            masker_loss = masker_criterion(masker_pred, train_labels_masker)

        classifier_loss.backward(retain_graph=True)
        masker_loss.backward()

        classifier_optimizer.step()
        masker_optimizer.step()
    return masker_loss, classifier_loss

def train(data_dir=data_dir, epochs=EPOCHS, batch_size=BATCH_SIZE,
          masker_path=None, classifier_path=None, masker_learn_rate=LEARN_RATE, classifier_learn_rate=LEARN_RATE,
          masker_dropout=DROPOUT, classifier_dropout=DROPOUT, masker_keep_layer=MASKER_KEEP_LAYER, fake_train=FAKE_TRAIN,
          generate_mask_image=GENERATE_MASK_IMAGE, save_dir='../model', show_model=False, describe=None,
          zipit=False, on_client=ON_CLIENT):
    """
    trainning...
    """
    screen_print = []
    def printf(str):
        if str is None:
            print()
            screen_print.append('\n')
        else:
            print(str)
            screen_print.append(str+'\n')

    data_dir = os.path.abspath(data_dir)
    if masker_path is not None:
        masker_path = os.path.abspath(masker_path)
    if classifier_path is not None:
        classifier_path = os.path.abspath(classifier_path)
    save_dir = os.path.abspath(save_dir)
    fake_train = bool(fake_train)

    if masker_keep_layer is not None:
        if isinstance(masker_keep_layer, int):
            masker_keep_layer = [masker_keep_layer]
        elif isinstance(masker_keep_layer, str):
            masker_keep_layer = [int(x) for x in masker_keep_layer.split(',')]
        elif isinstance(masker_keep_layer, list):
            pass
        else:
            masker_keep_layer = None

    if describe is not None:
        printf('===== DESCRIBE INFO =====')
        printf(describe)

    printf('===== TRAIN INFO =====\n')
    printf('data dir = {}\n'.format(data_dir))
    printf('masker path = {}\n'.format(
        masker_path if masker_path is not None else '[new model]'))
    printf('classifier path = {}\n'.format(
        classifier_path if classifier_path is not None else '[new model]'))
    printf('masker dropout = {}\n'.format(masker_dropout))
    printf('classifier dropout = {}\n'.format(classifier_dropout))
    printf('masker keep layers = {}\n'.format(masker_keep_layer))
    printf('epochs = {}\n'.format(epochs))
    printf('masker learn rate = {}\n'.format(masker_learn_rate))
    printf('classifier learn rate = {}\n'.format(classifier_learn_rate))
    printf('batch size = {}\n'.format(batch_size))
    printf('fake_train = {}\n'.format(fake_train))
    printf('save dir = {}\n'.format(save_dir))
    printf('======================\n')

    train_transorms, valid_transorms = make_transorms()
    train_dataiter, valid_datasets = get_datasets(
        data_dir, train_transorms, valid_transorms, batch_size)
    print('[INFO]generate datasets ok')

    masker_model, classifier_model = make_model(
        masker_path, classifier_path, masker_dropout, classifier_dropout,
        masker_keep_layer, on_client)
    masker_model.to(DEVICE_MASKER)
    classifier_model.to(DEVICE_CLASSIFIER)
    print('[INFO]load model ok')

    # for name, param in classifier_model.named_parameters():
    #    if param.requires_grad is True:
    #        print('trainning --- {}'.format(name))
    #    else:
    #        print(name)
    #    print(param.grad)

    if show_model is True:
        print(masker_model)
        print()
        print(classifier_model)

    # RMSprop
    masker_criterion = torch.nn.CrossEntropyLoss()
    masker_optimizer = torch.optim.Adam(
        masker_model.parameters(), lr=masker_learn_rate)

    classifier_criterion = torch.nn.CrossEntropyLoss()
    classifier_optimizer = torch.optim.Adam(classifier_model.parameters(),
                                            lr=classifier_learn_rate)

    best_valid_accuracy = -float('inf')
    for epoch in range(1, epochs+1):
        printf('Trainning: {}\n'.format(epoch))

        masker_train_loss, classifier_train_loss = train_nn(
            masker_model, classifier_model, masker_criterion, classifier_criterion,
            masker_optimizer, classifier_optimizer, train_dataiter, fake_train,
            generate_mask_image)

        printf('Train Epoch: {}\t Masker Loss: {:.6f}\t Classifier Loss: {:.6f}\n'.format(
            epoch, masker_train_loss.item(), classifier_train_loss.item()))

        valid_loss, correct, mask_loss, mask_correct = valid(
            masker_model, classifier_model, masker_criterion, classifier_criterion, valid_datasets)

        printf('Valid set: Average of masker loss: {:.4f}, Accuracy of masker: {}/{} ({:.0f}%)\n'.format(
            mask_loss, mask_correct, len(valid_datasets), 100. * mask_correct / len(valid_datasets)))

        printf('Valid set: Average of classifier loss: {:.4f}, Accuracy of classifier: {}/{} ({:.0f}%)\n'.format(
            valid_loss, correct, len(valid_datasets), 100. * correct / len(valid_datasets)))

        curr_valid_accuracy = correct / len(valid_datasets)
        if (curr_valid_accuracy > best_valid_accuracy):
            masker_save_path = os.path.join(
                save_dir, 'xvgg_masker_state_dict.pth')
            torch.save(masker_model.state_dict(), masker_save_path)

            classifier_save_path = os.path.join(
                save_dir, 'xvgg_classifier_state_dict.pth')
            torch.save(classifier_model.state_dict(), classifier_save_path)

    timestamp = time.asctime().replace(' ', '_').replace(':', '_')

    print_save_path = os.path.join(save_dir, timestamp+'.txt')
    f = open(print_save_path, 'w')
    for line in screen_print:
        f.write(line)
    f.close()

    new_masker_save_path = os.path.join(
        save_dir, 'xvgg_masker_'+timestamp+'.pth')
    os.rename(masker_save_path, new_masker_save_path)
    print('save masker model to {}'.format(new_masker_save_path))

    new_classifier_save_path = os.path.join(
        save_dir, 'xvgg_classifier_'+timestamp+'.pth')
    os.rename(classifier_save_path, new_classifier_save_path)
    print('save classifier model to {}'.format(new_classifier_save_path))

    new_checkpoint_path = os.path.join(save_dir, 'xvgg_'+timestamp+'.json')
    with open(new_checkpoint_path, "w") as f:
        state_dict = {
            'describe': describe,
            'masker_path': new_masker_save_path,
            'classifier_path': new_classifier_save_path,
            'info_path': print_save_path,
            'masker_dropout': masker_dropout,
            'classifier_dropout': classifier_dropout,
            'masker_keep_layer': masker_keep_layer,
            'fake_train': fake_train
        }
        json.dump(state_dict, f)

    mask_image_path = None
    if generate_mask_image is True:
        img_dict = {}
        mask_image_path = os.path.join(save_dir, 'xvgg_mask_image'+timestamp+'.npy')
        for i, img in enumerate(ddd_mask_images):
            img_dict['c{}'.format(i)] = img.numpy()
        np.save(mask_image_path, img_dict)

    if zipit is True:
        tarfile_path = os.path.join(save_dir, 'xvgg_'+timestamp+'.tar.gz')
        with tarfile.open(tarfile_path, "w:gz") as tar:
            tar.add(print_save_path)
            tar.add(new_checkpoint_path)
            tar.add(new_masker_save_path)
            tar.add(new_classifier_save_path)
            if mask_image_path is not None:
                tar.add(mask_image_path)
            tar.close()

    return True

def model_print():
    model = MaskNet()
    dummy_input = torch.randn(1, 3, 224, 224)
    with SummaryWriter(log_dir='../ddd_model', comment='Masknet')as w:
        w.add_graph(model, dummy_input)

# ========================================================================================================

class TestDataset(Dataset):
    def __init__(self, data_dir):
        super(TestDataset, self).__init__()
        self.image_paths = glob(os.path.join(data_dir, '*.jpg'))
        _, test_transorms = make_transorms()
        self.transorms = test_transorms
        print('{} number of test image loaded.'.format(len(self.image_paths)))

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

    def __getitem__(self, index):
        # assert index >= len(self.image_paths), 'index, out of range'

        path = self.image_paths[index]
        image = Image.open(path)
        image_tensor = self.transorms(image)
        # image_tensor = transforms.functional.to_tensor(image)
        filename = os.path.basename(path)

        return (image_tensor, filename)

def generate_csv_handle(csv_path='./test.csv'):
    header = ['img', 'c0', 'c1', 'c2', 'c3',
              'c4', 'c5', 'c6', 'c7', 'c8', 'c9']
    csv_file = open(csv_path, 'w')
    csv_write = csv.writer(csv_file)
    csv_write.writerow(header)
    return csv_file, csv_write

def test(data_dir=data_dir, masker_path=None, classifier_path=None, batch_size=256,
         csv_path='./result.csv', on_client=ON_CLIENT):
    data_dir = os.path.abspath(data_dir)
    if masker_path is not None:
        masker_path = os.path.abspath(masker_path)
    if classifier_path is not None:
        classifier_path = os.path.abspath(classifier_path)
    csv_path = os.path.abspath(csv_path)

    csv_file, csv_write = generate_csv_handle(csv_path)

    print('===== TEST INFO =====')
    print('data dir = {}'.format(data_dir))
    print('masker path = {}'.format(masker_path))
    print('classifier path = {}'.format(classifier_path))
    print('batch size = {}'.format(batch_size))
    print('csv path = {}'.format(csv_path))
    print('======================')

    print()

    test_gen = TestDataset(os.path.join(data_dir, 'test/'))

    if on_client is True:
        batch_size = 3
    test_dataiter = DataLoader(test_gen, batch_size=batch_size, shuffle=False)
    print('[INFO]generate datasets ok')

    masker_model, classifier_model = make_model(masker_path, classifier_path)

    # named_parameters, named_children, named_modules
    # for name, child in masker_model.named_parameters():
    #    print(name)

    masker_model.to(DEVICE_TEST)
    classifier_model.to(DEVICE_TEST)
    print('[INFO]load model ok')

    # extract_image(masker_model, 0)

    masker_model.eval()
    classifier_model.eval()

    with torch.no_grad():
        for i, data in enumerate(test_dataiter):
            print('Test {}th batch'.format(i))
            print()

            images = data[0]
            filenames = data[1]

            masker_pred = masker_model(images)
            masker_output = masker_model.mask

            # masker_images = masker_model.excessive(masker_output)
            # masker_image = transforms.functional.to_pil_image(masker_output[0])
            # masker_image.show()

            classifier_pred = classifier_model(masker_output)

            prob_list = torch.nn.functional.softmax(classifier_pred, dim=0)
            prob_list = prob_list.numpy()
            row = []
            for (filename, prob) in zip(filenames, prob_list):
                row.append(filename)
                row.extend(prob)
                csv_write.writerow(row)
                row.clear()

    csv_file.close()
    print('save csv to {}'.format(csv_path))

# ========================================================================================================

def run(command='train', data_dir=data_dir, checkpoint_path=None, epochs=EPOCHS,
        masker_learn_rate=LEARN_RATE, classifier_learn_rate=LEARN_RATE,
        masker_dropout=DROPOUT, classifier_dropout=DROPOUT, masker_keep_layer=MASKER_KEEP_LAYER,
        fake_train=FAKE_TRAIN, generate_mask_image=GENERATE_MASK_IMAGE, batch_size=BATCH_SIZE, save_dir='../model',
        show_model=False, describe=None, csv_path='./result.csv', zipit=False, on_client=ON_CLIENT):
    """
    command: train, test
    """
    device_setting(on_client)

    masker_path = None
    classifier_path = None
    if checkpoint_path is not None:
        with open(checkpoint_path, 'r') as load_f:
            config = json.load(load_f)
            masker_path = config['masker_path']
            classifier_path = config['classifier_path']
            masker_keep_layer = config['masker_keep_layer']

    if command == 'train':
        train(data_dir=data_dir, epochs=epochs, batch_size=batch_size,
              masker_path=masker_path, classifier_path=classifier_path,
              masker_learn_rate=masker_learn_rate, classifier_learn_rate=classifier_learn_rate,
              masker_dropout=masker_dropout, classifier_dropout=classifier_dropout,
              masker_keep_layer=masker_keep_layer, fake_train=fake_train,
              save_dir=save_dir, show_model=show_model, describe=describe, zipit=zipit, on_client=on_client)
        return True

    if command == 'test':
        test(data_dir=data_dir, masker_path=masker_path, classifier_path=classifier_path,
             batch_size=batch_size, csv_path=csv_path, on_client=on_client)
        return True

    print('invalid command')
    return False

# ========================================================================================================

#test(data_dir=data_dir+'../data', csv_path='./result.csv')
if __name__ == '__main__':
    fire.Fire()