#
#     dualloop
#
#        File:  dualloop.py
#
#     Authors: Deleted for purposes of anonymity
#
#     Proprietor: Deleted for purposes of anonymity --- PROPRIETARY INFORMATION
#
# The software and its source code contain valuable trade secrets and shall be maintained in
# confidence and treated as confidential information. The software may only be used for
# evaluation and/or testing purposes, unless otherwise explicitly stated in the terms of a
# license agreement or nondisclosure agreement with the proprietor of the software.
# Any unauthorized publication, transfer to third parties, or duplication of the object or
# source code---either totally or in part---is strictly prohibited.
#
#     Copyright (c) 2019 Proprietor: Deleted for purposes of anonymity
#     All Rights Reserved.
#
# THE PROPRIETOR DISCLAIMS ALL WARRANTIES, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS FOR A PARTICULAR PURPOSE AND THE WARRANTY AGAINST LATENT
# DEFECTS, WITH RESPECT TO THE PROGRAM AND ANY ACCOMPANYING DOCUMENTATION.
#
# NO LIABILITY FOR CONSEQUENTIAL DAMAGES:
# IN NO EVENT SHALL THE PROPRIETOR OR ANY OF ITS SUBSIDIARIES BE
# LIABLE FOR ANY DAMAGES WHATSOEVER (INCLUDING, WITHOUT LIMITATION, DAMAGES
# FOR LOSS OF BUSINESS PROFITS, BUSINESS INTERRUPTION, LOSS OF INFORMATION, OR
# OTHER PECUNIARY LOSS AND INDIRECT, CONSEQUENTIAL, INCIDENTAL,
# ECONOMIC OR PUNITIVE DAMAGES) ARISING OUT OF THE USE OF OR INABILITY
# TO USE THIS PROGRAM, EVEN IF the proprietor HAS BEEN ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGES.
#
# For purposes of anonymity, the identity of the proprietor is not given herewith.
# The identity of the proprietor will be given once the review of the
# conference submission is completed.
#
# THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY.
#

import numpy as np

np.seterr(divide='ignore', invalid='ignore')
import matplotlib.colors as clr
import matplotlib.pyplot as plt
import os
import pandas as pd
import pickle
import seaborn as sns
import sys
import time
import torch
import math
from scipy.stats import entropy
from tqdm.notebook import tqdm

from snorkel.analysis import Scorer
from sklearn.metrics import confusion_matrix

from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
import xgboost as xgb
from sklearn.neural_network import MLPClassifier

from weasul_label_model import WeaSulLabelModel as WeaSulLabelModel
from snorkel.labeling.model.label_model import LabelModel
from snorkel.labeling.model import MajorityLabelVoter
from scipy.optimize import minimize
from sklearn.metrics import f1_score

import ray

### for test

def show_lfs_stat2(my_df, my_lfs, balance, epochs):
    results = []

    y = my_df['label'].to_numpy()

    # add the result of each lf
    for lf in my_lfs:
        result = {}

        result['name'] = lf

        Y_pred = my_df[lf].to_numpy()

        coverage = (len(Y_pred) - list(Y_pred).count(-1)) / len(Y_pred)
        result['coverage'] = coverage

        tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
        result['tn'] = tn
        result['fp'] = fp
        result['fn'] = fn
        result['tp'] = tp

        result['precision'] = tp / (tp + fp)
        result['recall'] = tp / (tp + fn)
        result['f1'] = tp / (tp + 0.5 * (fp + fn))
        result['accuracy'] = (tp + tn) / (tp + fp + tn + fn)

        results.append(result)

        # add the result of snorkel
    label_matrix = my_df[my_lfs].to_numpy()

    # train the generative model
    label_model = LabelModel(cardinality=2, verbose=False)
    label_model.fit(L_train=label_matrix, class_balance=balance,
                    n_epochs=epochs, log_freq=100, seed=123)

    # apply the trained generative model to predict the result
    Y_pred, Y_prob = label_model.predict(L=label_matrix, return_probs=True, tie_break_policy='abstain')

    result = {}

    result['name'] = 'snorkel'

    coverage = (len(Y_pred) - list(Y_pred).count(-1)) / len(Y_pred)
    result['coverage'] = coverage

    tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
    result['tn'] = tn
    result['fp'] = fp
    result['fn'] = fn
    result['tp'] = tp

    result['precision'] = tp / (tp + fp)
    result['recall'] = tp / (tp + fn)
    result['f1'] = tp / (tp + 0.5 * (fp + fn))
    result['accuracy'] = (tp + tn) / (tp + fp + tn + fn)

    results.append(result)

    return pd.DataFrame.from_records(results)


def show_lfs_stat(my_df, my_lfs, balance, epochs):
    results = []

    y = my_df['label'].to_numpy()

    # add the result of each lf
    for lf in my_lfs:
        result = {}

        Y_pred = my_df[lf].to_numpy()

        coverage = (len(Y_pred) - list(Y_pred).count(-1)) / len(Y_pred)
        result['coverage'] = coverage

        tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
        result['tn'] = tn
        result['fp'] = fp
        result['fn'] = fn
        result['tp'] = tp

        result['precision'] = tp / (tp + fp)
        result['recall'] = tp / (tp + fn)
        result['f1'] = tp / (tp + 0.5 * (fp + fn))
        result['accuracy'] = (tp + tn) / (tp + fp + tn + fn)

        result['name'] = lf

        results.append(result)

        # add the result of snorkel
    label_matrix = my_df[my_lfs].to_numpy()

    # train the generative model
    label_model = LabelModel(cardinality=2, verbose=False)
    label_model.fit(L_train=label_matrix, class_balance=balance,
                    n_epochs=epochs, log_freq=100, seed=123)

    # apply the trained generative model to predict the result
    Y_pred, Y_prob = label_model.predict(L=label_matrix, return_probs=True, tie_break_policy='abstain')

    result = {}

    coverage = (len(Y_pred) - list(Y_pred).count(-1)) / len(Y_pred)
    result['coverage'] = coverage

    tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
    result['tn'] = tn
    result['fp'] = fp
    result['fn'] = fn
    result['tp'] = tp

    result['precision'] = tp / (tp + fp)
    result['recall'] = tp / (tp + fn)
    result['f1'] = tp / (tp + 0.5 * (fp + fn))
    result['accuracy'] = (tp + tn) / (tp + fp + tn + fn)

    result['name'] = 'snorkel'

    results.append(result)

    return results


### Step 1. define the strategy to cluster data points into group
ABSTAIN = -1
UNMATCHED = 0
MATCHED = 1


def count_positive_votes(row):
    count = 0
    for index, value in row.items():
        if value == MATCHED:
            count = count + 1
    return count


def group_by_positive_vote(df, lfs):
    df['group_value'] = df[lfs].apply(lambda r: count_positive_votes(r), axis=1)
    group_values = df['group_value'].unique()

    group_key = 1
    for group_value in group_values:
        df.loc[df['group_value'] == group_value, 'group_key'] = group_key
        group_key = group_key + 1

    return len(group_values)


def group_all_in_one(df, lfs):
    df['group_key'] = 1
    df['group_value'] = 1

    return 1


def count_num_abstain(row):
    count = 0

    for index, value in row.items():
        if value == ABSTAIN:
            count = count + 1

    return count


def group_by_abstain(df, lfs):
    df['group_value'] = df[lfs].apply(lambda r: count_num_abstain(r), axis=1)

    group_values = df['count'].unique()

    group_key = 1
    for group_value in group_values:
        df.loc[df['group_value'] == group_value, 'group_key'] = group_key
        group_key = group_key + 1

    return len(group_values)


def count_num_disagreement(row):
    num_matches = 0
    num_non_matches = 0

    for index, value in row.items():
        if value == MATCHED:
            num_matches = num_matches + 1
        elif value == UNMATCHED:
            num_non_matches = num_non_matches + 1

    return row.count() - abs(num_matches - num_non_matches)


def group_by_disagreement(df, lfs):
    df['group_value'] = df[lfs].apply(lambda r: count_num_disagreement(r), axis=1)

    group_values = df['group_value'].unique()

    group_key = 1
    for group_value in group_values:
        df.loc[df['group_value'] == group_value, 'group_key'] = group_key
        group_key = group_key + 1

    return len(group_values)


def create_votes_string(row):
    votes = ''
    for index, value in row.items():
        votes = votes + str(value)
    return votes


def group_by_uniqueness_votes(df, lfs):
    df['group_value'] = df[lfs].apply(lambda r: create_votes_string(r), axis=1)

    group_values = df['group_value'].unique()

    group_key = 1
    for group_value in group_values:
        df.loc[df['group_value'] == group_value, 'group_key'] = group_key
        group_key = group_key + 1

    return len(group_values)


def grouping_datapoints(strategy_name, df, lfs):
    if strategy_name == 'num_positive_votes':
        return group_by_positive_vote(df, lfs)
    elif strategy_name == 'none':
        return group_all_in_one(df, lfs)
    elif strategy_name == 'abstain':
        return group_by_abstain(df, lfs)
    elif strategy_name == 'disagreement':
        return group_by_disagreement(df, lfs)
    elif strategy_name == 'uniqueness_votes':
        return group_by_uniqueness_votes(df, lfs)


### Step 2. select a top group for query

def round_robin_selection(df, number_groups):
    groups = {}

    # init each group with its own group size
    group_sizes = df.groupby('group_key').size()
    for group_key, v in group_sizes.items():
        groups[group_key] = {}
        groups[group_key]['count'] = v
        groups[group_key]['annotated'] = 0

    # print(groups)

    # look up the annotated points in each group
    annotated_per_group = df[df['selected_index'] > 0].groupby('group_key').size().sort_values(ascending=True)
    for group_key, v in annotated_per_group.items():
        groups[group_key]['annotated'] = v

    # print(groups)

    # find the group with the least number of annotated points
    sorted_group_keys = sorted(groups, key=lambda x: groups[x]['annotated'])

    for k in sorted_group_keys:
        if (groups[k]['count'] - groups[k]['annotated']) == 0:
            continue
        # print(k)
        # print('\r\n')
        return k

    return None


def alway_first_group(df, number_groups):
    return 1


def max_group_value(df, number_groups):
    idx = df[df['selected_index'] == 0]['group_value'].idxmax()
    return df.iloc[idx]['group_key']


def min_group_value(df, number_groups):
    idx = df[df['selected_index'] == 0]['group_value'].idxmin()
    return df.iloc[idx]['group_key']


def KL(a, b):
    a = np.asarray(a, dtype=np.float)
    b = np.asarray(b, dtype=np.float)

    return np.sum(np.where(a != 0, a * np.log(a / b), 0))


def max_kl_group(df, number_groups):
    groups = []
    for gk in range(1, number_groups + 1):
        group = df[(df['selected_index'] > 0) & (df['group_key'] == gk)]
        group_profile = {'group_key': gk, 'group_value': 0}
        if len(group) > 0:
            annotation = group['label']
            prediction = group['snorkel']
            group_profile['group_value'] = KL(prediction, annotation)

        groups.append(group_profile)

    sorted_groups = sorted(groups, key=lambda i: i['group_value'], reverse=True)

    return sorted_groups[0]['group_key']


def select_group(strategy_name, df, number_groups):
    if strategy_name == 'round_robin':
        return round_robin_selection(df, number_groups)
    elif strategy_name == 'none':
        return alway_first_group(df, number_groups)
    elif strategy_name == 'max':
        return max_group_value(df, number_groups)
    elif strategy_name == 'min':
        return min_group_value(df, number_groups)
    elif strategy_name == 'max_kl':
        return max_kl_group(df, number_groups)

    ### Step 3. define the strategy to query a data point out of the selected group


def random_selection(df, gk, num):
    idx = df[(df['selected_index'] == 0) & (df['group_key'] == gk)].sample(n=num, random_state=123).index.values
    return idx


def calculate_margin(x):
    return 1 - abs(x["nonmatch-prob"] - x["match-prob"])


# margin is the difference in probability of positive and negative predictions
def query_by_margin(df, gk, num):
    mask = (df['selected_index'] == 0) & (df['group_key'] == gk)
    df.loc[mask, 'score'] = df.loc[mask].apply(lambda x: calculate_margin(x), axis=1)

    idx = df.loc[mask, 'score'].nlargest(num, keep='first').index.values

    return idx


def calculate_entropy(x):
    proba = [x["nonmatch-prob"], x["match-prob"]]
    return entropy(proba)


def query_by_entropy(df, gk, num):
    mask = (df['selected_index'] == 0) & (df['group_key'] == gk)

    df.loc[mask, 'score'] = df.loc[mask].apply(lambda x: calculate_entropy(x), axis=1)
    idx = df.loc[mask, 'score'].nlargest(num, keep='first').index.values

    if len(idx) == 0:
        idx = df.loc[mask].sample(n=num, random_state=123).index.values
        print(idx, ' because the original selection return null')

    return idx


def calculate_confidence(x):
    return 1 - max(x["nonmatch-prob"], x["match-prob"])


# Least Confidence score
def query_by_confidence(df, gk, num):
    mask = (df['selected_index'] == 0) & (df['group_key'] == gk)
    df.loc[mask, 'score'] = df.loc[mask].apply(lambda x: calculate_confidence(x), axis=1)
    idx = df.loc[mask, 'score'].nlargest(num, keep='first').index.values

    return idx


# probability of true match
def calculate_match_confidence(x):
    return x["match-prob"]


def query_by_match_confidence(df, gk, num):
    mask = (df['selected_index'] == 0) & (df['group_key'] == gk)
    df.loc[mask, 'score'] = df.loc[mask].apply(lambda x: calculate_match_confidence(x), axis=1)
    idx = df.loc[mask, 'score'].nlargest(num, keep='first').index.values

    return idx


def take_positive_first(df, gk, num):
    df['score'] = df[(df['selected_index'] == 0) & (df['group_key'] == gk)].apply(lambda x: calculate_entropy(x),
                                                                                  axis=1)
    sub_df = df[(df['selected_index'] == 0) & (df['group_key'] == gk) & (df['prediction'] == 1)]
    if len(sub_df) == 0:
        sub_df = df[(df['selected_index'] == 0) & (df['group_key'] == gk) & (df['prediction'] == 0)]

    idx = sub_df.nlargest(num, keep='first').index.values

    return idx


def select_sample_within_group(strategy_name, df, gk, num, active_learning):
    if active_learning:  # only for active learning case
        # make sure there are both matches and non-matches in the annotated group, otherwise, just do the random sampling
        num_unique_labels = len(df[(df['selected_index'] > 0) & (df['group_key'] == gk)]['prediction'].unique())
        if num_unique_labels < 2:
            # print("do random sampling before having some matches and non-matches in the annotated part")
            return random_selection(df, gk, num)
        #         num_annotations = len(df[(df['selected_index']> 0) & (df['group_key']==gk)]['prediction'])
    #         if num_annotations < 20:
    #             print("select positive matches first")
    #             return query_by_match_confidence(df, gk, num)

    if strategy_name == 'random':
        return random_selection(df, gk, num)
    elif strategy_name == 'minimal_margin':
        return query_by_margin(df, gk, num)
    elif strategy_name == 'entropy':
        return query_by_entropy(df, gk, num)
    elif strategy_name == 'least_confidence':
        return query_by_confidence(df, gk, num)
    elif strategy_name == 'match_confidence':
        return query_by_match_confidence(df, gk, num)
    elif strategy_name == 'positive_first':
        return take_positive_first(df, gk, num)

    ### Step 4. define the strategy to select a subset of all provided LFs to achieve the best performance


# define the strategy to select a subset of all provided LFs to achieve the best performance

def evaluate_lf_set(lfset, dataset, epochs):
    inputs = dataset[lfset].to_numpy()
    y = dataset.label.values

    metrics = ["f1", "accuracy", "precision", "recall", "coverage"]
    scorer = Scorer(metrics=metrics)

    label_model = LabelModel(cardinality=2, verbose=False)
    label_model.fit(L_train=inputs, n_epochs=epochs, log_freq=100, seed=123)

    Y_pred, Y_prob = label_model.predict(L=inputs, return_probs=True, tie_break_policy='abstain')

    results = scorer.score(y, Y_pred, Y_prob)
    return {'name': lfset, 'result': results}


#     metrics=["f1", "accuracy", "precision", "recall", "coverage"]
#     scorer = Scorer(metrics=metrics)

def select_by_ranked_f1(df, lfs, epochs):
    # take the annotated data points
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']

    total_num_annotation = len(df[df.selected_index > 0])
    if total_num_annotation < 500:
        return lfs

    # evaluate all LFs over the selected data points
    methods = []
    for lf in lfs:
        y_lf = df.loc[mask, lf]
        y_label = df.loc[mask, 'label']

        tn, fp, fn, tp = confusion_matrix(y_label, y_lf, labels=[0, 1]).ravel()
        coverage = (len(y_lf) - list(y_lf).count(-1)) / len(y_lf)

        result = {}

        result['precision'] = tp / (tp + fp)
        result['recall'] = tp / (tp + fn)
        result['f1'] = tp / (tp + 0.5 * (fp + fn))
        result['coverage'] = tp / (tp + 0.5 * (fp + fn))

        result['tn'] = tn
        result['fp'] = fp
        result['fn'] = fn
        result['tp'] = tp

        method = {}
        method['name'] = lf
        method['result'] = result
        methods.append(method)

    # sort out all LFs according to their performance over the selected data points
    sorted_methods = sorted(methods, key=lambda x: x['result']['f1'], reverse=True)

    # construct a new lfs set 
    sortedlfs = []
    for m in sorted_methods:
        sortedlfs.append(m['name'])

    # print(sortedlfs)
    # print(sorted_methods)

    # iterate and evaluate the LFs combination by adding a new one each time
    best_result = None
    for n in range(3, len(sortedlfs) + 1):
        selected_lfs = sortedlfs[:n]
        # print(selected_lfs)
        evaluation = evaluate_lf_set(selected_lfs, df, epochs)

        if best_result == None:
            best_result = evaluation
        elif best_result['result']['f1'] < evaluation['result']['f1']:
            best_result = evaluation
        else:
            break

    # print(best_result)

    return selected_lfs


def select_by_ranked_recall(df, lfs, epochs):
    return lfs


def select_by_ranked_score(df, lfs, epochs):
    return lfs


def all_lfs(df, lfs, epochs):
    return lfs


def select_best_lfs(strategy_name, df, lfs, epochs):
    # do the filtering of LFs only after the number of annotated data points > 10, at least one match and non-match
    mask = (df['selected_index'] > 0)
    num_annotated_datapoints = len(df.loc[mask])
    num_unique_labels = len(df.loc[mask, 'prediction'].unique())
    if num_annotated_datapoints <= 10 or num_unique_labels < 2:
        # print('no selection of lfs due to limited number of annotated samples')
        return all_lfs(df, lfs, epochs)

    if strategy_name == 'all':
        return all_lfs(df, lfs, epochs)
    elif strategy_name == 'f1_first':
        return select_by_ranked_f1(df, lfs, epochs)
    elif strategy_name == 'recall_first':
        return select_by_ranked_recall(df, lfs, epochs)
    elif strategy_name == 'combined_score':
        return select_by_ranked_score(df, lfs, epochs)

    ### Step 5. define the strategy to ensemble the results of the selected LFs


def apply_majority_voting(df, lfs):
    label_matrix = df[lfs].to_numpy()

    label_model = MajorityLabelVoter()

    # training is NOT required

    Y_pred = label_model.predict(L=label_matrix)
    Y_prob = label_model.predict_proba(L=label_matrix)

    df['snorkel'] = Y_pred
    df['nonmatch-prob'] = Y_prob[:, 0]
    df['match-prob'] = Y_prob[:, 1]

    # overwrite the annotated data points with their golden labels
    df['prediction'] = df['snorkel']
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def apply_none(df, selected_lfs):
    num_annotated_datapoints = len(df[df['selected_index'] > 0])
    labels = df[df['selected_index'] > 0]['label'].unique()
    if num_annotated_datapoints < 2 or len(labels) < 2:
        df['prediction'] = 0
        df['nonmatch-prob'] = 0
        df['match-prob'] = 0

        # overwrite the annotated data points with their golden labels
        mask = (df.selected_index > 0)

        if len(mask) > 0:
            df.loc[mask, 'prediction'] = df.loc[mask, 'label']

        return

    if selected_lfs in df.columns:
        df['prediction'] = df[selected_lfs]
        df['nonmatch-prob'] = df[selected_lfs + '_0']
        df['match-prob'] = df[selected_lfs + '_1']

        # overwrite the annotated data points with their golden labels        
        mask = (df.selected_index > 0)
        df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def apply_snorkel(df, lfs, balance, epochs):
    label_matrix = df[lfs].to_numpy()

    # train the generative model
    label_model = LabelModel(cardinality=2, verbose=False)
    label_model.fit(L_train=label_matrix, class_balance=balance,
                    n_epochs=epochs, log_freq=100, seed=123)

    # apply the trained generative model to predict the result
    Y_pred, Y_prob = label_model.predict(L=label_matrix, return_probs=True, tie_break_policy='abstain')

    df['snorkel'] = Y_pred
    df['nonmatch-prob'] = Y_prob[:, 0]
    df['match-prob'] = Y_prob[:, 1]

    # overwrite the annotated data points with their golden labels
    df['prediction'] = df['snorkel']
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def apply_snorkel_with_vote_correction(df, lfs, balance, epochs):
    # correct the votes of all LFs for the data points that have been annotated by human
    mask = (df.selected_index > 0)
    df.loc[mask, lfs] = df.loc[mask, 'label']

    # prepare the label matrix out of the dataframe
    label_matrix = df[lfs].to_numpy()

    # train the generative model
    label_model = LabelModel(cardinality=2, verbose=False)
    label_model.fit(L_train=label_matrix, class_balance=balance,
                    n_epochs=epochs, log_freq=100, seed=123)

    # apply the trained generative model to predict the result
    Y_pred, Y_prob = label_model.predict(L=label_matrix, return_probs=True, tie_break_policy='abstain')

    df['snorkel'] = Y_pred
    df['nonmatch-prob'] = Y_prob[:, 0]
    df['match-prob'] = Y_prob[:, 1]

    # overwrite the annotated data points with their golden labels
    df['prediction'] = df['snorkel']
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def apply_snorkel_with_annotated_labels(df, lfs, balance, epochs):
    # correct the votes of all LFs for the data points that have been annotated by human
    df['annotation'] = -1
    mask = (df.selected_index > 0)
    df.loc[mask, 'annotation'] = df.loc[mask, 'label']

    # prepare the label matrix out of the dataframe
    label_matrix = df[lfs].to_numpy()

    # initialize the groups of LFs
    cliques = []
    for i in range(len(lfs)):
        cliques.append([i])

    weasul_label_model = WeaSulLabelModel(
        n_epochs=epochs,
        active_learning=True,
        hide_progress_bar=True)

    weasul_label_model.fit(
        label_matrix=label_matrix,
        cliques=cliques,
        class_balance=np.array(balance),
        ground_truth_labels=df['annotation'].to_numpy()
    )

    probs = weasul_label_model.predict(label_matrix, weasul_label_model.mu, weasul_label_model.E_S)
    Y_prob = probs.cpu().detach().numpy()
    Y_pred = torch.argmax(probs, dim=1).cpu().detach().numpy()

    df['snorkel'] = Y_pred
    df['nonmatch-prob'] = Y_prob[:, 0]
    df['match-prob'] = Y_prob[:, 1]

    # overwrite the annotated data points with their golden labels
    df['prediction'] = df['snorkel']
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


from sklearn.metrics import precision_score


def get_avg_nonmatch_prob(df):
    return


def apply_snorkel_with_init_precision(df, lfs, balance, epochs):
    # calculate the precision of lfs over the annotated data points
    mask = (df.selected_index > 0)
    y_true = df.loc[mask, 'label']

    init_lf_precision = []
    for i in range(len(lfs)):
        lf = lfs[i]
        y_pred = df.loc[mask, lf]
        tn, fp, fn, tp = confusion_matrix(y_true, y_pred, labels=[0, 1]).ravel()
        # print(tn, fp, fn, tp)

        if (tp + fp) == 0:
            precision = 0.1
        else:
            precision = 0.7 * (tp / (tp + fp))

        init_lf_precision.append(precision)

    # print(init_lf_precision)

    # prepare the label matrix out of the dataframe
    mask = (df.selected_index == 0)
    label_matrix = df.loc[mask, lfs].to_numpy()
    if len(label_matrix) > 0:
        # train the generative model
        label_model = LabelModel(cardinality=2, verbose=False)
        label_model.fit(L_train=label_matrix, class_balance=balance,
                        prec_init=init_lf_precision,
                        n_epochs=epochs, log_freq=100, seed=123)

        # apply the trained generative model to predict the result
        Y_pred, Y_prob = label_model.predict(L=label_matrix, return_probs=True, tie_break_policy='abstain')

        df.loc[mask, 'snorkel'] = Y_pred
        df.loc[mask, 'prediction'] = Y_pred

        # if 'random_forest_0' in df.columns:
        #     df.loc[mask, 'nonmatch-prob'] = df.loc[mask, 'random_forest_0']
        #     df.loc[mask, 'match-prob'] = df.loc[mask, 'random_forest_1']
        # else:
        df.loc[mask, 'nonmatch-prob'] = Y_prob[:, 0]
        df.loc[mask, 'match-prob'] = Y_prob[:, 1]

        # overwrite the annotated data points with their golden labels
        df['prediction'] = df['snorkel']

    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def apply_normal_active_learning(df, lfs, features, balance, epochs, model_type):
    num_annotated_datapoints = len(df[df['selected_index'] > 0])
    labels = df[df['selected_index'] > 0]['label'].unique()
    if num_annotated_datapoints < 2 or len(labels) < 2:
        df['prediction'] = 0
        df['nonmatch-prob'] = 0
        df['match-prob'] = 0

        # overwrite the annotated data points with their golden labels
        mask = (df.selected_index > 0)

        if len(mask) > 0:
            df.loc[mask, 'prediction'] = df.loc[mask, 'label']

        return

    if model_type == 'random_forest':
        model = RandomForestClassifier()
    elif model_type == 'logistic_regression':
        model = LogisticRegression(C=1, penalty='l1', solver='liblinear')
    elif model_type == 'mlp':
        model = MLPClassifier(max_iter=200, random_state=21)

    train_x = df[df['selected_index'] > 0][features].astype(float).to_numpy()
    train_y = df[df['selected_index'] > 0]['label']

    # train the random forest model
    model.fit(train_x, train_y)

    # predict the rest
    test_x = df[features].astype(float).to_numpy()
    df['prediction'] = model.predict(test_x)
    Y_prob = model.predict_proba(test_x)
    df['nonmatch-prob'] = Y_prob[:, 0]
    df['match-prob'] = Y_prob[:, 1]

    # overwrite the annotated data points with their golden labels
    mask = (df.selected_index > 0)
    df.loc[mask, 'prediction'] = df.loc[mask, 'label']


def lfs_ensemble(strategy_name, df, lfs, features, balance, epochs):
    if len(lfs) == 0:
        return

    if strategy_name == 'majority':
        return apply_majority_voting(df, lfs)
    elif strategy_name == 'none':
        return apply_none(df, 'random_forest')
    elif strategy_name == 'snorkel':
        return apply_snorkel(df, lfs, balance, epochs)
    elif strategy_name == 'snorkel_with_corrected_votes':
        return apply_snorkel_with_vote_correction(df, lfs, balance, epochs)
    elif strategy_name == 'snorkel_with_annotated_labels':
        return apply_snorkel_with_annotated_labels(df, lfs, balance, epochs)
    elif strategy_name == 'snorkel_with_init_precision':
        # trigger this only after the number of annotated data points > 10, at least one match and non-match
        mask = (df['selected_index'] > 0)
        num_annotated_datapoints = len(df.loc[mask])
        num_unique_labels = len(df.loc[mask, 'label'].unique())
        if num_annotated_datapoints <= 10 or num_unique_labels < 2:
            print('no selection of lfs due to limited number of annotated samples')
            return apply_snorkel(df, lfs, balance, epochs)
        else:
            return apply_snorkel_with_init_precision(df, lfs, balance, epochs)
    elif strategy_name == 'normal_active_learning_rf':
        return apply_normal_active_learning(df, lfs, features, balance, epochs, 'random_forest')
    elif strategy_name == 'normal_active_learning_lg':
        return apply_normal_active_learning(df, lfs, features, balance, epochs, 'logistic_regression')
    elif strategy_name == 'normal_active_learning_mlp':
        return apply_normal_active_learning(df, lfs, features, balance, epochs, 'mlp')


### Step 6. evaluate the predicted result

def evaluate_result(df):
    if len(df[df['prediction'] >= 0]) == 0:
        return

    result = {}

    # calculate the performance metrics for the entire dataset
    Y_pred = df['prediction']
    y = df['label'].to_numpy()

    tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
    result['tn'] = tn
    result['fp'] = fp
    result['fn'] = fn
    result['tp'] = tp

    result['precision'] = tp / (tp + fp)
    result['recall'] = tp / (tp + fn)
    result['f1'] = tp / (tp + 0.5 * (fp + fn))

    # calculate the performance metrics for the annotated part
    Y_pred = df[df['selected_index'] > 0]['prediction']
    y = df[df['selected_index'] > 0]['label'].to_numpy()

    tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
    result['a-tn'] = tn
    result['a-fp'] = fp
    result['a-fn'] = fn
    result['a-tp'] = tp

    result['a-precision'] = tp / (tp + fp)
    result['a-recall'] = tp / (tp + fn)
    result['a-f1'] = tp / (tp + 0.5 * (fp + fn))

    # calculate the performance metrics for the predicted part
    Y_pred = df[df['selected_index'] == 0]['prediction']
    y = df[df['selected_index'] == 0]['label'].to_numpy()

    tn, fp, fn, tp = confusion_matrix(y, Y_pred, labels=[0, 1]).ravel()
    result['p-tn'] = tn
    result['p-fp'] = fp
    result['p-fn'] = fn
    result['p-tp'] = tp

    result['p-precision'] = tp / (tp + fp)
    result['p-recall'] = tp / (tp + fn)
    result['p-f1'] = tp / (tp + 0.5 * (fp + fn))

    # add the calculated effort for annotation and confirmation/checking
    number_annotation = df[df['selected_index'] > 0].shape[0]
    number_predicted_matches = df[(df['selected_index'] == 0) & (df['prediction'] == 1)].shape[0]
    result['human_effort'] = number_annotation + number_predicted_matches

    # calculate the number of true matched captured after this amount of effort
    number_annotated_true_matches = df[(df['selected_index'] > 0) & (df['label'] == 1)].shape[0]
    number_predicted_true_matches = \
    df[(df['selected_index'] == 0) & (df['prediction'] == 1) & (df['label'] == 1)].shape[0]
    result['num_true_matches'] = number_annotated_true_matches + number_predicted_true_matches

    return result


### Step 7. update the LF pool with newly created LFs

#### create new LFs based on the annotated golden labels and predicted weak labels

# random forest
@ray.remote(num_returns=3)
def random_forest(train_x, train_y, test_x):
    model = RandomForestClassifier()

    model.fit(train_x, train_y)

    pred = model.predict(test_x)
    prob = model.predict_proba(test_x)

    return pred, prob[:, 0], prob[:, 1]


# logistic regression
@ray.remote(num_returns=3)
def logistic_regression(train_x, train_y, test_x):
    model = LogisticRegression(C=1, penalty='l1', solver='liblinear')
    model.fit(train_x, train_y)

    pred = model.predict(test_x)
    prob = model.predict_proba(test_x)

    return pred, prob[:, 0], prob[:, 1]


# XGBoost
@ray.remote(num_returns=3)
def xgboost(train_x, train_y, test_x):
    model = xgb.XGBClassifier(subsample=0.5, use_label_encoder=False, eval_metric='auc')
    model.fit(train_x, train_y)

    pred = model.predict(test_x)
    prob = model.predict_proba(test_x)

    return pred, prob[:, 0], prob[:, 1]


# MLP
@ray.remote(num_returns=3)
def mlp(train_x, train_y, test_x):
    # model = MLPClassifier(max_iter=100, verbose=10, random_state=21)
    model = MLPClassifier(max_iter=200, random_state=21)

    model.fit(train_x, train_y)

    pred = model.predict(test_x)
    prob = model.predict_proba(test_x)

    return pred, prob[:, 0], prob[:, 1]

#
# def augmented_with_ml_lf(df, features):
#     num_annotation = len(df[df['annotation'] >= 0])
#     if num_annotation < 20:
#         train_x = df[df['prediction'] >= 0][features].astype(float).to_numpy()
#         train_y = df[df['prediction'] >= 0]['prediction']
#     #         train_x = df[features].astype(float).to_numpy()
#     #         train_y = df['prediction']
#     else:
#         train_x = df[df['annotation'] >= 0][features].astype(float).to_numpy()
#         train_y = df[df['annotation'] >= 0]['prediction']
#
#     test_x = df[features].astype(float).to_numpy()
#
#     train_x_ref = ray.put(train_x)
#     train_y_ref = ray.put(train_y)
#     test_x_ref = ray.put(test_x)
#
#     column_names = []
#     column_values_ref = []
#
#     models = ['random_forest', 'logistic_regression', 'mlp']
#     # models = ['random_forest', 'logistic_regression', 'xgboost', 'mlp']
#     for model in models:
#         if model == 'random_forest':
#             pred, prob_0, prob_1 = random_forest.remote(train_x_ref, train_y_ref, test_x_ref)
#         elif model == 'logistic_regression':
#             pred, prob_0, prob_1 = logistic_regression.remote(train_x_ref, train_y_ref, test_x_ref)
#         elif model == 'xgboost':
#             pred, prob_0, prob_1 = xgboost.remote(train_x_ref, train_y_ref, test_x_ref)
#         elif model == 'mlp':
#             pred, prob_0, prob_1 = mlp.remote(train_x_ref, train_y_ref, test_x_ref)
#
#         column_names.append(model)
#         column_values_ref.append(pred)
#         column_names.append(model + '_0')
#         column_values_ref.append(prob_0)
#         column_names.append(model + '_1')
#         column_values_ref.append(prob_1)
#
#         # fetch the result from ray
#     column_values = ray.get(column_values_ref)
#
#     # update the dataframe
#     for i in range(len(column_names)):
#         column_name = column_names[i]
#         df[column_name] = column_values[i]
#
#     return models


def augmented_with_ml_lf(df, features):
    train_x = df[df['prediction'] >= 0][features].astype(float).to_numpy()
    train_y = df[df['prediction'] >= 0]['prediction']
    test_x = df[features].astype(float).to_numpy()

    train_x_ref = ray.put(train_x)
    train_y_ref = ray.put(train_y)
    test_x_ref = ray.put(test_x)

    column_names = []
    column_values_ref = []
    models = ['random_forest', 'logistic_regression', 'xgboost', 'mlp']
    for model in models:
        if model == 'random_forest':
            pred, prob_0, prob_1 = random_forest.remote(train_x_ref, train_y_ref, test_x_ref)
        elif model == 'logistic_regression':
            pred, prob_0, prob_1 = logistic_regression.remote(train_x_ref, train_y_ref, test_x_ref)
        elif model == 'xgboost':
            pred, prob_0, prob_1 = xgboost.remote(train_x_ref, train_y_ref, test_x_ref)
        elif model == 'mlp':
            pred, prob_0, prob_1 = mlp.remote(train_x_ref, train_y_ref, test_x_ref)

        column_names.append(model)
        column_values_ref.append(pred)
        column_names.append(model + '_0')
        column_values_ref.append(prob_0)
        column_names.append(model + '_1')
        column_values_ref.append(prob_1)

        # fetch the result from ray
    column_values = ray.get(column_values_ref)

    # update the dataframe
    for i in range(len(column_names)):
        column_name = column_names[i]
        df[column_name] = column_values[i]

    return models

#### tune existing labeling functions

class _Param(object):
    def __init__(self, func, x0=[0], bounds=[(0, 100)]):
        self.func = func
        self.x0 = x0
        self.bounds = bounds

    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)


def Param(func=None, x0=[0], bounds=[(0, 100)]):
    if func:
        return _Param(func)
    else:
        def wrapper(func):
            return _Param(func, x0, bounds)

        return wrapper


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_class_name_similarity_ma(row, x0):
    T0 = x0[0]

    r = row['class_name_words_similarity_a']
    # print(r, T0)

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_class_name_similarity_mb(row, x0):
    T0 = x0[0]

    r = row['class_name_words_similarity_b']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_label_similarity_ma(row, x0):
    T0 = x0[0]

    r = row['label_words_similarity_a']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_label_similarity_mb(row, x0):
    T0 = x0[0]

    r = row['label_words_similarity_b']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_comment_similarity_ma(row, x0):
    T0 = x0[0]

    r = row['comment_similarity_a']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0.0, 1.0)])
def LF_comment_similarity_mb(row, x0):
    T0 = x0[0]

    r = row['comment_similarity_b']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[5], bounds=[(1, 5)])
def LF_num_common_words(row, x0):
    T0 = x0[0]

    r = row['num_common_words']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0, 1.0)])
def LF_random_forest_prob(row, x0):
    T0 = x0[0]

    r = row['random_forest_1']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0, 1.0)])
def LF_logistic_regression_prob(row, x0):
    T0 = x0[0]

    r = row['logistic_regression_1']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0, 1.0)])
def LF_xgboost_prob(row, x0):
    T0 = x0[0]

    r = row['xgboost_1']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0


@Param(x0=[1.0], bounds=[(0, 1.0)])
def LF_mlp_prob(row, x0):
    T0 = x0[0]

    r = row['mlp_1']

    if r == None:
        print("feature is None")
        return -1

    if r >= T0:
        return 1
    else:
        return 0




def optimize_wrapper(x0, annotated_df, predicted_df, lf_func, old_perf, batch_size):
    cur_performance = {}
    cur_performance['x0'] = x0

    # count the current number of true matches that we gain from the annotation
    annotation_gain = len(annotated_df[annotated_df['label'] == 1])

    # for the annotation part 
    y_df = annotated_df.apply(lf_func, x0=x0, axis=1).to_numpy()
    mask = (y_df != -1)
    y_pred = y_df[mask]
    y_ref = annotated_df.iloc[mask].prediction
    cur_performance['annotated'] = compute_performance(y_ref, y_pred)

    # for the predicted part
    y_df = predicted_df.apply(lf_func, x0=x0, axis=1).to_numpy()
    mask = (y_df != -1)
    y_pred = y_df[mask]
    y_ref = predicted_df.iloc[mask].prediction
    cur_performance['predicted'] = compute_performance(y_ref, y_pred)

    # calculate a score for each parameter accordingly
    score = compute_score(cur_performance, old_perf, annotation_gain, batch_size)

    #print(x0, score, annotation_gain, cur_performance)

    return score


def compute_score(cur_performance, old_performance, annotation_gain, batch_size):
    a_tp_old = old_performance['annotated']['tp']
    a_fp_old = old_performance['annotated']['fp']

    p_tp_old = old_performance['predicted']['tp']
    p_fp_old = old_performance['predicted']['fp']

    a_tp_cur = cur_performance['annotated']['tp']
    a_fp_cur = cur_performance['annotated']['fp']

    p_tp_cur = cur_performance['predicted']['tp']
    p_fp_cur = cur_performance['predicted']['fp']

    cur_predicted_precision = 0
    if (p_tp_cur + p_fp_cur) > 0:
        cur_predicted_precision = p_tp_cur / (p_tp_cur + p_fp_cur)

    if p_fp_cur > 0:
        fp_ratio = batch_size/p_fp_cur
    else:
        fp_ratio = 0

    # score = a_tp_old - a_tp_cur  - cur_predicted_precision
    # score = annotation_gain - a_tp_cur  - cur_predicted_precision
    # score = (a_tp_old + p_tp_old) - (a_tp_cur + p_tp_cur)  - cur_predicted_precision
    # score = annotation_gain - (a_tp_cur + p_tp_cur)  - cur_predicted_precision

    score = annotation_gain - a_tp_cur - cur_predicted_precision - fp_ratio

    #print(score)

    return score


def compute_performance(y_ref, y_pred):
    tn, fp, fn, tp = confusion_matrix(y_ref, y_pred, labels=[0, 1]).ravel()
    coverage = (len(y_pred) - list(y_pred).count(-1)) / len(y_pred)

    performance = {}
    performance['tp'] = tp
    performance['fp'] = fp
    performance['tn'] = tn
    performance['fn'] = fn
    performance['coverage'] = coverage

    return performance




def compute_optimal_parameter(df, lf_name, iteration, batch_size):
    if lf_name in locals():
        my_lf = locals()[lf_name]
    elif lf_name in globals():
        my_lf = globals()[lf_name]

    x0 = my_lf.x0
    bounds = my_lf.bounds
    # print('=====', lf_name, '=========')
    # print("x0", x0)
    # print("bounds", bounds)

    # calculate the estimated performance on both annotated part and the predicted part in a separated way
    annotated_df = df[df['selected_index'] > 0]
    predicted_df = df[df['selected_index'] == 0]

    old_performance = {}
    old_performance['x0'] = x0

    # for the annotated part
    y_df = annotated_df.apply(my_lf, x0=x0, axis=1).to_numpy()
    mask = (y_df != -1)
    y_pred = y_df[mask]
    y_ref = annotated_df.iloc[mask].prediction
    old_performance['annotated'] = compute_performance(y_ref, y_pred)

    # for the predicted part    
    y_df = predicted_df.apply(my_lf, x0=x0, axis=1).to_numpy()
    mask = (y_df != -1)
    y_pred = y_df[mask]
    y_ref = predicted_df.iloc[mask].prediction
    old_performance['predicted'] = compute_performance(y_ref, y_pred)

    # print('OLD Performance:', old_performance)

    # if there is no improvement for a certain runs, we low the boundary to do the parameter search
    num_predicted_matches = old_performance['predicted']['tp'] + old_performance['predicted']['fp']
    # print('num_predicted_matches = ', num_predicted_matches)
    if num_predicted_matches == 0:
        # lower the upper boundary
        if x0[0] - 0.02 > 0:
            bounds[0] = (0, x0[0] - 0.02)
            x0[0] = x0[0] - 0.02
        else:
            bounds[0] = (0, x0[0])

        my_lf.bounds = bounds
        # print('updated boundary ', bounds)

    start = time.time()
    # SLSQP, Powell
    res = minimize(optimize_wrapper, x0, method='Powell', bounds=bounds,
                   args=(annotated_df, predicted_df, my_lf, old_performance, batch_size))
    end = time.time()
    # print("Optimizing Time: ", end - start, "s")
    # print(res)

    new_x0 = res.x
    # print('new x0 = ', new_x0)
    #
    # new_performance = {}
    # new_performance['x0'] = [new_x0]
    #
    # # for the annotated part
    # y_df = annotated_df.apply(my_lf, x0=[new_x0], axis=1).to_numpy()
    # mask = (y_df != -1)
    # y_pred = y_df[mask]
    # y_ref = annotated_df.iloc[mask].prediction
    # new_performance['annotated'] = compute_performance(y_ref, y_pred)
    #
    # # for the predicted part
    # y_df = predicted_df.apply(my_lf, x0=[new_x0], axis=1).to_numpy()
    # mask = (y_df != -1)
    # y_pred = y_df[mask]
    # y_ref = predicted_df.iloc[mask].prediction
    # new_performance['predicted'] = compute_performance(y_ref, y_pred)
    #
    # print('NEW Performance: ', new_performance)
    # print('\r\n')

    # update the parameter of this label function
    my_lf.x0 = [new_x0]

    return [new_x0], my_lf
#
# @ray.remote
# def tuning_lf(df, lf_name, iteration):
#     return compute_optimal_parameter(df, lf_name, iteration)
#
# def augmented_with_tuned_lf(df, iteration):
#     updated_lfs = []
#
#     tunable_lfs = ['LF_class_name_similarity_ma', 'LF_class_name_similarity_mb',
#                    'LF_label_similarity_ma', 'LF_label_similarity_mb',
#                    'LF_comment_similarity_ma', 'LF_comment_similarity_mb',
#                    'LF_num_common_words',
#                    'LF_random_forest_prob',
#                    'LF_logistic_regression_prob',
#                    'LF_xgboost_prob',
#                    'LF_mlp_prob']
#
#     for lf_name in tunable_lfs:
#         # search for the best parameter
#         optimal_parameter, my_lf = compute_optimal_parameter(df, lf_name, iteration)
#
#         # apply the new parameter
#         y_pred = df.apply(my_lf, x0=optimal_parameter, axis=1).to_numpy()
#
#         # write back the updated labels
#         tuned_lf = lf_name + "_tuned"
#         df[tuned_lf] = y_pred
#
#         print(tuned_lf)
#
#         updated_lfs.append(tuned_lf)
#
#     return updated_lfs


@ray.remote
def tuning_lf(df, lf_name, iteration, batch_size):
    # search for the best parameter
    optimal_parameter, my_lf = compute_optimal_parameter(df, lf_name, iteration, batch_size)

    # apply the new parameter
    y_pred = df.apply(my_lf, x0=optimal_parameter, axis=1).to_numpy()

    return y_pred

def augmented_with_tuned_lf(df, iteration, tunable_lfs, batch_size):
    updated_lfs = []

    df_ref = ray.put(df)
    y_pred_refs = []
    tuned_lfs = []

    for lf_name in tunable_lfs:
        y_pred_ref = tuning_lf.remote(df_ref, lf_name, iteration, batch_size)
        y_pred_refs.append(y_pred_ref)
        tuned_lfs.append(lf_name + "_tuned")

    y_preds = ray.get(y_pred_refs)

    for i in range(len(tuned_lfs)):
        tuned_lf = tuned_lfs[i]
        df[tuned_lf] = y_preds[i]

        updated_lfs.append(tuned_lf)

    return updated_lfs


def combine_lf_sets(list_1, list_2):
    set_1 = set(list_1)
    set_2 = set(list_2)

    list_2_items_not_in_list_1 = list(set_2 - set_1)
    combined_list = list_1 + list_2_items_not_in_list_1

    return combined_list


def trigger_slow_loop(df, features, iteration, with_lf_tuning, tunable_lfs, batch_size):
    new_lfs = []

    # make sure there are both matches and non-matches in the annotated group, otherwise, just do the random sampling
    num_unique_labels = len(df[df['selected_index'] > 0]['prediction'].unique())
    if num_unique_labels < 2:
        # print("====== ", num_unique_labels)
        return new_lfs

    # train/retrain a machine learning model to create a new LF    
    ml_lfs = augmented_with_ml_lf(df, features)
    new_lfs += ml_lfs

    # reconfigure a tunable LF to add/update LF
    if with_lf_tuning:
        tuned_lfs = augmented_with_tuned_lf(df, iteration, tunable_lfs, batch_size)
        new_lfs += tuned_lfs

    return new_lfs


### Step 8. visualize the generated results

import seaborn as sns
import matplotlib.pyplot as plt


def plot_result(results, metric_name, scale):
    measurement_list = []
    for exp in results:
        experiment = results[exp]
        for iteration in experiment:
            r = experiment[iteration]
            measurement = {
                'methods': exp,
                'iteration': iteration,
                'metric': r[metric_name] * scale
            }

            measurement_list.append(measurement)

    exp_result_df = pd.DataFrame(measurement_list)

    fig, ax = plt.subplots(figsize=(10, 5))
    # plt.ylim(0, 100)

    sns.set(font_scale=1)
    p = sns.lineplot(x='iteration', y='metric', hue='methods', ax=ax, data=exp_result_df,
                     style="methods", markers=True)

    #    p.set_xlabel("#iteration")
    p.set_xlabel("percentage of annotated data (%)")
    p.set_ylabel(metric_name)

    plt.show()


def plot_cost_over_gain(results):
    measurement_list = []
    for exp in results:
        costMap = {}
        experiment = results[exp]
        previous_gain = 0
        for iteration in experiment:
            m = experiment[iteration]

            gain = m['num_true_matches']
            cost = m['human_effort']

            if gain == 0 and cost == 0:
                continue

            if gain not in costMap:
                for j in range(gain, previous_gain, -1):
                    measurement = {
                        'methods': exp,
                        'gain': j,
                        'cost': cost
                    }

                    measurement_list.append(measurement)

                costMap[gain] = cost

            previous_gain = gain

    exp_result_df = pd.DataFrame(measurement_list)

    fig, ax = plt.subplots(figsize=(10, 5))
    # plt.ylim(0, 100)

    sns.set(font_scale=1)
    p = sns.lineplot(x='gain', y='cost', hue='methods', ax=ax, data=exp_result_df,
                     style="methods", markers=True)

    p.set_xlabel("gain")
    p.set_ylabel("cost")

    plt.show()


def plot_comparison_result(results, methods, names, interval, filename):
    metrics = ['human_effort', 'num_true_matches', 'a-tp', 'p-tp']
    labels = ['cost', 'gain', '#annoate_matches', '#verified_matches']
    markers = ['o', '*', '+', 'x']
    fsize = 20

    num_fig = len(metrics)
    fig, axs = plt.subplots(1, num_fig, figsize=(28, 5))
    lines = []

    for i in range(len(metrics)):
        metric = metrics[i]
        for j in range(len(methods)):
            method = methods[j]
            rdf = pd.DataFrame(results[method])
            x = rdf.columns
            xticks = np.arange(min(x), max(x) + 1, interval)
            y = rdf.loc[metric]

            axs[i].set_xticks(xticks)
            line = axs[i].plot(x, y, linestyle='-', marker=markers[j])
            lines.append(lines)

        axs[i].set_ylabel(labels[i], fontsize=fsize)
        axs[i].grid(True)

    fig.legend(lines,  # The line objects
               labels=names,  # The labels for each line
               loc="upper center",  # Position of legend,
               bbox_to_anchor=(0.45, 1.0), ncol=len(methods)
               )

    plt.savefig(filename, transparent=False)
    plt.close(fig)


# Main logic of the fast loop

def run_experiment(experiment_conf, df, pool_lfs, tunable_lfs, feature_set, total_size, num_iteration, interval_slow_loop,
                   batch_size, balance, epochs, debug=False):
    if debug:
        my_df = df
    else:
        my_df = df.copy()

    my_df.update(my_df[feature_set].fillna(0))

    # initialization
    my_df['selected_index'] = 0

    measurements = {}

    new_lfs = []

    # initialize the prediction result given by weak supervision without human annotations     
    lfs_ensemble(experiment_conf['lf_ensemble'], my_df, pool_lfs, feature_set, balance, epochs)
    r = evaluate_result(my_df)
    measurements[0] = r

    pbar = tqdm(total=num_iteration)
    iteration_slow_loop = 0
    i = 1
    while (i <= num_iteration):
        # print("iteration = ", i)

        start = time.time()

        j = 0
        while j < batch_size:
            if (debug == True):
                print('\r\n===' + str(i) + '===')

            number_of_groups = grouping_datapoints(experiment_conf['datapoint_grouping'], my_df, pool_lfs)
            if (debug == True):
                print("# of GROUPs = " + str(number_of_groups))

                # group selection
            selected_group = select_group(experiment_conf['group_selection'], my_df, number_of_groups)
            if (debug == True):
                print("SELECTED GROUP = " + str(selected_group))

            # datapoint selection
            active_learning_flag = False
            if 'normal_active_learning' in experiment_conf['lf_ensemble']:
                active_learning_flag = True

            required_num = batch_size - j

            if len(pool_lfs) == 0:
                selected_sample_loc = select_sample_within_group('random',
                                                                 my_df, selected_group, required_num,
                                                                 active_learning_flag)
            else:
                selected_sample_loc = select_sample_within_group(experiment_conf['datapoint_selection'],
                                                                 my_df, selected_group, required_num,
                                                                 active_learning_flag)

            if (debug == True):
                print("SELECTED DATAPOINT = ", selected_sample_loc)

                # add it into the annotated data set
            my_df.loc[selected_sample_loc, 'selected_index'] = i
            # update its prediction with the label from the ground truth
            my_df.loc[selected_sample_loc, 'prediction'] = my_df.loc[selected_sample_loc, 'label']

            num_fetched_samples = len(selected_sample_loc)
            pbar.update(num_fetched_samples)

            j = j + num_fetched_samples

            i = i + num_fetched_samples
            if i > num_iteration:
                break

        end = time.time()
        # print("sample query time: ", end - start, "s")

        # select a subset of the lfs to produce the best prediction result
        best_fs = select_best_lfs(experiment_conf['lf_selection'], my_df, pool_lfs, epochs)
        if (debug == True):
            print("SELECTED set of LFs = " + str(best_fs))

        combined_lfs = combine_lf_sets(best_fs, new_lfs)
        if (debug == True):
            print("combined LFs = " + str(combined_lfs))

        start = time.time()

        # apply the ensembling algorithm to generate the predicted matches
        lfs_ensemble(experiment_conf['lf_ensemble'], my_df, combined_lfs, feature_set, balance, epochs)

        end = time.time()
        # print("LF ensemble time: ", end - start, "s")

        # record the evaluation performance of the ensemble model
        r = evaluate_result(my_df)

        measurements[i - 1] = r

        # plugin in the new lfs generated by the slow loop for the next iteration
        if experiment_conf['with_slow_loop'] == True and i < total_size:
            num_slow_loop = i // interval_slow_loop
            if num_slow_loop > iteration_slow_loop:
                # print("========= SLOW LOOP ==============")

                start = time.time()
                iteration_slow_loop = num_slow_loop
                new_lfs = trigger_slow_loop(my_df, feature_set, num_slow_loop, experiment_conf['with_lf_tuning'], tunable_lfs, batch_size)

                end = time.time()
                # print("LF augumentation time: ", end - start, "s")

    pbar.close()

    return measurements