import numpy as np
import pandas as pd
import pandas.api.types as pdtypes
from pygtftk.stats.intersect.modl import tree
from functools import partial
import math
from collections import Counter

from plotnine import ggplot, aes, labs, scale_color_gradient, geom_point, geom_abline, scale_x_log10, scale_y_log10, geom_violin, geom_boxplot, position_dodge, scale_y_log10, xlab, ylab, ggtitle

ROOT_PATH = "./output/ologram_result_scatacseq_pbmc/"
DATA_ROOT_PATH = "./output/sc_atac_seq_pbmc_data/"

# For manual execution
# ROOT_PATH = "../output/ologram_result_scatacseq_pbmc/"
# DATA_ROOT_PATH = "../output/sc_atac_seq_pbmc_data/"


## Superclusters: how to regroup the cell types
# For now, this is hardcoded
SUPERCLUSTERS = {
    "CD14+ Monocytes":  "cd14",
    "CD4 Naive":        "cd48",
    "CD8 Naive":        "cd48",
    "Quer":             "NA",
    "...":              "NA",
    "NA":               "NA",
    "pre-B cell":       "preB"
}


# ---------------------------------------------------------------------------- #

## Read the OLOGRAM-MODL result file
# Read dataframe to create the found_combis dictionary
df_res = pd.read_csv(ROOT_PATH+"merged_batches_result.tsv", sep='\t', header=0, index_col=None)
# Pval set to 0 or -1 are changed to 1e-320 and NaN respectively
df_res.loc[df_res['summed_bp_overlaps_pvalue'] == 0, 'summed_bp_overlaps_pvalue'] = 1e-320
df_res.loc[df_res['summed_bp_overlaps_pvalue'] == -1, 'summed_bp_overlaps_pvalue'] = np.nan

# Create a Library tree with the combinations, like in the MODL algorithm itself  
print("Creating a Library for these words. This can be very long for longer words.")
print("In this particular Snakemake, it should take about 12 minutes.")

L = tree.Library()
L.build_nodes_for_words_from_ologram_result_df(df_res)
L.assign_nodes()


# ---------------------------------------------------------------------------- #


# Remember the features names are in L.features_names
# And the interesting attributes of each node are: 
#   node.word, node.children, node.s pval and fc
def word_to_combi(word): return [L.features_names[i]  for i in range(len(word)) if word[i]]

# Retrieves all combis and enrichment
def work_on_this_node(node, graph):
    return word_to_combi(node.word), node.fc, node.pval, node.s

# Iterate over all nodes
mygraphfunc = partial(work_on_this_node, graph=L)
global_results = {}
tree.apply_recursively_to_all_nodes(L.root_node, mygraphfunc, global_results)

df = pd.DataFrame(global_results.values(), columns = ["combination","fc","pval","s"])

df['combi_length'] = (df['combination'].apply(len)) -2
# Substract 2 to the length to discard the "Query" and "..." terms of the combinations

# Now, get the enrichments of all 2-wise combis, all 3-wise, etc.
def entropy(data):
    """
    Compute Shannon entropy.
    """
    if len(data) <= 1: return 0

    # Data counts
    counts = Counter()
    for d in data: counts[d] += 1

    ent = 0
    probs = [float(c) / len(data) for c in counts.values()]
    for p in probs:
        if p > 0.: ent -= p * math.log(p, 2)

    return ent


# ---------------------------------------------------------------------------- #


# Read the translation table as a dictionary and convert
# If df_map is a pandas dataframe with two columns, 'name' and 'category'

df_map = pd.read_csv(DATA_ROOT_PATH + "cell_to_class.tsv", sep='\t', header=0, index_col=None)
df_map.set_index('id', inplace=True)
df_map.transpose()
dict_translate = df_map['class'].to_dict()

dict_translate["Query"] = "NA" # Add the 'Query' class
dict_translate["..."] = "NA"


def combination_to_combi_of_classes(combination):
    new_combination = []
    for element in combination:
        current_key = dict_translate[element]
        element_to_add = SUPERCLUSTERS[current_key]
        if element_to_add != "NA": 
            new_combination += [element_to_add]
    return new_combination

df['combination_classes'] = df['combination'].apply(combination_to_combi_of_classes)


## Compute entropy
df['entropy'] = df['combination_classes'].apply(entropy)
# Round to 2 decimals and make a category
df['entropy'] = df['entropy'].apply(lambda x: round(x,2)).astype(pdtypes.CategoricalDtype())


# Using only combis of each individual length

all_lengths = sorted(set(df['combi_length']))

# No point in going beyond, roughly, 12
all_lengths = [l for l in all_lengths if l <= 12]

for length in all_lengths:

    df_filtered = df.loc[df['combi_length'] == length,:]

    try:
        p = (ggplot(data=df_filtered, mapping = aes(x='entropy', y='fc'))
                + geom_violin(position = position_dodge(1), width = 1)
                + geom_boxplot(position = position_dodge(1), width = 0.25)
                + xlab("Entropy") + ylab("Fold change (log2)")
                + ggtitle("Order "+str(length)))

        p.save(filename = ROOT_PATH + "entropy_graph/entropy_length_" + str(length) + "_fc.png")



        p = (ggplot(data=df_filtered, mapping = aes(x='entropy', y='s'))
                + geom_violin(position = position_dodge(1), width = 1)
                + geom_boxplot(position = position_dodge(1), width = 0.25)
                + xlab("Entropy") + ylab("True total overlapping bp.")
                + ggtitle("Order "+str(length)))

        p.save(filename = ROOT_PATH + "entropy_graph/entropy_length_" + str(length) + "_s.png")


    except:
        print("Skipping length",length,"which caused an error.")

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.offsetbox import AnchoredText

# root_dir = './output/supp_fig4/'
# concat_file = root_dir + 'merged_ologram_test.tsv'
# res_file_1 = root_dir + 'S_mean_boxplot.pdf'
# res_file_2 = root_dir + 'S_var_boxplot.pdf'

concat_file = snakemake.input["merged"]         # Concatenated shuffles
truth_file = snakemake.input["truth"]           # Deep sampled empirical p-value
res_file = snakemake.output[0]                  # Result file


# ---------------------------- Read result files ----------------------------- #
def result_ologram_to_list_of_lists(file, is_truth = False):

    # Skip even rows except for 0, they are the headers of other files (non-truth data only)
    def logic(index):
        if (index % 2 == 0) and index != 0 : return True
        return False 

    if not is_truth: df = pd.read_csv(file, sep='\t', header = 0, index_col = None, skiprows = lambda x: logic(x))
    else: df = pd.read_csv(file, sep='\t', header = 0, index_col = None)

    # Iterate over each row and build a result object
    result = list()
    for i, row in df.iterrows():
        header = row['feature_type'].split('_')

        if not is_truth:
            r = {
                'nb_minibatches': int(header[2]),
                'try_id': int(header[5]),
                'p_val': row['summed_bp_overlaps_pvalue'],
                'log_10(p_val)': np.log10(row['summed_bp_overlaps_pvalue'] + 1E-320),
                'log_10(empirical_pval)': np.log10(row["summed_bp_overlaps_empirical_pvalue"] + 1E-320)
            }
        else:
            r = {
                'p_val': row['summed_bp_overlaps_pvalue'],
                'log_10(p_val)': np.log10(row['summed_bp_overlaps_pvalue'] + 1E-320),
                'log_10(empirical_pval)': np.log10(row["summed_bp_overlaps_empirical_pvalue"] + 1E-320)
                #'log_10(beta_pval)': np.log10(row["ad_hoc_beta_summed_bp_overlaps_pvalue"] + 1E-320)
            }

        result += [r]

    return result


result = result_ologram_to_list_of_lists(concat_file)
result_truth = result_ologram_to_list_of_lists(truth_file, is_truth = True)

# ------------------------------- Plot figure -------------------------------- #
result_df = pd.DataFrame(result)

MINIBATCH_SIZE = 10 # Hardcoded for now, keep it the same in the Snakefile
result_df['Number of shuffles'] = MINIBATCH_SIZE * result_df['nb_minibatches']

meanplot = result_df.boxplot(column=['log_10(p_val)'], by='Number of shuffles')
result_truth_df = pd.DataFrame(result_truth)

# Add a line for the true p-val (as estimated by very deep sampling)
true_pval = list(result_truth_df['log_10(empirical_pval)'])[0] # Hardcoded to read the first line, since there should be only one
plt.axhline(y=true_pval, color='g', linestyle='-')

#beta_pval = list(result_truth_df['log_10(beta_pval)'])[0] # Hardcoded to read the first line, since there should be only one
#plt.axhline(y=beta_pval, color='r', linestyle='-')

at = AnchoredText("True p-val = "+'{0:.4g}'.format(10**true_pval),
                  prop=dict(size=8), frameon=True,
                  loc='upper right',
                  )
at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
meanplot.add_artist(at)

plt.savefig(res_file)

import numpy as np
np.random.seed(42)

import pandas as pd
from plotnine import ggplot, aes, geom_point, geom_line, geom_smooth, scale_x_continuous, scale_y_log10, xlab, ylab

import time

from pygtftk.stats.intersect.modl.dict_learning import Modl, test_data_for_modl
from pygtftk.stats.intersect.modl.apriori import Apriori, matrix_to_list_of_transactions, apriori_results_to_matrix
from pygtftk.stats.intersect.modl.subroutines import learn_dictionary_and_encode
from pygtftk import utils


# Debug : plots are not shown, only printed, so use Agg backend for matplotlib
import matplotlib
matplotlib.use("Agg")

OUTPUT_ROOT = "output/benchmark/scaling/" # Hardcoded for now. It was necessary to launch this script in shell and not via snakemake.script to allow proper redirection of the log


## ---------------------------- Parameters ---------------------------------- ##
utils.VERBOSITY = 0 # We don't want to record DEBUG messages for these tests
REPEATS = range(5) # Repeat all operations N times to get the average

N_THREADS = 4

## MODL
SIZES = [6,9,12,18,24,30,40,50]  # Numbers of sets (columns)
STEPS = [3,5,8,10,12,15,18,20,22,25,30] # Numbers of queried atoms

# When varying only one parameter, which is the default for the others ?
DEFAULT_QUERIED_ATOMS = 8
DEFAULT_N_SETS = 8

# Data parameters
NFLAGS = 10000
NOISE = 0.1     # Previously 0.5


## ---------------------------- Time benchmarks ----------------------------- ##

## Number of sets
df_bench = pd.DataFrame(columns = ['nb_sets','time'])   # Prepare dataframe

for _ in REPEATS:

    for size in SIZES:
        X = test_data_for_modl(nflags = NFLAGS, number_of_sets = size, noise = NOISE)

        start_time = time.time()
        combi_miner = Modl(X, multiple_overlap_max_number_of_combinations = DEFAULT_QUERIED_ATOMS, nb_threads = N_THREADS)
        modl_interesting_combis = combi_miner.find_interesting_combinations()
        stop_time = time.time()

        df_bench = df_bench.append({'nb_sets':size, 'time': stop_time-start_time}, ignore_index = True)

df_bench['nb_sets'] = df_bench['nb_sets'].astype(int)
p = (ggplot(df_bench) + aes('nb_sets', 'time')
 + geom_point() + geom_smooth(span=.3) + scale_x_continuous()
 + xlab("Number of sets") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig1")


## Number of queried words
df_bench = pd.DataFrame(columns = ['step','time'])

X = test_data_for_modl(nflags = NFLAGS, number_of_sets = DEFAULT_N_SETS, noise = NOISE)

for _ in REPEATS:
    for step in STEPS:

        start_time = time.time()
        combi_miner = Modl(X, multiple_overlap_max_number_of_combinations = step, nb_threads = N_THREADS)
        modl_interesting_combis = combi_miner.find_interesting_combinations()
        stop_time = time.time()

        df_bench = df_bench.append({'step':step, 'time': stop_time-start_time}, ignore_index = True)

df_bench['step'] = df_bench['step'].astype(int)
p = (ggplot(df_bench) + aes('step', 'time')
 + geom_point() + geom_smooth(span=.3) + scale_x_continuous()
 + xlab("Queried nb. of atoms") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig2")

import numpy as np
np.random.seed(42)

import pandas as pd
from plotnine import ggplot, aes, geom_point, geom_line, geom_smooth, scale_x_continuous, scale_y_log10, xlab, ylab

import time

from pygtftk.stats.intersect.modl.dict_learning import Modl, test_data_for_modl
from pygtftk.stats.intersect.modl.apriori import Apriori, matrix_to_list_of_transactions, apriori_results_to_matrix
from pygtftk.stats.intersect.modl.subroutines import learn_dictionary_and_encode
from pygtftk import utils

from mlxtend.frequent_patterns import fpgrowth, apriori

# Debug : plots are not shown, only printed, so use Agg backend for matplotlib
import matplotlib
matplotlib.use("Agg")

OUTPUT_ROOT = "output/benchmark/scaling/" # Hardcoded for now. It was necessary to launch this script in shell and not via snakemake.script to allow proper redirection of the log


## ---------------------------- Parameters ---------------------------------- ##
utils.VERBOSITY = 0 # We don't want to record debug messages for these tests
REPEATS = range(5) # Repeat all operations N times to get the average

## Elementary operation (DL) vs other itemset miners

# Number of words, and 1/min_support for the comparisons
SCALING_FACTORS =  [1,2,5,10,25,30,40,50,75,100,150,200,250,300,350,400,500]    

# Data generation parameters
NOISE = 0.5
NLINES = 5000
NSETS = 13

# DL parameters
ALPHA = 0.5

## -------------- Elementary operation benchmark : apriori vs fpgrowth vs dict learning

## Support and number of queried words
df_bench = pd.DataFrame(columns = ['scaling_factor','algo','time'])   # Prepare df

for _ in REPEATS:

    for k in SCALING_FACTORS:

        # NOTE Scaling factor of k means:
        # - Apriori and FP-growth will have min support of 1/k
        # - Dictionary learning will learn k atoms
        current_min_support = 1/k

        X = test_data_for_modl(nflags = NLINES, number_of_sets = NSETS, noise = NOISE)
        names = [str(i) for i in range(X.shape[1])]
        transactions = matrix_to_list_of_transactions(X, names)
        X_as_dataframe = pd.DataFrame(X)

        # Apriori - my pure Python implementation
        start_time = time.time()
        myminer = Apriori(min_support = current_min_support)
        myminer.run_apriori(transactions)
        results = myminer.produce_results()
        stop_time = time.time()

        df_bench = df_bench.append({'scaling_factor':k, 'algo':'apriori_pure_python', 'time': stop_time-start_time}, ignore_index = True)  


        # Apriori 
        # It seems this particular apriori implementation reserves a very large amount of memory when support is very low, so cap it.
        if current_min_support >= 0.02: 
            start_time = time.time()
            apriori(X_as_dataframe, min_support = current_min_support)
            stop_time = time.time()

            df_bench = df_bench.append({'scaling_factor':k, 'algo':'apriori', 'time': stop_time-start_time}, ignore_index = True)  

        # FP-Growth
        start_time = time.time()
        result = fpgrowth(X_as_dataframe, min_support=current_min_support)
        stop_time = time.time()

        df_bench = df_bench.append({'scaling_factor':k, 'algo':'fpgrowth', 'time': stop_time-start_time}, ignore_index = True)  

        # Dict learning
        start_time = time.time()
        U_df, V_df, error = learn_dictionary_and_encode(X, n_atoms = k, alpha = ALPHA, n_jobs = 1)
        stop_time = time.time()

        df_bench = df_bench.append({'scaling_factor':k, 'algo':'DL', 'time': stop_time - start_time}, ignore_index = True)   


df_bench['scaling_factor'] = df_bench['scaling_factor'].astype(int)
p = (ggplot(df_bench) + aes('scaling_factor', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Scaling factor (k)") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig3")

p = (ggplot(df_bench) + aes('scaling_factor', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous() + scale_y_log10()
 + xlab("Scaling factor (k)") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig3_log10")


# Normalized time to scaling factor of 1
minimum_time = df_bench[df_bench['scaling_factor']==1][['algo','time']]
df_bench['time_relative'] = df_bench['time'] # Placeholder
for index, row in df_bench.iterrows():
    my_algo = row['algo']
    my_minimum_time = pd.to_numeric(minimum_time.loc[minimum_time['algo'] == my_algo]['time']).min()
    df_bench.at[index,'time_relative'] = row['time']/my_minimum_time

p = (ggplot(df_bench) + aes('scaling_factor', 'time_relative', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Scaling factor (k)") + ylab("Time (relative)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig3_relative")


p = (ggplot(df_bench) + aes('scaling_factor', 'time_relative', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous() + scale_y_log10()
 + xlab("Scaling factor (k)") + ylab("Time (relative)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig3_relative_log10")

import numpy as np
np.random.seed(42)

import pandas as pd
from plotnine import ggplot, aes, geom_point, geom_line, geom_smooth, scale_x_continuous, scale_y_log10, xlab, ylab

import time

from pygtftk.stats.intersect.modl.dict_learning import Modl, test_data_for_modl
from pygtftk.stats.intersect.modl.apriori import Apriori, matrix_to_list_of_transactions, apriori_results_to_matrix
from pygtftk.stats.intersect.modl.subroutines import learn_dictionary_and_encode
from pygtftk import utils


from mlxtend.frequent_patterns import fpgrowth, apriori

# Debug : plots are not shown, only printed, so use Agg backend for matplotlib
import matplotlib
matplotlib.use("Agg")

OUTPUT_ROOT = "output/benchmark/scaling/" # Hardcoded for now. It was necessary to launch this script in shell and not via snakemake.script to allow proper redirection of the log



## ---------------------------- Parameters ---------------------------------- ##
utils.VERBOSITY = 0 # We don't want to record debug messages for these tests
REPEATS = range(5) # Repeat all operations N times to get the average

## Elementary operation (DL) vs other itemset miners

# Number of sets (columns), minimum of 6
SETS_NB = [6,7,8,9,10,11,12,13,14,16,18,20,22,24,26]    

# Data generation parameters
NOISE = 0.5
LOW_NOISE = 0.01
N_FLAGS = 5000

# DL parameters
ALPHA = 0
N_ATOMS = 120

# Other algorithms' parameters
MIN_SUPPORT = 1/100 

## -------------- Elementary operation benchmark : apriori vs fpgrowth vs dict learning



## Number of sets
df_bench = pd.DataFrame(columns = ['set_nb','algo','time'])   # Prepare df

for _ in REPEATS:

    for size in SETS_NB:

        # Generate data
        X = test_data_for_modl(nflags = N_FLAGS, number_of_sets = size, noise = NOISE)
        names = [str(i) for i in range(X.shape[1])]
        transactions = matrix_to_list_of_transactions(X, names)
        X_as_dataframe = pd.DataFrame(X)

        X_low_noise = test_data_for_modl(nflags = N_FLAGS, number_of_sets = size, noise = LOW_NOISE)
        X_noiseless = test_data_for_modl(nflags = N_FLAGS, number_of_sets = size, noise = 0)

        # Apriori
        # Cap it to the max number that is not unreasonable
        if size <= 15:
            start_time = time.time()
            myminer = Apriori(min_support = MIN_SUPPORT)
            myminer.run_apriori(transactions)
            results = myminer.produce_results()
            stop_time = time.time()

            df_bench = df_bench.append({'set_nb':size, 'algo':'apriori_pure_python', 'time': stop_time-start_time}, ignore_index = True)   


        # Apriori
        # NOTE: this implementation has high RAM cost for large data or low min_support, be careful
        if size <= 15:
            start_time = time.time()
            result = apriori(X_as_dataframe, min_support = MIN_SUPPORT)
            stop_time = time.time()
            df_bench = df_bench.append({'set_nb':size, 'algo':'apriori', 'time': stop_time-start_time}, ignore_index = True)  


        # FP-Growth
        start_time = time.time()
        result = fpgrowth(X_as_dataframe, min_support = MIN_SUPPORT)
        stop_time = time.time()

        df_bench = df_bench.append({'set_nb':size, 'algo':'fpgrowth', 'time': stop_time-start_time}, ignore_index = True)  


        # Dict learning
        start_time = time.time()
        U_df, V_df, error = learn_dictionary_and_encode(X, n_atoms = N_ATOMS, alpha = ALPHA, n_jobs = 1)
        stop_time = time.time()

        df_bench = df_bench.append({'set_nb':size, 'algo':'DL', 'time': stop_time - start_time}, ignore_index = True)   


        # Dict learning - low noise data
        start_time = time.time()
        U_df, V_df, error = learn_dictionary_and_encode(X_low_noise, n_atoms = N_ATOMS, alpha = ALPHA, n_jobs = 1)
        stop_time = time.time()

        df_bench = df_bench.append({'set_nb':size, 'algo':'DL_low_noise_data', 'time': stop_time - start_time}, ignore_index = True)  


        # Dict learning - noiseless data
        start_time = time.time()
        U_df, V_df, error = learn_dictionary_and_encode(X_noiseless, n_atoms = N_ATOMS, alpha = ALPHA, n_jobs = 1)
        stop_time = time.time()

        df_bench = df_bench.append({'set_nb':size, 'algo':'DL_noiseless_data', 'time': stop_time - start_time}, ignore_index = True)  



df_bench['set_nb'] = df_bench['set_nb'].astype(int)
p = (ggplot(df_bench) + aes('set_nb', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Number of sets") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig4")

p = (ggplot(df_bench) + aes('set_nb', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous() + scale_y_log10()
 + xlab("Number of sets") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig4_log10")


# Normalized time to minimum number of sets
min_nb_sets = min(SETS_NB)
minimum_time = df_bench[df_bench['set_nb'] == min_nb_sets][['algo','time']]
df_bench['time_relative'] = df_bench['time'] # Placeholder
for index, row in df_bench.iterrows():
    my_algo = row['algo']
    my_minimum_time = pd.to_numeric(minimum_time.loc[minimum_time['algo'] == my_algo]['time']).min()
    df_bench.at[index,'time_relative'] = row['time']/my_minimum_time

p = (ggplot(df_bench) + aes('set_nb', 'time_relative', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Number of sets") + ylab("Time (relative)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig4_relative")


p = (ggplot(df_bench) + aes('set_nb', 'time_relative', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous() + scale_y_log10()
 + xlab("Number of sets") + ylab("Time (relative)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig4_relative_log10")

import numpy as np
np.random.seed(42)

import pandas as pd
from plotnine import ggplot, aes, geom_point, geom_line, geom_smooth, scale_x_continuous, scale_y_log10, xlab, ylab

import time

from pygtftk.stats.intersect.modl.dict_learning import Modl, test_data_for_modl
from pygtftk.stats.intersect.modl.apriori import Apriori, matrix_to_list_of_transactions, apriori_results_to_matrix
from pygtftk.stats.intersect.modl.subroutines import learn_dictionary_and_encode
from pygtftk import utils

from mlxtend.frequent_patterns import fpgrowth, apriori

# Debug : plots are not shown, only printed, so use Agg backend for matplotlib
import matplotlib
matplotlib.use("Agg")

OUTPUT_ROOT = "output/benchmark/scaling/" # Hardcoded for now. It was necessary to launch this script in shell and not via snakemake.script to allow proper redirection of the log

## ---------------------------- Parameters ---------------------------------- ##
utils.VERBOSITY = 0 # We don't want to record debug messages for these tests
REPEATS = range(5) # Repeat all operations N times to get the average


## Elementary operation (DL) vs other itemset miners

# Number of lines (in thousands)
LINES_NB = [1,2,5,8,10,15,20,25,40,50]
NB_SETS = 15  

# Data generation parameters
NOISE = 0.5

# DL parameters
ALPHA = 0
N_ATOMS = 120

# Other algo parameters
MIN_SUPPORT = 1/100


## -------------- Elementary operation benchmark : apriori vs fpgrowth vs dict learning

## Number of lines
df_bench = pd.DataFrame(columns = ['lines','algo','time'])   # Prepare df

for _ in REPEATS:

    for lines in LINES_NB:

        true_line_nb = lines * 1000

        # Generate data
        # Scaling factor is in the thouands
        X = test_data_for_modl(nflags = true_line_nb, number_of_sets = NB_SETS, noise = NOISE)
        names = [str(i) for i in range(X.shape[1])]
        transactions = matrix_to_list_of_transactions(X, names)
        X_as_dataframe = pd.DataFrame(X)

        # Apriori
        start_time = time.time()
        myminer = Apriori(min_support = MIN_SUPPORT)
        myminer.run_apriori(transactions)
        results = myminer.produce_results()
        stop_time = time.time()

        df_bench = df_bench.append({'lines':lines, 'algo':'apriori_pure_python', 'time': stop_time-start_time}, ignore_index = True)   


        # Apriori
        # NOTE: this implementation has high RAM cost for large data or low min_support, be careful with scaling
        if true_line_nb <= 20000:
            start_time = time.time()
            result = apriori(X_as_dataframe, min_support = MIN_SUPPORT)
            stop_time = time.time()
            df_bench = df_bench.append({'lines':lines, 'algo':'apriori', 'time': stop_time-start_time}, ignore_index = True)  


        # FP-Growth
        start_time = time.time()
        result = fpgrowth(X_as_dataframe, min_support = MIN_SUPPORT)
        stop_time = time.time()

        df_bench = df_bench.append({'lines':lines, 'algo':'fpgrowth', 'time': stop_time-start_time}, ignore_index = True)  


        # Dict learning
        start_time = time.time()
        U_df, V_df, error = learn_dictionary_and_encode(X, n_atoms = N_ATOMS, alpha = ALPHA, n_jobs = 1)
        stop_time = time.time()

        df_bench = df_bench.append({'lines':lines, 'algo':'DL', 'time': stop_time - start_time}, ignore_index = True)   


df_bench['lines'] = df_bench['lines'].astype(int)
p = (ggplot(df_bench) + aes('lines', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Number of lines") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig5")

p = (ggplot(df_bench) + aes('lines', 'time', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous() + scale_y_log10()
 + xlab("Number of lines") + ylab("Time (seconds)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig5_log10")


# Normalized time to minimum line number
min_nb_lines = min(LINES_NB)
minimum_time = df_bench[df_bench['lines'] == min_nb_lines][['algo','time']]
df_bench['time_relative'] = df_bench['time'] # Placeholder
for index, row in df_bench.iterrows():
    my_algo = row['algo']
    my_minimum_time = pd.to_numeric(minimum_time.loc[minimum_time['algo'] == my_algo]['time']).min()
    df_bench.at[index,'time_relative'] = row['time']/my_minimum_time

p = (ggplot(df_bench) + aes('lines', 'time_relative', color='algo', group='algo')
 + geom_point() + geom_smooth(method='gpr', span=.3) + scale_x_continuous()
 + xlab("Number of lines") + ylab("Time (relative)"))
p.save(filename = OUTPUT_ROOT + "scaling_fig5_relative")