# -*- coding: utf-8 -*-
"""
This file provides almost all plot functions for the paper titled `The spread
of low-credibility content by social bots`.

Functions are named by the appearing order of the generated figures in the
paper, conjuncting with the key words of the figure captions. For example,
function `m1a_article_distribution_by_tweets` is the plot function for
Figure 1 (a) in the main paper and function
`s13_14_bots_retweeter_and_tweeter` is the plot function for Supplementary
Figure 13 and 14 in the sumpplementary paper.

Please note that the input for almost all the functions are intermediate
results that could be produced from the raw BIG ONE file. We don't provide these
data processing functions, either because they are fetched by SQL query on Hoaxy
database server, or because they need additional tweet information that could
violate Twitter Policies. However, we do give the instructions of how you could
produce these intermediate results.

Another big file is about the calibrated bot score for the accounts used in
our analysis. This big file is a union of several analysis. As some accounts
were suspended during our analysis, in order to keep the exact number as the
paper, for each specified analysis, we use the bot scores when we were
doing the analysis instead of join this union set (at that time, we did not
know this union bot score set). Otherwise, the statistic patterns are still the
same, but the exact numbers would be a slightly different.

There are a few figures that are generated by Gephi, we would provide the gml
files.

As time goes, there are more sources are listed by more organizations. We
compiled a large source table to represent it, which is much large then the
sources we used to do the analysis. The consensus number is base on this
large table, located at 'dataset/consensus.csv'.

Finally, we test these scripts on python 2.7. To use this scripts, you also
need to install the following packages:
    - matplotlib
    - numpy
    - pandas
    - scipy
    - python-ternary (https://github.com/marcharper/python-ternary)
"""

#
# written by Chengcheng Shao <sccotte@gmail.com>

import logging
from datetime import date, timedelta
from os.path import dirname, join, realpath

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import gridspec
from scipy.stats import ks_2samp, mannwhitneyu, ttest_ind

import ternary
from ccdf import ccdf
from ternary_heatmap_helper import colorbar_hack, heatmap_density

logger = logging.getLogger(__name__)
# project root
BASE_DIR = dirname(dirname(realpath(__file__)))
DATA_DIR = join(BASE_DIR, 'dataset')
FIGURES_DIR = join(BASE_DIR, 'figures')
# default figure size
FIGSIZE = (4, 3)
# color for low-credibility
C1 = '#1F78B4'
# color for fact-checking
C2 = '#FF7F00'

def get_data_file(fn, subfolder=None):
    """Return the absolute path of a dataset file.
    """
    if subfolder is not None:
        return join(DATA_DIR, subfolder, fn)
    else:
        return join(DATA_DIR, fn)

def get_output_file(fn):
    """Return the absolute path of a figure file.
    """
    return join(FIGURES_DIR, fn)

def nplog(a, base=10):
    """Return the log value on the base of `base`.
    """
    return (np.log(a) / np.log(base))

# Figure 1a in the main paper
def m1a_article_distribution_by_tweets(
        fn1='M1a_ntweets_per_article.lc.csv',
        fn2='M1a_ntweets_per_article.ff.csv',
        ofn='M1a-article-distribution-by-tweets.pdf',
        density=True,
        nbins=15,
        base=2,
        figsize=(4, 3)):
    """
    Caption in the main paper:
    --------------------------
    Figure 1: Online virality of content. (a) Probability distribution
    (density function) of the number of tweets for articles from both
    low-credibility and fact-checking sources. The distributions of the number
    of accounts sharing an article are very similar (see Supplementary
    Information).

    Input data and instructions:
    ----------------------------
    `fn1`: number of tweets for each article from low-credibility sources,
    `fn2`: number of tweets for each article from fact-checking sources.
    """
    fn1 = get_data_file(fn1, subfolder='consensus1')
    fn2 = get_data_file(fn2, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    # Note, set max bin points slightly larger than max data value
    # because the calculation may make right most bin points
    # a little smaller than max data value
    bins1 = np.logspace(
        nplog(df1.counts.min(), base),
        nplog(df1.counts.max() + 1, base),
        nbins + 1,
        base=base)
    bins2 = np.logspace(
        nplog(df2.counts.min(), base),
        nplog(df2.counts.max() + 1, base),
        nbins + 1,
        base=base)
    if density is not True:
        w1 = np.ones_like(df1.counts.values) / len(df1.counts.values)
        w2 = np.ones_like(df2.counts.values) / len(df2.counts.values)
    else:
        w1 = None
        w2 = None

    h1, bins1 = np.histogram(
        df1.counts.values, bins=bins1, weights=w1, density=density)
    h2, bins2 = np.histogram(
        df2.counts.values, bins=bins2, weights=w2, density=density)
    f, ax = plt.subplots(figsize=figsize)
    ax.set_xlabel('(a) Number of Tweets')
    ax.set_ylabel('PDF')
    ax.set_xscale('log')
    ax.set_yscale('log')
    ax.plot(
        bins1[:-1], h1, c=C1, marker='o', alpha=0.6, label='Low-credibility')
    ax.plot(bins2[:-1], h2, c=C2, marker='s', alpha=0.6, label='Fact checking')
    ax.legend()
    plt.tight_layout()
    plt.savefig(ofn)

# help function for ternary plot (m2a_ternary_types_of_tweets_low_credibility)
def prepare_ternary(figsize, scale):
    """Help function to ternary plot"""
    fig, ax = plt.subplots(figsize=figsize)
    tax = ternary.TernaryAxesSubplot(ax=ax, scale=scale)
    ax.axis('off')
    gm = 0.1 * scale
    blw = 1
    tlw = 1
    # Draw Boundary and Gridlines
    tax.boundary(linewidth=blw)
    tax.gridlines(color='grey', multiple=gm, alpha=0.8)
    # Set Axis labels and Title
    tax.bottom_axis_label(
        r"Retweets $\rightarrow$", offset=-0.08, fontsize='small')
    tax.right_axis_label(r"$\leftarrow$Replies", offset=0.2, fontsize='small')
    tax.left_axis_label(r"$\leftarrow$Tweets", offset=0.2, fontsize='small')

    # ticks
    locations = range(0, scale + 1, 4)
    ticks = ['{}'.format(x * 10) for x in range(0, 11, 2)]
    tax.ticks(
        axis='lbr',
        ticks=ticks,
        locations=locations,
        linewidth=tlw,
        offset=0.03,
        fsize=9,
        clockwise=False)
    return tax

# Figure2a in the main paper
def m2a_ternary_types_of_tweets_low_credibility(
        fn='M2a_types_of_tweets.lc.csv',
        ofn='M2a-types-of-tweets-low-credibility.pdf',
        ntop=None,
        least_ntweets=100,
        scale=20,
        style='hexagonal',
        figsize=(5, 3.5)):
    """
    Caption in the main paper:
    --------------------------
    Figure 2: Anomalies. The distribution of types of tweet spreading articles
    from (a) low-credibility ...

    Input data and instructions:
    ----------------------------
    `fn`: providing the statistics of the types of tweets for each article.
    You should have the knowledge of how to recogonize the different types
    of tweets.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    df.set_index('id', inplace=True)
    df['ntotal'] = df.sum(axis=1)
    df = df.sort_values('ntotal', ascending=False)
    logger.info('Total artilces %s', len(df))
    if ntop is not None:
        df = df.iloc[:ntop]
    if least_ntweets is not None:
        df = df.loc[df.ntotal >= least_ntweets]
    logger.info('Number of articles to visualize, %s', len(df))

    X = ((df.nretweets + df.nquotes) / df.ntotal * scale).values
    Y = (df.nreplies / df.ntotal * scale).values
    Z = (df.norigins / df.ntotal * scale).values

    data = heatmap_density(X, Y, Z, scale)
    data = pd.Series(data)
    vmin = data.min()
    vmax = data.max()
    data = data.apply(lambda x: np.log10(x + 1))
    if vmin <= 0:
        vmin = 1
    logger.info('vmin=%s, vmax=%r', vmin, vmax)
    data = data.to_dict()
    cmap = plt.cm.get_cmap('gnuplot2_r')
    tax = prepare_ternary(figsize, scale)
    tax.heatmap(
        data,
        cmap=cmap,
        style=style,
        colorbar=False,
    )
    colorbar_hack(
        ax=tax.ax,
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        log_norm=True,
        cbarlabel='Density')
    # tax.scatter(zip(X, Y), s=49, alpha=0.2)
    plt.tight_layout()
    plt.savefig(ofn)

# Figure 2b in the main paper
def m2b_types_of_tweets_factchecking(
        fn='m2b_types_of_tweets.ff.csv',
        ofn='M2b-types-of-tweets-factchecking.pdf',
        ntop=None,
        least_ntweets=100,
        scale=20,
        style='hexagonal',
        figsize=(5, 3.5)):
    """
    Caption in the main paper:
    --------------------------
    Figure 2: Anomalies. The distribution of types of tweet spreading articles
    from (a) low-credibility ...

    Input data and instructions:
    ----------------------------
    `fn`: providing the statistics of the types of tweets for each article.
    You should have the knowledge of how to recogonize the different types
    of tweets.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    df.set_index('id', inplace=True)
    df['ntotal'] = df.sum(axis=1)
    df = df.sort_values('ntotal', ascending=False)
    logger.info('Total artilces %s', len(df))
    if ntop is not None:
        df = df.iloc[:ntop]
    if least_ntweets is not None:
        df = df.loc[df.ntotal >= least_ntweets]
    logger.info('Number of articles to visualize, %s', len(df))

    X = ((df.nretweets + df.nquotes) / df.ntotal * scale).values
    Y = (df.nreplies / df.ntotal * scale).values
    Z = (df.norigins / df.ntotal * scale).values

    data = heatmap_density(X, Y, Z, scale)
    data = pd.Series(data)
    vmin = data.min()
    vmax = data.max()
    data = data.apply(lambda x: np.log10(x + 1))
    if vmin <= 0:
        vmin = 1
    logger.info('vmin=%s, vmax=%r', vmin, vmax)
    data = data.to_dict()
    cmap = plt.cm.get_cmap('gnuplot2_r')
    tax = prepare_ternary(figsize, scale)
    tax.heatmap(
        data,
        cmap=cmap,
        style=style,
        colorbar=False,
    )
    colorbar_hack(
        ax=tax.ax,
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        log_norm=True,
        cbarlabel='Density')
    # tax.scatter(zip(X, Y), s=49, alpha=0.2)
    plt.tight_layout()
    plt.savefig(ofn)

# Figure 2c in the main text
def m2c_article_gini_distribution_violin(
        fn='M2c_article_gini_coef.csv',
        ofn='M2c-article-gini-distribution-violin.pdf',
        nbins=10,
        base=10,
        figsize=FIGSIZE):
    """
    Caption in the main paper:
    --------------------------
    Figure 2: Anomalies. The distribution of types of tweet spreading articles
    from (a) low-credibility ...

    Input data and instructions:
    ----------------------------
    `fn`: the gini coefficient (in terms of how the users share the same
    article) for each low-credibility article.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    bins = np.logspace(
        nplog(df.ntweets.min(), base),
        nplog(df.ntweets.max() + 1, base),
        nbins + 1,
        base=base)
    # pdb.set_trace()
    data = []
    for i in range(nbins):
        data.append(df.loc[(df.ntweets > bins[i])].gini.values)

    positions = bins[:-1] + np.true_divide(bins[1:] - bins[:-1], 2)
    f, ax = plt.subplots(figsize=FIGSIZE)
    ax.set_ylabel('Gini coefficient')
    ax.set_xlabel('Minimum number of tweets')
    ax.violinplot(
        data,
        positions=positions,
        vert=True,
        widths=1.6 * bins[:-1],
        showmeans=False,
        showextrema=False,
        showmedians=True,
        points=100,
        bw_method='scott',
    )

    ax.set_xscale('log')
    ax.set_xlim([1, 50000])
    ax.set_ylim([-0.05, 0.65])
    plt.tight_layout()
    plt.savefig(ofn)

# Figure 2d in the main paper
def m2d_bots_of_users(fn1='M2d_top1000_users.lc.csv',
                      fn2='M2d_rand1000_users.lc.csv',
                      ofn='M2d-bots-of-users.pdf',
                      nbins=20,
                      normed=True,
                      figsize=FIGSIZE):
    """
    Caption in the main paper:
    --------------------------
    Figure 2: Anomalies. The distribution of types of tweet spreading articles
    from (a) low-credibility ...

    Input data and instructions:
    ----------------------------
    `fn1`: the top active (in terms of number of tweets) 1000 accounts with bot
    score that post tweets with a link to low-credibility sources,
    `fn2`: the random sampled 1000 accounts with bot score that post tweets
    with a link to low-credibility sources.
    """
    fn1 = get_data_file(fn1, subfolder='consensus1')
    fn2 = get_data_file(fn2, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    a1 = df1.loc[df1.bot_score.notnull()].bot_score.values
    a2 = df2.loc[df2.bot_score.notnull()].bot_score.values
    mu1 = np.mean(a1)
    sigma1 = np.std(a1, ddof=1)
    mu2 = np.mean(a2)
    sigma2 = np.std(a2, ddof=1)
    logger.info('Number of Non-nan values: len(top)=%s, len(rand)=%s', len(a1),
                len(a2))
    logger.info('Ratio of >0.5: r(top)=%s, r(rand)=%s',
                (a1 > 0.5).sum() / float(len(a1)),
                (a2 > 0.5).sum() / float(len(a2)))
    logger.info('mu1=%s, mu2=%s', mu1, mu2)
    logger.info('sigma1=%s, sigma2=%s', sigma1, sigma2)
    logger.info('Welch\'s t-test: %s', ttest_ind(a1, a2, equal_var=False))
    logger.info('Kolmogorov-Smirnov test: %s', ks_2samp(a1, a2))
    logger.info('Mann Whitney U test: %s', mannwhitneyu(a1, a2))
    fig, ax = plt.subplots(figsize=FIGSIZE)
    bins = np.linspace(0, 1, nbins + 1)
    if normed is False:
        w1 = np.ones_like(a1) / len(a1)
        w2 = np.ones_like(a2) / len(a2)
    else:
        w1 = None
        w2 = None
    ax.set_xlim([0, 1])
    ax.hist(
        a1,
        bins,
        weights=w1,
        normed=normed,
        alpha=1,
        label='Most Active',
        histtype='step',
        color=C1)
    ax.hist(
        a2,
        bins,
        weights=w2,
        normed=normed,
        alpha=1,
        label='Random Sample',
        histtype='step',
        color=C2)
    plt.legend(loc='upper right', fontsize='small')
    ax.set_xlabel('Bot Score')
    if normed is True:
        ax.set_ylabel('Probability Density')
    else:
        ax.set_ylabel('$proportion$')
    plt.tight_layout()
    plt.savefig(ofn)

# Figure 3a in main paper
def m3a_bots_in_first_hour_spreading_violin(
        fn1='M3a_first_day_tweets_for_top1000_articles.lc.csv',
        fn2='M3a_sampled_users_with_score.csv',
        ofn='M3a-bots-in-first-hour-spreading-violin.pdf',
        include_theonion=False,
        keep_first_accounts_only=False,
        base=2,
        figsize=FIGSIZE):
    """
    Caption in the main paper:
    --------------------------
    Figure 3: Bot strategies. (a) Early bot support after a viral
    low-credibility article is first shared. We consider a sample of 60,000
    accounts that participate in the spread of the 1,000 most viral stories
    from low-credibility sources. We align the times when each article
    first appears. We focus on a one-hour early spreading phase following
    each of these events, and divide it into logarithmic lag intervals.
    The plot shows the bot score distribution for accounts sharing the articles
    during each of these lag intervals.

    Input data and instructions:
    ----------------------------
    `fn1`: the first day tweets for the top 1000 low-credibility articles.
    `fn2`: the sampled accounts from `fn1` of size=60,000 and the bot score
    is attached.
    """
    fn1 = get_data_file(fn1, subfolder='consensus1')
    fn2 = get_data_file(fn2, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    df1 = df1.loc[df1.seconds <= 3600]
    if include_theonion is False:
        # site_id=129 for theonion.com
        df1 = df1.loc[df1.site_id != 129]
    df = pd.merge(df1, df2, on='user_raw_id', how='inner')
    df.sort_values('seconds', inplace=True, ascending=True)
    if keep_first_accounts_only is True:
        logger.info('Keep first appeared tweets only!')
        logger.info('Before drop duplicates: %s', len(df))
        df = df.drop_duplicates(['article_id', 'user_raw_id'], keep='first')
        logger.info('After drop duplicates: %s', len(df))
    X = df.seconds.values + 1
    logger.info('Maximum seconds is: %s', X.max())
    if len(df.loc[df.seconds == 1]) > 0:
        logger.info('bin [1,2) needed!')
        nxbins = int(np.ceil(nplog(X.max(), base=base))) + 1
        xbins = np.ones(nxbins + 1)
        xbins[1:] = 1 + base**np.arange(0, nxbins)
    else:
        logger.info('bin [1,2) not needed!')
        nxbins = int(np.ceil(nplog(X.max(), base=base)))
        xbins = base**np.arange(0, nxbins + 1)
    logger.info('xbins are %s', xbins)
    data = []
    # import pdb; pdb.set_trace()
    for i in range(nxbins - 1):
        dfi = df.loc[(X >= xbins[i]) & (X < xbins[i + 1])]
        data.append(dfi.bot_score.values)
    dfi = df.loc[(X >= xbins[nxbins - 1]) & (X <= xbins[nxbins])]
    data.append(dfi.bot_score.values)
    f, ax = plt.subplots(figsize=figsize)
    positions = xbins[:-1] + np.true_divide(xbins[1:] - xbins[:-1], 2)
    logger.info('positions=%s', positions)
    ax.violinplot(
        data,
        positions=positions,
        vert=True,
        widths=0.6 * xbins[:-1],
        showmeans=False,
        showextrema=False,
        showmedians=True,
        points=100,
    )
    ax.set_xlabel('Lag $+1$ (seconds)')
    ax.set_ylabel('Bot Score')
    ax.set_ylim([-0.05, 0.85])
    ax.set_xscale('log')
    ax.set_xlim([1, 4e3])
    plt.tight_layout()
    plt.savefig(ofn)

# help function for m3b_bots_targeting
def assign_segregated_groups(df, ngroups=3):
    g_size, _ = divmod(len(df), ngroups)
    df['gn'] = -1
    windows = []
    left = 0
    df = df.reset_index(drop=True)
    for i in range(ngroups):
        cur_p = (i + 1) * g_size
        right = df.iloc[cur_p - 1].from_bot_score
        windows.append((left, right))
        left = right
    left, right = windows[-1]
    windows[-1] = (left, 1)
    for i, window in enumerate(windows):
        left, right = window
        df.loc[((df.from_bot_score > left) & (df.from_bot_score < right)
                ), 'gn'] = i
    df = df.loc[df.gn != -1]
    return df

def m3b_bots_targeting_3bins(fn='M3b_sampled_mentions.60000.csv',
                             ofn='M3b-bots-targeting-3bins.pdf',
                             nbins=3):
    """
    Caption in the main paper:
    --------------------------
    Figure 3: Bot strategies. (a) Early bot support after a viral
    low-credibility article is first shared. We consider a sample of 60,000
    accounts that participate in the spread of the 1,000 most viral stories
    from low-credibility sources. We align the times when each article
    first appears. We focus on a one-hour early spreading phase following
    each of these events, and divide it into logarithmic lag intervals.
    The plot shows the bot score distribution for accounts sharing the articles
    during each of these lag intervals.

    Input data and instructions:
    ----------------------------
    `fn`: sampled mention tweets that both the mentioning accounts and
    mentioned accounts are in file 'M3a_sampled_users_with_score.csv'.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    df = df.loc[~df.is_via_word_in_tweet_text]
    df = df[['from_raw_id', 'from_bot_score', 'to_followers_count']].copy()
    df = df.groupby('from_raw_id').mean()
    df = df.sort_values('from_bot_score', ascending=True)
    df = df.reset_index(drop=True)
    df = assign_segregated_groups(df, nbins)
    gpb = df.groupby('gn')
    d_mean = gpb.to_followers_count.mean()
    d_std = gpb.to_followers_count.std()
    d_size = gpb.size()
    x = np.array(range(3))
    y1 = d_mean.values / 1e6
    y1err = d_std / d_size.apply(np.sqrt) / 1e6
    fig, ax = plt.subplots(figsize=FIGSIZE)
    ax.errorbar(x, y1, yerr=y1err, fmt='--o', capsize=4)
    ax.set_xticks(x)
    ax.set_xticklabels(['Lower Third', 'Middle Third', 'Top Third'])
    ax.set_xlabel('Bot Score Percentile of Mentioning Accounts')
    ax.set_ylabel('Followers of Mentioned Accounts')
    ax.set_xlim([-0.5, 2.5])
    ax.set_ylim([6.55, 8.15])
    ax.text(
        0,
        1.02,
        r'$\times10^6$',
        transform=ax.transAxes,
        horizontalalignment='left')
    iax = plt.axes([.54, .24, .4, .32])
    df.loc[:, 'to_followers_count'] = np.log10(df.to_followers_count)
    data = []
    for name, gp in df.groupby('gn'):
        data.append(gp['to_followers_count'].values)
    iax.violinplot(
        data,
        positions=x + 0.15,
        vert=True,
        widths=0.8,
        showmeans=False,
        showextrema=False,
        showmedians=True,
        points=100,
    )
    iax.set_xticks(x + 0.15)
    iax.set_xticklabels([])
    iax.set_yticks([0, 4, 8])
    iax.yaxis.set_major_formatter(
        mpl.ticker.FuncFormatter(lambda x, y: r'$10^%d$' % x))
    plt.tight_layout()
    plt.savefig(ofn)

# Figure 4 in the main paper
def m4_bots_retweeter_and_tweeter(fn='M4_retweeter_and_tweeter.csv',
                                  ofn='M4-bots-retweeter-and-tweeter.pdf',
                                  nbins=20,
                                  figsize=(6, 4.5)):
    """
    Caption in the main paper:
    --------------------------
    Figure 4: Impact of bots on humans. (a) Joint distribution of bot scores
    of accounts that retweeted links to low-credibility articles and accounts
    that had originally posted the links. Color represents the number of
    retweeted messages in each bin, on a log scale. (b) The top projection
    shows the distributions of bot scores for retweeters, who are mostly human.
    (c) The left projection shows the distributions of bot scores for accounts
    retweeted by likely humans (score below 0.5), with a significant portion
    of likely bots.

    Input data and instructions:
    ----------------------------
    `fn`: retweeter and tweeter with bot score.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    x = df.retweeter_score.values
    y = df.tweeter_score.values
    H, xedges, yedges = np.histogram2d(
        x, y, bins=np.linspace(0.0, 1.0, nbins + 1))
    # X1, Y1, horizontal curve
    df = pd.DataFrame(H)
    X1 = xedges[:-1] + np.true_divide(xedges[1:] - xedges[:-1], 2)
    Y1 = df.sum(axis=1) / df.sum(axis=1).sum()
    # X2, Y2, vertical curve
    df2 = df.iloc[:int(0.4 * nbins)]
    X21 = df2.sum(axis=0) / df.sum(axis=0).sum()
    df2 = df.iloc[:int(0.5 * nbins)]
    X22 = df2.sum(axis=0) / df.sum(axis=0).sum()
    df2 = df.iloc[:int(0.6 * nbins)]
    X23 = df2.sum(axis=0) / df.sum(axis=0).sum()
    Y2 = X1
    # X3, Y3, heatmap
    X3, Y3 = np.meshgrid(xedges, yedges)

    plt.figure(figsize=figsize)
    # Display -> Figure
    # figtr = fig.transFigure.inverted()
    gs = gridspec.GridSpec(
        2,
        4,
        wspace=0.3,
        hspace=0.2,
        width_ratios=[2.8, 6, 0.2, 0.5],
        height_ratios=[0.8, 2])
    # gs[0, 1].update(hspace=0.1)
    ax1 = plt.subplot(gs[0, 1])
    ax2 = plt.subplot(gs[1, 0])
    ax3 = plt.subplot(gs[1, 1])
    ax4 = plt.subplot(gs[1, 3])
    # horizontal curve
    ax1.plot(X1, Y1)
    # verticla curve
    ax2.plot(X21, Y2, color='k', marker='x', label='$0.4$')
    ax2.plot(X22, Y2, color='#781C6D', marker='*', label='$0.5$')
    ax2.plot(X23, Y2, color='#ED6925', marker='.', label='$0.6$')
    ax2.legend(fontsize='x-small')
    # heatmap
    im = ax3.pcolormesh(
        X3,
        Y3,
        H.T,
        norm=mpl.colors.LogNorm(vmin=H.min() + 1, vmax=H.max()),
        cmap='gnuplot2_r')
    # color bar
    cb = plt.colorbar(im, ax=ax4, orientation='vertical', fraction=0.9)
    cb.set_label('Retweets')
    # disable xticklabels of ax1
    ax1.set_xticklabels([])
    # disable xticklabels of ax2
    ax2.set_yticklabels([])
    # set tick on right yaxis of ax2
    ax2.yaxis.tick_right()
    # add vertical dashed line on ax1
    ax1.axvline(0.5, color='k', linestyle='--', lw=0.5)
    # add xlabel
    ax1.set_ylabel('$\Pr(x)$')
    ax1.set_ylim((0, 0.2))
    # add horizontal dashed line on ax2
    ax2.axhline(0.5, color='k', linestyle='--', lw=0.5)
    # set inverse xlim of ax2, making xaxis increasing on inverse direction
    ax2.set_xlim((0.1, 0))
    # add ylabel
    ax2.set_xlabel('$\Pr(y|x\leq threshold)$')
    # fontsize=rcParams['axes.titlesize'],
    # fontweight=rcParams['axes.titleweight'])
    # set xlabel of ax3
    ax3.set_xlabel('Bot Score of Retweeter, $x$')
    # set ylabel of ax3
    ax3.yaxis.tick_right()
    ax3.yaxis.set_label_position("right")
    ax3.set_ylabel('Bot Score of Tweeter, $y$')
    # disable axis
    ax4.axis('off')
    # add rectangel for ax3
    rect = mpl.patches.Rectangle(
        (0.015, 0.008),
        0.49,
        0.98,
        alpha=1,
        ec='black',
        facecolor=None,
        fill=False,
        ls='-.',
        lw=1.5)
    ax3.add_patch(rect)
    plt.savefig(ofn)

# help function for m5_bots_robustness
def remaining_claim_url(tweet_df, user_centrality_series):
    tweet_df = tweet_df.copy()
    ordered_uid = [
        str(uid)
        for uid in user_centrality_series.sort_values(ascending=False).index
    ]
    user_rank_dict = {uid: rank + 1 for rank, uid in enumerate(ordered_uid)}
    retweet_removal_order = \
        np.minimum(
            tweet_df.from_raw_id.apply(lambda uid: user_rank_dict[uid]),
            tweet_df.to_raw_id.apply(lambda uid: user_rank_dict[uid])
        )
    tweet_df['retweet_removal_order'] = retweet_removal_order
    link_removal_order = tweet_df.groupby('article_id')[
        'retweet_removal_order'].max()
    counts = {user_rank_dict[uid]: 0.0 for uid in ordered_uid}
    counts.update(link_removal_order.value_counts().to_dict())
    removed = pd.Series(counts).sort_index()
    remaining = 1.0 - removed.cumsum() / removed.sum()
    remaining = pd.concat([pd.Series([1.0], index=[0.0]), remaining])
    remaining.index += 1
    return remaining

# help function for m5_bots_robustness
def remaining_percentage(tweet_df, user_centrality_series):
    ordered_uid = [
        str(uid)
        for uid in user_centrality_series.sort_values(ascending=False).index
    ]
    user_rank_dict = {uid: rank + 1 for rank, uid in enumerate(ordered_uid)}
    removal_order = \
        np.minimum(
            tweet_df.from_raw_id.apply(lambda uid: user_rank_dict[uid]),
            tweet_df.to_raw_id.apply(lambda uid: user_rank_dict[uid])
        )
    counts = {user_rank_dict[uid]: 0.0 for uid in ordered_uid}
    counts.update(removal_order.value_counts().to_dict())
    removed = pd.Series(counts).sort_index()
    remaining = 1.0 - removed.cumsum() / removed.sum()
    remaining = pd.concat([pd.Series([1.0], index=[0.0]), remaining])
    remaining.index += 1
    return remaining

def m5_bots_robustness(fn1='M5_retweet.20171108.lc.csv',
                       fn2='M5_centralities.csv',
                       output='M5-bots-robustness.pdf'):
    """
    Caption in the main paper:
    --------------------------
    Figure 5: Dismantling the low-credibility content diffusion network.
    The priority of disconnected nodes is determined by ranking accounts on
    the basis of different characteristics. The remaining fraction of
    (a) unique articles from low-credibility sources and
    (b) retweets linking to those articles is plotted versus the number of
    disconnected nodes.

    Input data and instructions:
    ----------------------------
    `fn1`: retweets with links to low-credibility sources before 2017-11-08.
    `fn2`: three centralities for users with large k-core number: in-strengh,
    out-strengh, and bot-score.
    """
    fn1 = get_data_file(fn1, subfolder='consensus1')
    fn2 = get_data_file(fn2, subfolder='consensus1')
    output = get_output_file(output)

    df1 = pd.read_csv(fn1, dtype=dict(from_raw_id=str, to_raw_id=str))
    df2 = pd.read_csv(fn2, index_col=0)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(4, 6))
    plot_kw = dict(ax=ax1, logx=True)
    remaining_claim_url(df1, df2['in_s']).plot(
        label='Activity', linestyle='-.', **plot_kw)
    remaining_claim_url(df1, df2['out_s']).plot(
        label='Influence', linestyle='--', **plot_kw)
    remaining_claim_url(df1, df2['bot_score']).plot(
        label='Bot Score', linestyle=':', **plot_kw)
    ax1.set_xlabel('number of spreaders disconnected')
    ax1.set_ylabel('fraction of claim links remaining')
    ax1.legend()

    plot_kw = dict(ax=ax2, logx=True)
    remaining_percentage(df1, df2['in_s']).plot(
        label='Activity', linestyle='-.', **plot_kw)
    remaining_percentage(df1, df2['out_s']).plot(
        label='Influence', linestyle='--', **plot_kw)
    remaining_percentage(df1, df2['bot_score']).plot(
        label='Bot Score', linestyle=':', **plot_kw)
    ax2.set_xlabel('number of spreaders disconnected')
    ax2.set_ylabel('fraction of claim links remaining')
    # ax2.legend(loc=3, fontsize=9)
    ax2.set_xlabel('number of spreaders disconnected')
    ax2.set_ylabel('fraction of claim posts remaining')
    plt.tight_layout()
    plt.savefig(output)

# Figure 6 in the main paper
def m6_sources_popularity(fn='M6_top_sources_popularity.csv',
                          ofn='M6-top-sources-popularity.pdf',
                          figsize=(12, 5)):
    """
    Caption in the main paper:
    --------------------------
    Figure 6: Popularity and bot support for the top sources. Satire websites
    are shown in orange, fact-checking sites in blue, and low-credibility
    sources in red. Popularity is measured by total tweet volume (horizontal
    axis) and median number of tweets per article (circle area). Bot support
    is gauged by the median bot score of the 100 most active accounts posting
    links to articles from each source (vertical axis). Low-credibility
    sources have greater support by bots, as well as greater median and/or
    total volume in many cases.

    Input data and instructions:
    ----------------------------
    `fn`: the statistics for the top 20 most active sources (in terms of number
    of tweets), including total number of tweets, median bot score of the 100
    most active accoutns, median number of tweets per article.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    df.rename(
        columns=dict(
            median_of_tweets_per_article='at_medians',
            median_of_bot_score_of_top_users='bs_medians',
            total_tweets='ntweets'),
        inplace=True)
    df = df.sort_values('ntweets', ascending=False)
    # trim domains top label
    df.domain = df.domain.apply(lambda x: x[:-4])
    # size configuration
    base_v = df.at_medians.min()
    scale = np.sqrt(df.at_medians.values / base_v)
    base_s = 180
    sizes = scale * base_s
    # color configuration
    color_low_credibility = 'red'
    color_satire = 'orange'
    color_fact_checking = 'deepskyblue'
    df['color'] = color_low_credibility
    df.loc[df.site_type == 'fact_checking', 'color'] = color_fact_checking
    df.loc[df.site_type == 'satire', 'color'] = color_satire
    # start ploting
    fig, ax = plt.subplots(figsize=figsize)
    ax.scatter(
        df.ntweets.values,
        df.bs_medians.values,
        s=sizes,
        color=df.color.values,
        edgecolors='k',
        alpha=0.8)
    # Set lables for the scattered point
    for txt, x, y, sc in zip(
            range(len(df)), df.ntweets.tolist(), df.bs_medians.tolist(),
            scale):
        ax.annotate(
            txt + 1,
            xy=(x, y),
            # xytext=(0, sc * base_yt),
            # textcoords='offset points',
            # fontsize=nplog(sc + 1, 2) * base_fs,
            ha='center',
            va='center')
    # set legend
    handlers = [
        mpl.lines.Line2D(
            [0], [0],
            linestyle='none',
            mfc='black',
            mec='none',
            marker=r'$\mathregular{{{}.}}$'.format(i),
            markersize=7.5) for i in range(1, 10)
    ]
    handlers += [
        mpl.lines.Line2D(
            [0], [0],
            linestyle='none',
            mfc='black',
            mec='none',
            marker=r'$\mathregular{{{}.}}$'.format(i),
            markersize=12.5) for i in range(10,
                                            len(df) + 1)
    ]
    plt.legend(handlers, df.domain.tolist(), ncol=3, loc=2)
    # set colors of legend text
    # pdb.set_trace()
    for lh, c in zip(ax.get_legend().get_texts(), df.color.tolist()):
        lh.set_color(c)
    ax.set_xlabel('Total Tweet Volume', fontsize=14)
    ax.set_ylabel('Median Bot Score of Active Accounts', fontsize=14)
    ax.tick_params(labelsize=14)
    ax.set_xscale('log')
    ax.set_xlim(1e4, 1e7)
    plt.tight_layout()
    plt.savefig(ofn)

# Extended Data Figure 1 in the supplementary paper
def s1_timeline_stats(fn1='S1_new_articles.weekly.lc.csv',
                      fn2='S1_tweeted_articles.weekly.lc.csv',
                      fn3='S1_new_tweets.weekly.lc.csv',
                      ofn='S1-timeline-stats.pdf',
                      trim_edge=True,
                      trim_broken=False,
                      figsize=(5, 3)):
    """
    Caption in the supplementary paper:
    --------------------------
    Extended Data Figure 1: Weekly tweeted low-credibility articles,
    tweets/article ratio and articles/site ratio. The collection was briefly
    interrupted in October 2016. In December 2016 we expanded the set of
    low-credibility sources, from 70 to 120 websites.

    Input data and instructions:
    ----------------------------
    `fn1`: weekly new coming articles from low-credibility sources,
    `fn2`: weekly tweeted articles from low-credibility sources,
    `fn3`: weekly new coming tweets with links to low-credibility sources.
    """
    fn1 = get_data_file(fn1, subfolder='consensus1')
    fn2 = get_data_file(fn2, subfolder='consensus1')
    fn3 = get_data_file(fn3, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1, parse_dates=['wtimeline'])
    df2 = pd.read_csv(fn2, parse_dates=['wtimeline'])
    df3 = pd.read_csv(fn3, parse_dates=['wtimeline'])
    df1.columns = ['timeline', 'nclaim']
    df2.columns = ['timeline', 'wclaim']
    df3.columns = ['timeline', 'ntweet']
    df = pd.merge(df1, df2, on='timeline', how='inner')
    df = pd.merge(df, df3, on='timeline', how='inner')
    df['site'] = 71
    df.set_index('timeline', inplace=True)
    dt = date(2016, 12, 20)
    dt = dt - timedelta(days=dt.weekday())
    df.loc[dt:, 'site'] = 121
    df['tpa'] = df.ntweet / df.wclaim
    df['aps'] = df.nclaim / df.site
    df = df[['nclaim', 'tpa', 'aps']]
    if trim_edge is True:
        df = df.iloc[1:-1]
    if trim_broken is True:
        df.loc['2016-09-26', :] = np.nan
        df.loc['2016-10-03', :] = np.nan
    colors = ['black', 'red', 'blue']
    # pdb.set_trace()
    ax = df.plot(
        secondary_y=['tpa', 'aps'],
        color=colors,
        figsize=figsize,
        legend=False)
    ax.set_ylabel('Articles')
    ax.set_xlabel('')
    ax.right_ax.set_ylabel('Ratio')
    ax.right_ax.set_ylim([0, 200])
    ln1, _ = ax.get_legend_handles_labels()
    ln2, _ = ax.right_ax.get_legend_handles_labels()
    ax.legend(
        ln1 + ln2, ('Articles', 'Tweets/article ratio', 'Articles/site ratio'),
        fontsize='small')
    mf = mpl.ticker.ScalarFormatter(useMathText=True)
    mf.set_powerlimits((-1, 2))
    ax.yaxis.set_major_formatter(mf)
    plt.tight_layout()
    plt.savefig(ofn)

def s2_article_distribution_by_tweets_and_users(
        fn11='M1a_ntweets_per_article.lc.csv',
        fn12='M1a_ntweets_per_article.ff.csv',
        fn21='S2_nusers_per_article.lc.csv',
        fn22='S2_nusers_per_article.ff.csv',
        ofn='S2-article-distribution-by-tweets-and-users.pdf',
        density=True,
        nbins=15,
        base=2,
        figsize=(4, 5)):
    """
    Caption in the supplementary paper:
    --------------------------
    Extended Data Figure 2: Probability distributions of popularity of articles
    from low-credibility and fact-checking sources, measured by (a) the number
    of tweets and (b) the number of accounts sharing links to an article.

    Input data and instructions:
    ----------------------------
    `fn11`: number of tweets per article for low-credibility sources,
    `fn12`: number of tweets per article for fact-checking sources,
    `fn21`: number of users per article for fact-checking sources,
    `fn22`: number of users per article for fact-checking sources,
    """
    fn11 = get_data_file(fn11, subfolder='consensus1')
    fn12 = get_data_file(fn12, subfolder='consensus1')
    fn21 = get_data_file(fn21, subfolder='consensus1')
    fn22 = get_data_file(fn22, subfolder='consensus1')
    ofn = get_output_file(ofn)
    df11 = pd.read_csv(fn11)
    df12 = pd.read_csv(fn12)
    df21 = pd.read_csv(fn21)
    df22 = pd.read_csv(fn22)
    # Note, set max bin points slightly larger than max data value
    # because the calculation may make right most bin points
    # a little smaller than max data value
    bins11 = np.logspace(
        nplog(df11.counts.min(), base),
        nplog(df11.counts.max() + 1, base),
        nbins + 1,
        base=base)
    bins12 = np.logspace(
        nplog(df12.counts.min(), base),
        nplog(df12.counts.max() + 1, base),
        nbins + 1,
        base=base)
    bins21 = np.logspace(
        nplog(df21.counts.min(), base),
        nplog(df21.counts.max() + 1, base),
        nbins + 1,
        base=base)
    bins22 = np.logspace(
        nplog(df22.counts.min(), base),
        nplog(df22.counts.max() + 1, base),
        nbins + 1,
        base=base)
    if density is not True:
        w11 = np.ones_like(df11.counts.values) / len(df11.counts.values)
        w12 = np.ones_like(df12.counts.values) / len(df12.counts.values)
        w21 = np.ones_like(df21.counts.values) / len(df21.counts.values)
        w22 = np.ones_like(df22.counts.values) / len(df22.counts.values)
    else:
        w11 = None
        w12 = None
        w21 = None
        w22 = None

    h11, bins11 = np.histogram(
        df11.counts.values, bins=bins11, weights=w11, density=density)
    h12, bins12 = np.histogram(
        df12.counts.values, bins=bins12, weights=w12, density=density)
    h21, bins21 = np.histogram(
        df21.counts.values, bins=bins21, weights=w21, density=density)
    h22, bins22 = np.histogram(
        df22.counts.values, bins=bins22, weights=w22, density=density)
    f, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize)
    ax1.set_xlabel('(a) Number of Tweets')
    ax1.set_ylabel('PDF')
    ax1.set_xscale('log')
    ax1.set_yscale('log')
    ax2.set_xlabel('(b) Number of Users')
    ax2.set_ylabel('PDF')
    ax2.set_xscale('log')
    ax2.set_yscale('log')
    ax1.plot(
        bins11[:-1], h11, c=C1, marker='o', alpha=0.6, label='Low-credibility')
    ax1.plot(
        bins12[:-1], h12, c=C2, marker='s', alpha=0.6, label='Fact checking')
    ax2.plot(
        bins21[:-1], h21, c=C1, marker='s', alpha=0.6, label='Low-credibility')
    ax1.legend(fontsize='small')
    ax2.plot(
        bins22[:-1], h22, c=C2, marker='s', alpha=0.6, label='Fact Checking')
    ax2.legend(fontsize='small')

    plt.tight_layout()
    plt.savefig(ofn)

# Extended Data Figure 3 in the supplementary paper

# Extended Data Figure 4 in the supplementary paper

# Extended Data Figure 5 in the supplementary paper

# Extended Data Figure 6 in the supplementary paper
def s6_ccdf_ntweets_per_user_article(
        fn='S6_ntweets_per_user_article.lc.csv',
        ofn='S6-ccdf-ntweets-per-user-article.pdf',
        figsize=FIGSIZE):
    """
    Caption in the supplementary paper:
    -----------------------------------
    Supplementary Figure 6: Cumulative distribution of repetitions, i.e.,
    the number of times a single account tweets the same link to an article
    from a low-credibility source.

    Input data and instructions:
    ----------------------------
    `fn`: number of tweets for each article (from low-credibility sources)
        that is shared by the same user.
    """
    fn = get_data_file(fn, subfolder='consensus1')
    df = pd.read_csv(fn)
    ofn = get_output_file(ofn)
    s = ccdf(df.ncounts)

    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(s.index, s.values)
    ax.set_xlabel('Tweets per Account Sharing the Same Article')
    ax.set_ylabel('CCDF')
    ax.set_xscale('log')
    ax.set_yscale('log')
    # ax.set_xlim([0.8, 200])
    plt.tight_layout()
    plt.savefig(ofn)

# Extended Data Figure 7 in the supplementary paper

# Extended Data Figure 8 in the supplementary paper

# Extended Data Figure 9 in the supplementary paper

# Extended Data Figure 9 in the supplementary paper

# Supplementary Figure 10 in the supplementary paper
def s10_bots_of_users(fn1='S10_top1000_users.lc.csv',
                      fn2='S10_rand1000_users.lc.csv',
                      fn3='sub_user_bot_score.calibrated.csv',
                      ofn='S10-bots-of-users.pdf',
                      nbins=20,
                      normed=True,
                      figsize=FIGSIZE):
    """
    Caption in the supplementary paper:
    --------------------------
    Supplementary Figure 10: Bot score distributions for super-spreaders v.s.
    randomly selected sharers of links to low-credibility sources selected by
    the consensus criterion. The random sample includes 992 accounts who posted
    at least one link to an article from a low-credibility source. Their bot
    scores are compared to 997 accounts that most actively share such links
    The two groups have significant different socres (p<10e-4 according to
    a Mann-Whitney U test).

    Input data and instructions:
    ----------------------------
    `fn1`: the top 1000 most active accounts that share tweets with links to
    low-credibility sources,
    `fn2`: the random sampled 1000 accounts that share tweets with links to
    low-credibility sources,
    `fn3`: bot score for the accounts.
    """
    # pdb.set_trace()
    fn1 = get_data_file(fn1, subfolder='consensus3')
    fn2 = get_data_file(fn2, subfolder='consensus3')
    fn3 = get_data_file(fn3, subfolder='consensus3')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    df3 = pd.read_csv(fn3, usecols=['user_raw_id', 'bot_score'])
    df3 = df3.loc[df3.bot_score.notnull()]
    s1 = df3.loc[df3.user_raw_id.isin(df1.user_raw_id.values)].bot_score
    s2 = df3.loc[df3.user_raw_id.isin(df2.user_raw_id.values)].bot_score
    a1 = s1.values
    a2 = s2.values
    mu1 = np.mean(a1)
    sigma1 = np.std(a1, ddof=1)
    mu2 = np.mean(a2)
    sigma2 = np.std(a2, ddof=1)
    logger.info('Number of Non-nan values: len(top)=%s, len(rand)=%s', len(a1),
                len(a2))
    logger.info('Ratio of >0.5: r(top)=%s, r(rand)=%s',
                (a1 > 0.5).sum() / float(len(a1)),
                (a2 > 0.5).sum() / float(len(a2)))
    logger.info('mu1=%s, mu2=%s', mu1, mu2)
    logger.info('sigma1=%s, sigma2=%s', sigma1, sigma2)
    logger.info('Welch\'s t-test: %s', ttest_ind(a1, a2, equal_var=False))
    logger.info('Kolmogorov-Smirnov test: %s', ks_2samp(a1, a2))
    logger.info('Mann Whitney U test: %s', mannwhitneyu(a1, a2))
    fig, ax = plt.subplots(figsize=FIGSIZE)
    bins = np.linspace(0, 1, nbins + 1)
    if normed is False:
        w1 = np.ones_like(a1) / len(a1)
        w2 = np.ones_like(a2) / len(a2)
    else:
        w1 = None
        w2 = None
    ax.set_xlim([0, 1])
    ax.hist(
        a1,
        bins,
        weights=w1,
        normed=normed,
        alpha=1,
        label='Most Active',
        histtype='step',
        color=C1)
    ax.hist(
        a2,
        bins,
        weights=w2,
        normed=normed,
        alpha=1,
        label='Random Sample',
        histtype='step',
        color=C2)
    plt.legend(loc='upper right', fontsize='small')
    ax.set_xlabel('Bot Score')
    if normed is True:
        ax.set_ylabel('Probability Density')
    else:
        ax.set_ylabel('$proportion$')
    plt.tight_layout()
    plt.savefig(ofn)

# Supplementary Figure 11 in the supplementary paper
def s11_bots_in_first_hour_spreading_violin(
        fn1='S11_first_hour_tweets_for_top1000_articles.lc.csv',
        fn2='S11_sampled_first_hour_users_for_top1000_articles.lc.csv',
        fn3='sub_user_bot_score.calibrated.csv',
        ofn='S10-bots-in-first-hour-spreading-violin.pdf',
        with_theonion=False,
        first=False,
        base=2,
        figsize=FIGSIZE):
    """
    Caption in the supplementary paper:
    --------------------------
    Supplementary Figure 11: Temporal evolution of bot support after the first
    share of a viral story from a consensus low-credibility source. We consider
    a random sample of 20,000 accounts out of the 163,563 accounts that
    participate in the spread of the 1,000 most viral articles. After articles
    from The Onion are excluded, we are left with 42,202 tweets from 13,926
    accounts. We align the times when each link first appears. We focus on a
    one-hour early spreading phase following each of these events, and divide
    it into logarithmic lag intervals. The plot shows the bot score
    distribution for accounts sharing the links during each of thee lag
    intervals.

    Input data and instructions:
    ----------------------------
    `fn1`: first hour tweets for top 1000 most shared low-credibility articles,
    `fn2`: 20,000 accounts sampled from `fn1`
    `fn3`: bot score for the accounts.
    """
    fn1 = get_data_file(fn1, subfolder='consensus3')
    fn2 = get_data_file(fn2, subfolder='consensus3')
    fn3 = get_data_file(fn3, subfolder='consensus3')
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    df3 = pd.read_csv(fn3, usecols=['user_raw_id', 'bot_score'])
    if with_theonion is False:
        df1 = df1.loc[df1.site_id != 129]
    df3 = df3.loc[df3.bot_score.notnull()]
    df3 = df3.loc[df3.user_raw_id.isin(df2.user_raw_id.values)]
    df = pd.merge(df1, df3, on='user_raw_id', how='inner')
    logger.info('%s tweets from %s accounts', df.tweet_id.nunique(),
                df.user_raw_id.nunique())
    df.sort_values('second', inplace=True)
    if first is True:
        logger.info('Keep first appeared tweets only!')
        logger.info('Before drop duplicates: %s', len(df))
        df = df.drop_duplicates(['article_id', 'user_raw_id'], keep='last')
        logger.info('After drop duplicates: %s', len(df))
    X = df.second.values + 1
    logger.info('Maximum seconds is: %s', X.max())
    if len(df.loc[df.second == 1]) > 0:
        logger.info('bin [1,2) needed!')
        nxbins = int(np.ceil(nplog(X.max(), base=base))) + 1
        xbins = np.ones(nxbins + 1)
        xbins[1:] = 1 + base**np.arange(0, nxbins)
    else:
        logger.info('bin [1,2) not needed!')
        nxbins = int(np.ceil(nplog(X.max(), base=base)))
        xbins = base**np.arange(0, nxbins + 1)
    logger.info('xbins are %s', xbins)
    data = []
    for i in range(nxbins - 1):
        dfi = df.loc[(X >= xbins[i]) & (X < xbins[i + 1])]
        data.append(dfi.bot_score.values)
    dfi = df.loc[(X >= xbins[nxbins - 1]) & (X <= xbins[nxbins])]
    data.append(dfi.bot_score.values)
    f, ax = plt.subplots(figsize=figsize)
    positions = xbins[:-1] + np.true_divide(xbins[1:] - xbins[:-1], 2)
    logger.info('positions=%s', positions)
    ax.violinplot(
        data,
        positions=positions,
        vert=True,
        widths=0.6 * xbins[:-1],
        showmeans=False,
        showextrema=False,
        showmedians=True,
        points=100,
    )
    ax.set_xlabel('Lag $+1$ (seconds)')
    ax.set_ylabel('Bot Score')
    ax.set_ylim([-0.05, 1.0])
    ax.set_xscale('log')
    ax.set_xlim([1, 4e3])
    plt.tight_layout()
    plt.savefig(ofn)

# Supplementary Figure 12 in the supplementary paper
def s12_bots_targeting(fn='S12_sampled_mentions.20000.csv',
                       discard_via=True,
                       nbins=3,
                       ofn='S12-bots-targeting-3bins.pdf'):
    """
    Caption in the supplementary paper:
    --------------------------
    Supplementary Figure 12: Average number of followers for Twitter users who
    are mentioned (or replied to) by a sample of 20,000 accounts that link to
    the 1,000 most viral articles from consensus low-credibility sources. We
    obtained bot scores for 4,006 unique mentioning accounts and 4,965 unique
    mentioned accounts, participating in 33,112 mention/reply pairs. We
    excluded 13,817 of these pairs using the "via @screen_name" mentioning
    pattern. The mentioning accounts are ag- gregated into three groups by bot
    score percentile. Error bars indicate standard errors. Inset: Distributions
    of follower counts for users mentioned by accounts in each percentile
    group.

    Input data and instructions:
    ----------------------------
    `fn1`: mentions from these random sampled 20,000 accounts. (see S11)
    """
    fn = get_data_file(fn, subfolder='consensus3')
    ofn = get_output_file(ofn)
    df = pd.read_csv(fn)
    if discard_via is True:
        df = df.loc[~df.is_via_word_in_tweet_text,
                    ['from_raw_id', 'from_bot_score', 'to_followers_count']]
    df = df.groupby('from_raw_id').mean()
    df = df.sort_values('from_bot_score', ascending=True)
    df = df.reset_index(drop=True)
    df = assign_segregated_groups(df, nbins)
    gpb = df.groupby('gn')
    d_mean = gpb.to_followers_count.mean()
    d_std = gpb.to_followers_count.std()
    d_size = gpb.size()
    x = np.array(range(3))
    y1 = d_mean.values / 1e6
    y1err = d_std / d_size.apply(np.sqrt) / 1e6
    fig, ax = plt.subplots(figsize=FIGSIZE)
    ax.errorbar(x, y1, yerr=y1err, fmt='--o', capsize=4)
    ax.set_xticks(x)
    ax.set_xticklabels(['Lower Third', 'Middle Third', 'Top Third'])
    ax.set_xlabel('Bot Score Percentile of Mentioning Accounts')
    ax.set_ylabel('Followers of Mentioned Accounts')
    ax.set_xlim([-0.5, 2.5])
    ax.set_ylim([6.5, 11.8])
    ax.text(
        0,
        1.02,
        r'$\times10^6$',
        transform=ax.transAxes,
        horizontalalignment='left')
    iax = plt.axes([.54, .24, .4, .32])
    df.loc[:, 'to_followers_count'] = np.log10(df.to_followers_count)
    data = []
    for name, gp in df.groupby('gn'):
        data.append(gp['to_followers_count'].values)
    iax.violinplot(
        data,
        positions=x + 0.15,
        vert=True,
        widths=0.8,
        showmeans=False,
        showextrema=False,
        showmedians=True,
        points=100,
    )
    iax.set_xticks(x + 0.15)
    iax.set_xticklabels([])
    iax.set_yticks([0, 4, 8])
    iax.yaxis.set_major_formatter(
        mpl.ticker.FuncFormatter(lambda x, y: r'$10^%d$' % x))
    plt.tight_layout()
    plt.savefig(ofn)

# Supplementary Figure 13 and 14 in the supplementary paper
def s13_14_bots_retweeter_and_tweeter(fn1='S13_retweeter_and_tweeter.csv',
                                      fn2='sub_user_bot_score.calibrated.csv',
                                      ofn='S13-bots-retweeter-and-tweeter.pdf',
                                      exclude_account_theonion=True,
                                      nbins=20,
                                      figsize=(6, 4.5)):
    """
    Caption in the supplementary paper:
    -----------------------------------
    Supplementary Figure 13: Joint distribution of the bot scores of accounts
    that retweeted links to articles from consensus low-credibility sources
    and accounts that had originally posted the links. We considered retweets
    by a sample of 20,000 accounts that posted the 1,000 most viral articles.
    We obtained bot scores for 12,792 tweeting accounts and 17,664 retweeting
    accounts, participating in 229,725 retweet pairs. Color represents the
    number of retweeted messages in each bin, on a log scale. Projections show
    the distributions of bot scores for retweeters (top) and for accounts
    retweeted by likely humans (left).

    Input data and instructions:
    ----------------------------
    `fn1`: retweeter and tweeter from sampled 20, 000 accounts,
    `fn2`: bot score for retweeter and tweeter.
    `exclude_account_theonion`: the account `theonion` contributes so many
    retweets that it overloads others. Set this flag to exclude this account.
    """
    fn1 = get_data_file(fn1, subfolder='consensus3')
    fn2 = get_data_file(fn2, subfolder='consensus3')
    if exclude_account_theonion is True:
        ofn = 'S14-bots-retweeter-and-tweeter.pdf'
    else:
        ofn = 'S13-bots-retweeter-and-tweeter.pdf'
    ofn = get_output_file(ofn)
    df1 = pd.read_csv(fn1)
    df2 = pd.read_csv(fn2)
    df2 = df2.loc[df2.bot_score.notnull()]
    if exclude_account_theonion is True:
        # exlude theonion accounts
        df1 = df1.loc[df1.tweeter_id != 18643437]
    df = pd.merge(
        df2, df1, left_on='user_raw_id', right_on='retweeter_id', how='inner')
    df.rename(columns=dict(bot_score='retweeter_score'), inplace=True)
    df = pd.merge(
        df2, df, left_on='user_raw_id', right_on='tweeter_id', how='inner')
    df.rename(columns=dict(bot_score='tweeter_score'), inplace=True)
    df = df[['retweeter_score', 'tweeter_score']].copy()
    x = df.retweeter_score.values
    y = df.tweeter_score.values
    H, xedges, yedges = np.histogram2d(
        x, y, bins=np.linspace(0.0, 1.0, nbins + 1))
    # X1, Y1, horizontal curve
    df = pd.DataFrame(H)
    X1 = xedges[:-1] + np.true_divide(xedges[1:] - xedges[:-1], 2)
    Y1 = df.sum(axis=1) / df.sum(axis=1).sum()
    # X2, Y2, vertical curve
    df2 = df.iloc[:10]
    X2 = df2.sum(axis=0) / df.sum(axis=0).sum()
    Y2 = X1
    # X3, Y3, heatmap
    X3, Y3 = np.meshgrid(xedges, yedges)

    plt.figure(figsize=figsize)
    # Display -> Figure
    # figtr = fig.transFigure.inverted()
    gs = gridspec.GridSpec(
        2,
        4,
        wspace=0.3,
        hspace=0.2,
        width_ratios=[2.5, 6, 0.2, 0.5],
        height_ratios=[0.8, 2])
    # gs[0, 1].update(hspace=0.1)
    ax1 = plt.subplot(gs[0, 1])
    ax2 = plt.subplot(gs[1, 0])
    ax3 = plt.subplot(gs[1, 1])
    ax4 = plt.subplot(gs[1, 3])
    # horizontal curve
    ax1.plot(X1, Y1)
    # verticla curve
    ax2.plot(X2, Y2)
    # heatmap
    im = ax3.pcolormesh(
        X3,
        Y3,
        H.T,
        norm=mpl.colors.LogNorm(vmin=H.min() + 1, vmax=H.max()),
        cmap='gnuplot2_r')
    # color bar
    cb = plt.colorbar(im, ax=ax4, orientation='vertical', fraction=0.9)
    cb.set_label('Retweets')
    # disable xticklabels of ax1
    ax1.set_xticklabels([])
    # disable xticklabels of ax2
    ax2.set_yticklabels([])
    # set tick on right yaxis of ax2
    ax2.yaxis.tick_right()
    # add vertical dashed line on ax1
    ax1.axvline(0.5, color='k', linestyle='--', lw=0.5)
    # add xlabel
    ax1.set_ylabel('$\Pr(x)$')
    ax1.set_ylim((0, 0.2))
    # add horizontal dashed line on ax2
    ax2.axhline(0.5, color='k', linestyle='--', lw=0.5)
    # set inverse xlim of ax2, making xaxis increasing on inverse direction
    if exclude_account_theonion is True:
        ax2.set_xlim((0.2, 0))
    else:
        ax2.set_xlim((0.4, 0))

    # add ylabel
    ax2.set_xlabel('$\Pr(y|x\leq 0.5)$')
    # fontsize=rcParams['axes.titlesize'],
    # fontweight=rcParams['axes.titleweight'])
    # set xlabel of ax3
    ax3.set_xlabel('Bot Score of Retweeter, $x$')
    # set ylabel of ax3
    ax3.yaxis.tick_right()
    ax3.yaxis.set_label_position("right")
    ax3.set_ylabel('Bot Score of Tweeter, $y$')
    # disable axis
    ax4.axis('off')
    # add rectangel for ax3
    rect = mpl.patches.Rectangle(
        (0.015, 0.008),
        0.49,
        0.98,
        alpha=1,
        ec='black',
        facecolor=None,
        fill=False,
        ls='-.',
        lw=1.5)
    ax3.add_patch(rect)
    plt.savefig(ofn)

# Supplementary Figure 15 in the supplementary paper


import base64
import json
import urllib.parse
import urllib.request

import cv2
import numpy as np
from PIL import Image

from Thumbnail import LowerThumbnail
from Thumbnail import UpperThumbnail

std_len = 63

class ImageInput:
    def __init__(self, image_path):
        self.image_path = image_path
        self.image_to_upload = Image.open(image_path)
        self.width = None
        self.height = None
        self.image_to_upload_base64 = None
        self.score_map = None
        self.gray_scale = None
        self.thumbnail = None
        self.width_thumbnail = None
        self.height_thumbnail = None

    def generateThumbnail(self):
        self.width = self.image_to_upload.size[0]
        self.height = self.image_to_upload.size[1]
        scale_width = self.width / 590
        scale_height = self.height / 260
        if scale_width > scale_height:
            scale_mode = 1
        else:
            scale_mode = 2

        if scale_mode == 1:
            self.width_thumbnail = 590
            self.height_thumbnail = int(self.height * 590 / self.width)
        else:
            self.height_thumbnail = 260
            self.width_thumbnail = int(self.width * 260 / self.height)
        self.thumbnail = UpperThumbnail(self.image_to_upload, self.width_thumbnail, self.height_thumbnail).thumbnail
        return self.thumbnail

    def encodeImage(self):
        image_to_upload_binary = open(self.image_path, "rb")
        self.image_to_upload_base64 = base64.b64encode(image_to_upload_binary.read())

    def request(self):
        params = {"image": self.image_to_upload_base64}
        params = urllib.parse.urlencode(params).encode(encoding="UTF8")

        request_url = "https://aip.baidubce.com/rest/2.0/image-classify/v1/body_seg"
        access_token = ""
        request_url = request_url + "?access_token=" + access_token

        request = urllib.request.Request(url=request_url, data=params)
        request.add_header("Content-Type", "application/x-www-form-urlencoded")
        response = urllib.request.urlopen(request)
        content = response.read()
        if content:
            print(content)

        content = json.loads(content)  # content为通过接口获取的返回json
        self.score_map = base64.b64decode(content["scoremap"])
        print(self.score_map)

    def getGrayscale(self):
        nparr = np.frombuffer(self.score_map, np.uint8)
        labelimg = cv2.imdecode(nparr, 1)
        labelimg = cv2.resize(labelimg, (self.width, self.height), interpolation=cv2.INTER_NEAREST)
        im_new = np.where(labelimg == 1, 255, labelimg)
        cv2.imwrite("score_map.jpg", im_new)
        self.gray_scale = Image.open("score_map.jpg")

class ImageOutput:
    ascii_array = list("$@B%8&WM#ZO0QLCJUYXvuxrjft/|()1{}[]?*-_+~<>i!lI;:,^`.")
    ascii_array_length = len(ascii_array)
    unit = (256 + 1) / ascii_array_length

    def __init__(self, image, gray_scale, width, height, width_thumbnail, height_thumbnail):
        self.image = image
        self.gray_scale = gray_scale
        self.width = width
        self.height = height
        self.thumbnail = None
        self.width_thumbnail = width_thumbnail
        self.height_thumbnail = height_thumbnail

    def mapGrayScale(self):
        if self.width < self.height:
            width_out = std_len
            height_out = int(std_len * self.height / self.width)
        else:
            height_out = std_len
            width_out = int(std_len * self.width / self.height)

        scale_2 = max(self.width / width_out, self.height / height_out)
        self.width = int(self.width / scale_2 * 2)
        self.height = int(self.height / scale_2)
        self.gray_scale = self.gray_scale.resize((self.width, self.height), Image.NEAREST)
        self.image = self.image.resize((self.width, self.height), Image.NEAREST)
        self.writeOutputFile()

    def map_char(self, r, g, b, alpha=0):
        if alpha == 256:
            return '.'
        else:
            pos = int((256 - r) / self.unit)
            return self.ascii_array[pos]

    def writeOutputFile(self):
        file = open("TEXT_IMAGE.txt", 'w')
        for i in range(self.height):
            for j in range(self.width):
                char_0 = self.map_char(*self.gray_scale.getpixel((j, i)))
                if char_0 == '.':
                    file.write('.')
                else:
                    char_1 = self.map_char(*self.image.getpixel((j, i)))
                    file.write(char_1)
            file.write('\n')
        file.close()

    def generateThumbnail(self):
        self.thumbnail = LowerThumbnail(self.width_thumbnail, self.height_thumbnail).thumbnail
        return self.thumbnail


import cv2
import random
import numpy as np
import os

class shotgenerator():
    def __init__(self, msbpath, fps, num_video = -1, threshold_frame = 80):
        '''
        This class generates and stores individual shots of a video given the following
        inputs :
        INPUT : msbpath, fps, num_video = -1, thereshold_frame
        msbpath : Directory path of msb files containing shot boundary data
        fps : Set the Frames per Second of the generated shot
        threshold_frame - Set the threshold value for shot generation
        '''
        self.here = os.path.dirname(os.path.abspath(__file__))
        self.msbpath = msbpath
        self.check_dir()
        self.fps = fps
        self.threshold_frame = threshold_frame

    def check_dir(self):
        ########################################################
        #Create result directory where video shots will be stored
        ########################################################
        shotspath = os.path.join(self.here, "shots")
        if not os.path.isdir(shotspath):
            os.mkdir(shotspath)

    def get_boundaries(self, timestamps):
        '''
        Returns timestamps of video in list format given the msb timestamp file
        INPUT : timestamps
        timestamps - msb timestamp file path of the concerened video.
        OUPUT : times
        times - a list containg timestamps of the video in list format.
        '''
        timestamps = os.path.splitext(timestamps)[0] + '.msb'
        with open(timestamps, encoding = "ISO-8859-1") as f:
            times = f.readlines()[2:]
        f.close()
        return times

    def add_noise(self, frame):
        '''
        Adds random noise to an image.
        INPUT - frame
        frame:input frame where noise is to be added
        OUTPUT - noisy
        noisy:frame with added noise
        '''
        mean = 0
        var = np.random.uniform(0, 1.0)
        sigma = var**0.5
        row,col= frame.shape
        gauss = np.random.normal(mean,sigma,(row,col))
        gauss = gauss.reshape(row,col)
        noisy = frame + gauss
        noisy = np.uint8(noisy)
        return noisy

    def shotmaker(self, start, end, videopath, outpath, noise):
        '''
        Generates and stores individual shots in a video given start and end frame
        of the shots, and the videopath of the parent video. Shots are named with
        the starting frame of the shot.
        '''
        video = cv2.VideoCapture(videopath)
        framecount = start
        height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        outpath = os.path.join(outpath, str(start) + '.avi')
        out = cv2.VideoWriter(outpath,cv2.VideoWriter_fourcc(*'XVID'), self.fps, (width, height), 0)
        while framecount <= end:
            video.set(1, framecount-1)
            res, frame = video.read()
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
            if noise:
                frame = self.add_noise(frame)
            out.write(frame)
            framecount = framecount + 1
        out.release()

    def run(self, vidname, outpath, noise):
        '''
        Check if the length of individual shots are greater than a threshold value.
        If yes, runs the shotmaker function to generate and store shots.
        INPUT - vidname, outpath
        vidname : name of the video accoriding to TRECVID IACC.3 dataset xml file
        outpath : outpath path of the shots
        '''
        timestamps = os.path.join(self.msbpath, vidname)
        videopath = os.path.join(self.here, "videos", vidname)

        times = self.get_boundaries(timestamps)
        count = 1
        print("Generating Shots of Video-{} ...".format(vidname.split('._-o-_.')[1].rsplit('.', 1)[0]))
        for boundary in times:
            start = int(boundary.split(" ")[0])
            end = int(boundary.split(" ")[1][:-1])
            if (end - start) > self.threshold_frame:
                self.shotmaker(start, end, videopath, outpath, noise)
                count+=1

    def iterator(self, num_video=-1, noise=False):
        '''
        Iterates videos in 'videos' folder created by the viddownlaod.py module.
        Call run function to create shots from all videos or specified number of videos.
        INPUT - num_video
        num_video : Number of videos in the videos folder to be considered to create shots. Defauult value = -1, all videos.
        '''
        downvideos_path = os.path.join(self.here, "videos")
        video_names = []
        for file in os.listdir("/home/tre3x/Python/FilmEditsDetection/data/synthetic_data/videos"):
            if file.endswith(".mp4"):
                video_names.append(file)

        if num_video != -1:
            video_names = random.sample(video_names, num_video)

        for video in video_names:
            outpath = os.path.join(self.here, "shots", video.split('._-o-_.')[1].rsplit('.', 1)[0])
            if not os.path.exists(outpath):
                os.makedirs(outpath)
                self.run(video, outpath, noise)

if __name__ == '__main__':
    shotgenerator("/home/tre3x/Python/FilmEditsDetection/data/synthetic_data/msb", 24).iterator(2, True)


import numpy as np
import os
import glob
import cv2
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import config
import time

sys.path.append("..")

from object_detection.utils import label_map_util

# set to one GPU machine
os.environ["CUDA_VISIBLE_DEVICES"]="1"

def detect(input_dir,
           output_dir,
	   model_name,
           use_sample = config.DetectConfig.use_sample,
           use_relative_size = config.DetectConfig.use_relative_size):
    detection_metadata_fn = os.path.join(output_dir, config.config.DETS_FILE)
    print("Starting detection {}, output {}, scale {}". format(input_dir, detection_metadata_fn, config.DetectConfig.SCALE_RATIO))

    # Path to frozen detection graph. This is the actual model that is used for the object detection.
    PATH_Fold = "C:\\Users\\olgali\\repo\\models\\research\\fine_tuned_model_set1_10k\\"
    PATH_TO_CKPT = PATH_Fold + 'frozen_inference_graph.pb'

    # List of the strings that is used to add correct label for each box.
    PATH_TO_LABELS = os.path.join('C:\\data\\woodknots\\board_images_set1_output', 'pascal_label_map.pbtxt')

    NUM_CLASSES = 2

    ## Load a (frozen) Tensorflow model into memory
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')

    ## Loading label map
    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
    category_index = label_map_util.create_category_index(categories)

    ## Helper code
    def load_image_into_numpy_array(image):
      (im_width, im_height) = image.size
      return np.array(image.getdata()).reshape(
          (im_height, im_width, 3)).astype(np.uint8)

    #============================ Detection ============================
    start_time = time.time()
    raw_img_name = next(os.walk(input_dir))[2][0]
    fn_img, ext = os.path.splitext(os.path.basename(raw_img_name))
    raw_img = cv2.imread(os.path.join(input_dir, raw_img_name))
    raw_im_height, raw_im_width = raw_img.shape[0:2]

    print("Reading frames...")
    if config.DetectConfig.SCALE_RATIO < 1:
        im_height, im_width = int(raw_im_height*config.DetectConfig.SCALE_RATIO), int(raw_im_width*config.DetectConfig.SCALE_RATIO)
        images, image_names = zip(*[(cv2.resize(cv2.imread(file), (im_height, im_width)), file.split('/')[-1]) for file in glob.glob(input_dir + '/*' + ext)])
    else:
        im_height, im_width = raw_im_height, raw_im_width
        images, image_names = zip(*[(cv2.imread(file), file.split('/')[-1]) for file in glob.glob(input_dir + '/*' + ext)])

    images = np.array(images)
    image_names = list(image_names)
    if (use_sample == 1):
        print("Using sample size: ", config.DetectConfig.SAMPLE_SIZE)
        images = images[:config.DetectConfig.SAMPLE_SIZE]
        image_names = image_names[:config.DetectConfig.SAMPLE_SIZE]

    boxes_total = np.zeros(shape = (len(images), 300, 4))
    scores_total =  np.zeros(shape = (len(images), 300))
    classes_total = np.zeros(shape = (len(images), 300))

    with detection_graph.as_default():
        with tf.Session(graph=detection_graph) as sess:
            # Definite input and output Tensors for detection_graph
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            # Each box represents a part of the image where a particular object was detected.
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            # Each score represent how level of confidence for each of the objects.
            # Score is shown on the result image, together with the class label.
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            for i in range(0, len(images), config.DetectConfig.BATCH_SIZE):
                feed_images = np.array(images[i:i+config.DetectConfig.BATCH_SIZE])
                (boxes, scores, classes, num) = sess.run(
                  [detection_boxes, detection_scores, detection_classes, num_detections],
                  feed_dict={image_tensor: feed_images})
                boxes_total[i:i+config.DetectConfig.BATCH_SIZE] = boxes
                scores_total[i:i+config.DetectConfig.BATCH_SIZE] = scores
                classes_total[i:i+config.DetectConfig.BATCH_SIZE] = classes

    print("--- %s seconds ---" % (time.time() - start_time))
    #============================ Post-processing ============================
    all_dets = []
    skipped_dets = 0
    for i in range(len(images)):
        image_name = os.path.basename(image_names[i])
        per_score = scores_total[i]
        per_class = classes_total[i]
        per_box = boxes_total[i]
        ind = np.where(per_score >= config.DetectConfig.MIN_SCORE_THRESH)
        #print("Found {} boxes for class {}".format(len(ind[0]), config.config.CLASSES_DICT))

        saved_detection = 0
        for i in ind[0]:
            bbox_height = abs(per_box[i][0] - per_box[i][2])*im_height
            if bbox_height < config.DetectConfig.MIN_DETECT_HEIGHT:
                skipped_dets = skipped_dets + 1
                continue
            if use_relative_size:
                bbox = per_box[i]
            else:
                bbox = [ per_box[i][1]*im_width, per_box[i][0]*im_height, per_box[i][3]*im_width,per_box[i][2]*im_height]

            all_dets.append(image_name+','+str(config.config.CLASSES_DICT[per_class[i]]) + ',' + str(per_score[i]) + ',' +
                            ','.join(map(str,bbox)) + ','
                            + str(im_width) + ',' + str(im_height))
            saved_detection = 1
        # Handle images that did not have any detections
        if not saved_detection:
            bbox = [0, 0, 0, 0]
            all_dets.append(
                image_name + ',' + config.DetectConfig.CLASS_NA + ',' + str(0) + ',' +
                ','.join(map(str, bbox)) + ','
                + str(im_width) + ',' + str(im_height))

    ## Save to log file
    out = open(detection_metadata_fn, 'w')
    line = 'image,class,score,bbox_0,bbox_1,bbox_2,bbox_3,im_width,im_height'
    out.write("%s\n" % line)

    for line in all_dets:
        out.write("%s\n" % line)
    out.flush()

    out.close()

    print("Skipped {0} detections which heights is less than {1}". format(skipped_dets, config.DetectConfig.MIN_DETECT_HEIGHT))


# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the 'License'). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the 'license' file accompanying this file. This file is
# distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import logging
import os
import signal

import numpy as np
import xgboost as xgb

from sklearn.model_selection import RepeatedKFold, RepeatedStratifiedKFold
from sagemaker_algorithm_toolkit import exceptions as exc
from sagemaker_algorithm_toolkit.channel_validation import Channel
from sagemaker_xgboost_container.data_utils import get_content_type, get_dmatrix, get_size, validate_data_file_path
from sagemaker_xgboost_container import distributed
from sagemaker_xgboost_container import checkpointing
from sagemaker_xgboost_container.algorithm_mode import channel_validation as cv
from sagemaker_xgboost_container.algorithm_mode import hyperparameter_validation as hpv
from sagemaker_xgboost_container.algorithm_mode import metrics as metrics_mod
from sagemaker_xgboost_container.algorithm_mode import train_utils
from sagemaker_xgboost_container.callback import add_debugging
from sagemaker_xgboost_container.constants.xgb_constants import CUSTOMER_ERRORS, XGB_MAXIMIZE_METRICS

MODEL_NAME = "xgboost-model"

logger = logging.getLogger(__name__)

def add_sigterm_handler(model_dir, is_master):
    """Stop training and cleanup model directory when SIGTERM is received.

    Model directory is only cleaned if is_master is True. Otherwise program terminates.

    :param model_dir: Directory where model is saved
    :param is_master: True if single node training, or the current node is the master node in distributed training
    """
    def _terminate():
        os._exit(0)

    def _cleanup_files(signo, frame):
        if is_master:
            train_utils.cleanup_dir(model_dir, MODEL_NAME)

        _terminate()

    signal.signal(signal.SIGTERM, _cleanup_files)

def get_validated_dmatrices(train_path, validate_path, content_type, csv_weights=0, is_pipe=False,
                            combine_train_val=False):
    """Get training and validation Data Matrices for XGBoost training.

    Check size and format of both training and validation data channels, and return parsed
    Data Matrices.

    :param train_path:
    :param validate_path:
    :param content_type: Content type of data. Supports 'libsvm' or 'csv'
    :param csv_weights: 1 if instance weights are in the second column of csv data files; otherwise, 0
    :param is_pipe: Boolean to indicate if data is being read in pipe mode
    :combine_train_val: Boolean to indicate if returns a DMatrix combining train and validation data
    :return: Parsed xgb.DMatrix
    """
    train_files_size = get_size(train_path, is_pipe) if train_path else 0
    val_files_size = get_size(validate_path, is_pipe) if validate_path else 0

    if not is_pipe:
        logging.debug("File size need to be processed in the node: {}mb.".format(
            round((train_files_size + val_files_size) / (1024 * 1024), 2)))

        if train_files_size > 0:
            validate_data_file_path(train_path, content_type)
        if val_files_size > 0:
            validate_data_file_path(validate_path, content_type)

    train_dmatrix = get_dmatrix(train_path, content_type, csv_weights=csv_weights, is_pipe=is_pipe) \
        if train_files_size > 0 else None
    val_dmatrix = get_dmatrix(validate_path, content_type, csv_weights=csv_weights, is_pipe=is_pipe) \
        if val_files_size > 0 else None

    train_val_dmatrix = train_dmatrix
    if combine_train_val and train_dmatrix is not None and val_dmatrix is not None:
        logging.info("Read both train and validation data into one DMatrix")
        train_val_dmatrix = get_dmatrix([train_path, validate_path], content_type,
                                        csv_weights=csv_weights, is_pipe=is_pipe)

    return train_dmatrix, val_dmatrix, train_val_dmatrix

def sagemaker_train(train_config, data_config, train_path, val_path, model_dir, sm_hosts, sm_current_host,
                    checkpoint_config):
    """Train XGBoost in a SageMaker training environment.

    Validate hyperparameters and data channel using SageMaker Algorithm Toolkit to fail fast if needed.
    If running with more than one host, check if the current host has data and run train_job() using
    rabit_run.

    :param train_config:
    :param data_config:
    :param train_path:
    :param val_path:
    :param model_dir:
    :param sm_hosts:
    :param sm_current_host:
    :param checkpoint_config:
    """
    metrics = metrics_mod.initialize()

    hyperparameters = hpv.initialize(metrics)
    validated_train_config = hyperparameters.validate(train_config)
    if validated_train_config.get("updater"):
        validated_train_config["updater"] = ",".join(validated_train_config["updater"])

    channels = cv.initialize()
    validated_data_config = channels.validate(data_config)

    logging.debug("hyperparameters {}".format(validated_train_config))
    logging.debug("channels {}".format(validated_data_config))

    # Get Training and Validation Data Matrices
    file_type = get_content_type(validated_data_config['train'].get("ContentType"))
    input_mode = validated_data_config['train'].get("TrainingInputMode")
    csv_weights = validated_train_config.get("csv_weights", 0)
    is_pipe = (input_mode == Channel.PIPE_MODE)

    validation_channel = validated_data_config.get('validation', None)
    combine_train_val = '_kfold' in validated_train_config
    train_dmatrix, val_dmatrix, train_val_dmatrix = get_validated_dmatrices(train_path, val_path, file_type,
                                                                            csv_weights, is_pipe, combine_train_val)

    checkpoint_dir = checkpoint_config.get("LocalPath", None)

    train_args = dict(
        train_cfg=validated_train_config,
        train_dmatrix=train_dmatrix,
        val_dmatrix=val_dmatrix,
        train_val_dmatrix=train_val_dmatrix,
        model_dir=model_dir,
        checkpoint_dir=checkpoint_dir)

    # Obtain information about training resources to determine whether to set up Rabit or not
    num_hosts = len(sm_hosts)

    if num_hosts > 1:
        # Wait for hosts to find each other
        logging.info("Distributed node training with {} hosts: {}".format(num_hosts, sm_hosts))
        distributed.wait_hostname_resolution(sm_hosts)

        if not train_dmatrix:
            logging.warning("Host {} does not have data. Will broadcast to cluster and will not be used in distributed"
                            " training.".format(sm_current_host))
        distributed.rabit_run(exec_fun=train_job, args=train_args, include_in_training=(train_dmatrix is not None),
                              hosts=sm_hosts, current_host=sm_current_host, update_rabit_args=True)
    elif num_hosts == 1:
        if train_dmatrix:
            if validation_channel:
                if not val_dmatrix:
                    raise exc.UserError("No data in validation channel path {}".format(val_path))
            logging.info("Single node training.")
            train_args.update({'is_master': True})
            train_job(**train_args)
        else:
            raise exc.UserError("No data in training channel path {}".format(train_path))
    else:
        raise exc.PlatformError("Number of hosts should be an int greater than or equal to 1")

def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_dir, checkpoint_dir, is_master):
    """Train and save XGBoost model using data on current node.

    If doing distributed training, XGBoost will use rabit to sync the trained model between each boosting iteration.
    Trained model is only saved if 'is_master' is True.

    :param train_cfg: Training hyperparameter configurations
    :param train_dmatrix: Training Data Matrix
    :param val_dmatrix: Validation Data Matrix
    :param train_val_dmatrix: Training + Validation Data Matrix
    :param model_dir: Directory where model will be saved
    :param is_master: True if single node training, or the current node is the master node in distributed training.
    """
    # Parse arguments for train() API
    num_round = train_cfg.pop("num_round")
    # Parse arguments for intermediate model callback
    save_model_on_termination = train_cfg.pop('save_model_on_termination', "false")

    # Evaluation metrics to use with train() API
    tuning_objective_metric_param = train_cfg.pop("_tuning_objective_metric", None)
    eval_metric = train_cfg.get("eval_metric")
    cleaned_eval_metric, configured_feval, tuning_objective_metric = train_utils.get_eval_metrics_and_feval(
        tuning_objective_metric_param, eval_metric)
    if cleaned_eval_metric:
        train_cfg['eval_metric'] = cleaned_eval_metric
    else:
        train_cfg.pop('eval_metric', None)

    early_stopping_rounds = train_cfg.pop('early_stopping_rounds', None)
    early_stopping_data_name = 'validation' if val_dmatrix else None
    early_stopping_metric = None
    if early_stopping_rounds:
        if tuning_objective_metric:
            early_stopping_metric = tuning_objective_metric[-1]
        elif eval_metric:
            early_stopping_metric = eval_metric[-1]

    logging.info("Train matrix has {} rows and {} columns".format(train_dmatrix.num_row(), train_dmatrix.num_col()))
    if val_dmatrix:
        logging.info("Validation matrix has {} rows".format(val_dmatrix.num_row()))

    try:
        kfold = train_cfg.pop("_kfold", None)

        if kfold is None:
            xgb_model, iteration, callbacks, watchlist = get_callbacks_watchlist(
                train_dmatrix=train_dmatrix, val_dmatrix=val_dmatrix, model_dir=model_dir,
                checkpoint_dir=checkpoint_dir, early_stopping_data_name=early_stopping_data_name,
                early_stopping_metric=early_stopping_metric, early_stopping_rounds=early_stopping_rounds,
                save_model_on_termination=save_model_on_termination, is_master=is_master)
            add_debugging(callbacks=callbacks, hyperparameters=train_cfg, train_dmatrix=train_dmatrix,
                          val_dmatrix=val_dmatrix)

            bst = xgb.train(train_cfg, train_dmatrix, num_boost_round=num_round-iteration, evals=watchlist,
                            feval=configured_feval, callbacks=callbacks, xgb_model=xgb_model, verbose_eval=False)

        else:
            num_cv_round = train_cfg.pop("_num_cv_round", 1)
            logging.info("Run {}-round of {}-fold cross validation with {} rows".format(num_cv_round,
                                                                                kfold,
                                                                                train_val_dmatrix.num_row()))

            bst = []
            evals_results = []

            num_class = train_cfg.get("num_class", None)
            objective = train_cfg.get("objective", None)
            # RepeatedStratifiedKFold expects X as array-like of shape (n_samples, n_features)
            X = range(train_val_dmatrix.num_row())
            y = train_val_dmatrix.get_label() if num_class or objective.startswith("binary:") else None
            rkf = RepeatedStratifiedKFold(n_splits=kfold, n_repeats=num_cv_round) if y is not None \
                else RepeatedKFold(n_splits=kfold, n_repeats=num_cv_round)

            for train_index, val_index in rkf.split(X=X, y=y):
                cv_train_dmatrix = train_val_dmatrix.slice(train_index)
                cv_val_dmatrix = train_val_dmatrix.slice(val_index)

                xgb_model, iteration, callbacks, watchlist = get_callbacks_watchlist(
                    train_dmatrix=cv_train_dmatrix, val_dmatrix=cv_val_dmatrix, model_dir=model_dir,
                    checkpoint_dir=checkpoint_dir, early_stopping_data_name=early_stopping_data_name,
                    early_stopping_metric=early_stopping_metric, early_stopping_rounds=early_stopping_rounds,
                    save_model_on_termination=save_model_on_termination, is_master=is_master, fold=len(bst))
                add_debugging(callbacks=callbacks, hyperparameters=train_cfg, train_dmatrix=cv_train_dmatrix,
                              val_dmatrix=cv_val_dmatrix)

                evals_result = {}
                logging.info("Train cross validation fold {}".format((len(bst) % kfold) + 1))
                booster = xgb.train(train_cfg, cv_train_dmatrix, num_boost_round=num_round-iteration,
                                    evals=watchlist, feval=configured_feval, evals_result=evals_result,
                                    callbacks=callbacks, xgb_model=xgb_model, verbose_eval=False)
                bst.append(booster)
                evals_results.append(evals_result)

                if len(bst) % kfold == 0:
                    logging.info("The metrics of round {} cross validation".format(int(len(bst) / kfold)))
                    print_cv_metric(num_round, evals_results[-kfold:])

            if num_cv_round > 1:
                logging.info("The overall metrics of {}-round cross validation".format(num_cv_round))
                print_cv_metric(num_round, evals_results)
    except Exception as e:
        for customer_error_message in CUSTOMER_ERRORS:
            if customer_error_message in str(e):
                raise exc.UserError(str(e))

        exception_prefix = "XGB train call failed with exception"
        raise exc.AlgorithmError("{}:\n {}".format(exception_prefix, str(e)))

    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    if is_master:
        if type(bst) is not list:
            model_location = os.path.join(model_dir, MODEL_NAME)
            bst.save_model(model_location)
            logging.debug("Stored trained model at {}".format(model_location))
        else:
            for fold in range(len(bst)):
                model_location = os.path.join(model_dir, f"{MODEL_NAME}-{fold}")
                bst[fold].save_model(model_location)
                logging.debug("Stored trained model {} at {}".format(fold, model_location))

def get_callbacks_watchlist(train_dmatrix, val_dmatrix, model_dir, checkpoint_dir,
                            early_stopping_data_name, early_stopping_metric, early_stopping_rounds,
                            save_model_on_termination, is_master, fold=None):
    if checkpoint_dir and fold is not None:
        checkpoint_dir = os.path.join(checkpoint_dir, f"model-{fold}")

    # Set callbacks
    xgb_model, iteration = checkpointing.load_checkpoint(checkpoint_dir)
    if xgb_model is not None:
        if fold is not None:
            xgb_model = f"{xgb_model}-{fold}"
        logging.info("Checkpoint loaded from %s", xgb_model)
        logging.info("Resuming from iteration %s", iteration)

    callbacks = []
    callbacks.append(xgb.callback.EvaluationMonitor())
    if checkpoint_dir:
        save_checkpoint = xgb.callback.TrainingCheckPoint(directory=checkpoint_dir,
                                                          iterations=iteration,
                                                          name=checkpointing.CHECKPOINT_FILENAME)
        callbacks.append(save_checkpoint)

    if save_model_on_termination == "true":
        model_name = f"{MODEL_NAME}-{fold}" if fold is not None else MODEL_NAME
        save_intermediate_model = checkpointing.SaveIntermediateModelCallBack(model_dir, model_name, is_master)
        callbacks.append(save_intermediate_model)
        add_sigterm_handler(model_dir, is_master)

    if early_stopping_data_name and early_stopping_metric and early_stopping_rounds:
        maximize = early_stopping_metric in XGB_MAXIMIZE_METRICS
        early_stop = xgb.callback.EarlyStopping(rounds=early_stopping_rounds,
                                                data_name=early_stopping_data_name,
                                                metric_name=early_stopping_metric,
                                                maximize=maximize,
                                                save_best=True)
        callbacks.append(early_stop)

    watchlist = [(train_dmatrix, 'train')]
    if val_dmatrix is not None:
        watchlist.append((val_dmatrix, 'validation'))

    return xgb_model, iteration, callbacks, watchlist

def print_cv_metric(num_round, evals_results):
    cv_eval_report = f"[{num_round}]"
    for metric_name in evals_results[0]['train']:
        for data_name in ["train", "validation"]:
            metric_val = [evals_result[data_name][metric_name][-1] for evals_result in evals_results]
            cv_eval_report += '\t{0}-{1}:{2:.5f}'.format(data_name, metric_name, np.mean(metric_val))
    print(cv_eval_report)