import pandas as pd
import random
import sklearn.preprocessing as pp
from datetime import datetime
import itertools
import re
from scipy.stats import entropy
import uuid
import pickle
import os
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans
import sqlite3
from collections import Counter
import numpy as np
from sklearn.linear_model import SGDClassifier

pd.options.display.max_columns = 50
pd.options.display.max_colwidth = 200
pd.options.display.max_colwidth = 200
pd.set_option('display.max_rows', None)

RND_SEED = 45822
random.seed(RND_SEED)
np.random.seed(RND_SEED)

connection = sqlite3.connect('database.db', timeout=30)
with open("schema.sql") as f:
    connection.executescript(f.read())

def get_db_connection():
    conn = sqlite3.connect('database.db', timeout=30)
    conn.row_factory = sqlite3.Row
    return conn

def populate_texts_table_sql(texts_list, table_name="texts", reset_labels=True):
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM " + table_name + ";")
    conn.commit()
    if reset_labels:
        for text_record in texts_list:
            cur.execute("INSERT INTO " + table_name + " (id, text, label) VALUES (?, ?, ?)",
                        (text_record["id"], text_record["text"], "-"))
    else:
        for text_record in texts_list:
            cur.execute("INSERT INTO " + table_name + " (id, text, label) VALUES (?, ?, ?)",
                        (text_record["id"], text_record["text"], text_record["label"]))
    conn.commit()
    conn.close()
    return None

def get_decimal_value(name):
    conn = get_db_connection()
    query = "SELECT name, value FROM decimalValues WHERE name = '%s' ;" % name
    sql_table = conn.execute(query).fetchall()
    decimal_value = [dict(row)["value"] for row in sql_table][0]
    conn.close()
    return decimal_value

def set_decimal_value(name, value):
    conn = get_db_connection()
    query = "UPDATE decimalValues SET value = %s WHERE name = '%s' ;" % (value, name)
    conn.execute(query)
    conn.commit()
    conn.close()
    return None

def update_overall_quality_scores(value):
    current_score = get_decimal_value(name="OVERALL_QUALITY_SCORE_DECIMAL")
    set_decimal_value(name="OVERALL_QUALITY_SCORE_DECIMAL_PREVIOUS", value=current_score)
    set_decimal_value(name="OVERALL_QUALITY_SCORE_DECIMAL", value=value)
    return None

def set_pkl(name, pkl_data, reset=False):
    conn = get_db_connection()
    cur = conn.cursor()

    if not reset:
        test_query = cur.execute('SELECT * FROM pkls WHERE name = ?', (name,)).fetchall()
        if len(test_query) > 0:
            cur.execute('DELETE FROM pkls WHERE name = ?', (name,))

        query = "INSERT INTO pkls (name, data) VALUES (?, ?)"
        pkl_data_ = pickle.dumps(pkl_data)
        cur.execute(query, (name, pkl_data_))

        # test_query = cur.execute('SELECT * FROM pkls WHERE name = ?', (name,)).fetchall()
        # test_data = pickle.loads([dict(row)["data"] for row in test_query][0])
    else:
        cur.execute("DELETE FROM pkls WHERE name = '" + name + "';")

    conn.commit()
    conn.close()
    return None

def get_pkl(name):
    try:
        conn = get_db_connection()
        cur = conn.cursor()
        query = "SELECT * FROM pkls WHERE name = '" + name + "';"
        pkl_table = cur.execute(query).fetchall()

        data = [dict(row)["data"] for row in pkl_table]
        if len(data) > 0:
            pkl_data = pickle.loads(data[0])
        else:
            pkl_data = None

        conn.close()
        return pkl_data
    except:
        return None

def get_text_list(table_name="texts"):
    conn = get_db_connection()
    text_list_sql = conn.execute("SELECT * FROM " + table_name).fetchall()
    text_list_sql = [dict(row) for row in text_list_sql]
    conn.close()
    return text_list_sql

def set_text_list(label, table_name="searchResults"):
    conn = get_db_connection()
    conn.execute("UPDATE " + table_name + " SET label = '" + label + "'")
    conn.commit()
    conn.close()
    return None

def clear_text_list(table_name="searchResults"):
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM " + table_name + ";")
    conn.commit()
    conn.close()
    return None

def get_appropriate_text_list_list(text_list_full_sql, total_pages_sql, search_results_length_sql, table_limit_sql):
    if search_results_length_sql > 0:
        text_list_full_sql = get_text_list(table_name="searchResults")
        total_pages_sql = get_variable_value(name="SEARCH_TOTAL_PAGES")
    text_list_list_sql = create_text_list_list(text_list_full_sql=text_list_full_sql, sub_list_limit=table_limit_sql)
    return text_list_list_sql, text_list_full_sql, total_pages_sql

def get_y_classes():
    conn = get_db_connection()
    y_classes_sql = conn.execute('SELECT className FROM yClasses;').fetchall()
    y_classes_sql = [dict(row)["className"] for row in y_classes_sql]
    conn.close()
    return y_classes_sql

def clear_y_classes():
    conn = get_db_connection()
    conn.execute('DELETE FROM yClasses;')
    conn.commit()
    conn.close()
    return []

def add_y_classes(y_classses_list, begin_fresh=True):
    conn = get_db_connection()
    cur = conn.cursor()

    if begin_fresh:
        cur.execute("DELETE FROM yClasses;")

    for i, value in enumerate(y_classses_list):
        cur.execute("INSERT INTO yClasses (classId, className) VALUES (?, ?)", (i, value))

    conn.commit()
    conn.close()

    return 1

def get_click_log():
    conn = get_db_connection()
    sql_table = \
        conn.execute('SELECT click_id, click_location, click_type, click_object, click_date_time FROM clickRecord;').fetchall()

    click_log_sql = list()
    for row in sql_table:
        dict_row = {"click_id": dict(row)["click_id"], "click_location": dict(row)["click_location"],
                    "click_type": dict(row)["click_type"], "click_object": dict(row)["click_object"],
                    "click_date_time" : dict(row)["click_date_time"]}
        click_log_sql.append(dict_row)
    conn.close()
    return click_log_sql

def get_value_log():
    conn = get_db_connection()
    sql_table = \
        conn.execute('SELECT click_id, value_type, value FROM valueRecord;').fetchall()

    click_log_sql = list()
    for row in sql_table:
        dict_row = {"click_id": dict(row)["click_id"], "value_type": dict(row)["value_type"],
                    "value": dict(row)["value"]}
        click_log_sql.append(dict_row)
    conn.close()
    return click_log_sql

def reset_log_click_record_sql():
    conn = get_db_connection()
    conn.execute("DELETE FROM clickRecord")
    conn.commit()
    conn.close()
    return None

def reset_log_click_value_sql():
    conn = get_db_connection()
    conn.execute("DELETE FROM valueRecord")
    conn.commit()
    conn.close()
    return None

def add_log_click_record_sql(records):
    conn = get_db_connection()
    cur = conn.cursor()

    for record in records:
        cur.execute("""INSERT INTO clickRecord (click_id, click_location, click_type, click_object, click_date_time)
                        VALUES (?, ?, ?, ?, ?)""", (record["click_id"], record["click_location"], record["click_type"],
                                                    record["click_object"], record["click_date_time"]))

    conn.commit()
    conn.close()

    return None

def add_log_click_value_sql(records):
    conn = get_db_connection()
    cur = conn.cursor()

    for record in records:
        cur.execute("""INSERT INTO valueRecord (click_id, value_type, value)
                        VALUES (?, ?, ?)""", (record["click_id"], record["value_type"], record["value"]))

    conn.commit()
    conn.close()

    return None

def get_panel_flags():
    conn = get_db_connection()
    sql_table = conn.execute('SELECT name, value FROM initializeFlags;').fetchall()
    panel_flags = {dict(row)["name"]: dict(row)["value"] for row in sql_table}
    conn.close()
    return panel_flags

def update_panel_flags_sql(update_flag):
    conn = get_db_connection()
    cur = conn.cursor()
    update_query = "UPDATE initializeFlags SET value = ? WHERE name = ?;"

    for name, value in update_flag.items():
        cur.execute(update_query, (value, name))
    conn.commit()
    conn.close()
    return None

def get_texts_group_x(table_name="group1Texts"):
    conn = get_db_connection()
    sql_table = conn.execute("SELECT id, text, label FROM " + table_name + ";").fetchall()
    conn.close()
    if len(sql_table) > 0:
        texts_group_2 = [{"id": dict(row)["id"], "text": dict(row)["text"], "label": dict(row)["label"]} for row in sql_table]
    else:
        texts_group_2 = []
    return texts_group_2

def set_texts_group_x(top_texts, table_name="group1Texts"):
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM " + table_name + ";")
    if top_texts:
        for record in top_texts:
            cur.execute("INSERT INTO " + table_name + " (id, text, label) VALUES (?, ?, ?)",
                        (record["id"], record["text"], record["label"]))
    conn.commit()
    conn.close()
    return None

def get_total_summary_sql():
    conn = get_db_connection()
    sql_table = conn.execute('SELECT name, number, percentage FROM totalSummary;').fetchall()
    total_summary = [{"name": dict(row)["name"],
                      "number": dict(row)["number"],
                      "percentage": dict(row)["percentage"]} for row in sql_table]
    conn.close()
    return total_summary

def set_total_summary(text_lists):
    labels = [text_obj["label"] for text_obj in text_lists]
    label_counter = Counter(labels)
    total_texts = len(text_lists)
    number_unlabeled = label_counter["-"]
    number_labeled = total_texts - number_unlabeled
    total_texts_percentage = "100.00%"

    if total_texts > 0:
        number_unlabeled_percentage = "{:.2%}".format(number_unlabeled / total_texts)
        number_labeled_percentage = "{:.2%}".format(number_labeled / total_texts)
    else:
        number_unlabeled_percentage = "{:.2%}".format(1.0)
        number_labeled_percentage = "{:.2%}".format(0.0)

    total_summary = list()
    total_summary.append({"name": "Total Texts",
                          "number": "{:,}".format(total_texts),
                          "percentage": total_texts_percentage})
    total_summary.append({"name": "Total Unlabeled",
                          "number": "{:,}".format(number_unlabeled),
                          "percentage": number_unlabeled_percentage})
    total_summary.append({"name": "Total Labeled",
                          "number": "{:,}".format(number_labeled),
                          "percentage": number_labeled_percentage})

    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM totalSummary;")
    for record in total_summary:
        cur.execute("INSERT INTO totalSummary (name, number, percentage) VALUES (?, ?, ?)",
                    (record["name"], record["number"], record["percentage"]))

    conn.commit()
    conn.close()
    return None

def get_label_summary_sql():
    conn = get_db_connection()
    sql_table = conn.execute('SELECT name, number, percentage FROM labelSummary;').fetchall()
    label_summary = [{"name": dict(row)["name"],
                      "number": dict(row)["number"],
                      "percentage": dict(row)["percentage"]} for row in sql_table]
    conn.close()
    return label_summary

def set_label_summary(text_lists):
    labels = [text_obj["label"] for text_obj in text_lists]
    label_counter = Counter(labels)
    total_texts = len(text_lists)

    label_summary = []
    for key, value in label_counter.items():
        label_summary.append({"name": key,
                              "number": "{:,}".format(value),
                              "percentage": "{:.2%}".format(value / total_texts)})

    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM labelSummary;")
    for record in label_summary:
        cur.execute("INSERT INTO labelSummary (name, number, percentage) VALUES (?, ?, ?)",
                    (record["name"], record["number"], record["percentage"]))
    conn.commit()
    conn.close()
    return None

def get_selected_text(selected_text_id, text_list_full_sql):
    selected_text_test = [text["text"] for text in text_list_full_sql if text["id"] == selected_text_id]
    if selected_text_id:
        if len(selected_text_test) == 0:
            selected_text = ""
        else:
            selected_text = selected_text_test[0]
    else:
        selected_text = ""
    return selected_text

def create_text_list_list(text_list_full_sql, sub_list_limit):
    texts_list_list = \
        [text_list_full_sql[i:i + sub_list_limit] for i in range(0, len(text_list_full_sql), sub_list_limit)]
    return texts_list_list

def update_texts_list_by_id_sql(update_objs=None, selected_label=None, update_ids=None, sub_list_limit=10,
                                update_in_place=True):
    conn = get_db_connection()
    cur = conn.cursor()

    if selected_label and update_ids and not update_objs:
        if update_in_place:
            update_query = "UPDATE texts SET label = ? WHERE id IN (%s)" % ",".join("?"*len(update_ids))
            update_values = [selected_label]
            update_values.extend(update_ids)
            cur.execute(update_query, update_values)
            conn.commit()
            conn.close()

        else:
            cur.execute("DROP TABLE IF EXISTS temp_table;")
            cur.execute("""
                        CREATE TABLE temp_table (
                            id TEXT NOT NULL,
                            text TEXT NOT NULL,
                            label TEXT NOT NULL
                        );
                        """)
            query = "INSERT INTO temp_table SELECT * FROM texts WHERE id IN (%s)" % ",".join("?" * len(update_ids))
            cur.execute(query, update_ids)
            cur.execute("UPDATE temp_table SET label = ?", (selected_label, ))
            cur.execute("DELETE FROM texts WHERE id IN (%s)" % ",".join("?" * len(update_ids)), update_ids)
            cur.execute("INSERT INTO texts SELECT * FROM temp_table;")
            conn.commit()
            conn.close()

    elif update_objs and not selected_label and not update_ids:
        if update_in_place:
            labels = set([obj["label"] for obj in update_objs])
            for label in labels:
                update_ids = [obj["id"] for obj in update_objs if obj["label"] == label]
                update_ids_sql = ", ".join(update_ids)
                update_query = "UPDATE texts SET label = ? WHERE id IN (%s)" % update_ids_sql

                conn.execute(update_query, (label, ))
            conn.commit()
            conn.close()
        else:
            cur.execute("DROP TABLE IF EXISTS temp_table;")
            cur.execute("""
                           CREATE TABLE temp_table (
                               id TEXT NOT NULL,
                               text TEXT NOT NULL,
                               label TEXT NOT NULL
                           );
                           """)
            all_update_ids = [obj["id"] for obj in update_objs]
            query = "INSERT INTO temp_table SELECT * FROM texts WHERE id IN (%s)" % ",".join("?" * len(all_update_ids))
            cur.execute(query, all_update_ids)
            labels = set([obj["label"] for obj in update_objs])
            for label in labels:
                update_ids = [obj["id"] for obj in update_objs if obj["label"] == label]
                update_ids_sql = ", ".join(update_ids)
                update_query = "UPDATE temp_table SET label = ? WHERE id IN (%s)" % update_ids_sql
                conn.execute(update_query, (label,))
            delete_query = "DELETE FROM texts WHERE id IN (%s)" % ",".join("?" * len(all_update_ids))
            cur.execute(delete_query, all_update_ids)
            cur.execute("INSERT INTO texts SELECT * FROM temp_table;")

            conn.commit()
            conn.close()

    text_list_full = get_text_list(table_name="texts")
    texts_list_list = create_text_list_list(text_list_full_sql=text_list_full, sub_list_limit=sub_list_limit)

    return text_list_full, texts_list_list

def label_all_sql(fitted_classifier, sparse_vectorized_corpus, corpus_text_ids, texts_list,
                  label_only_unlabeled=True, sub_list_limit=50, update_in_place=True):

    texts_list_df = pd.DataFrame(texts_list)

    if not label_only_unlabeled:
        predictions = fitted_classifier.predict(sparse_vectorized_corpus)
        predictions_df = pd.DataFrame(predictions)
        predictions_df["id"] = corpus_text_ids
        labeled_text_ids = corpus_text_ids
        number_to_label = len(labeled_text_ids)
    else:
        label_only_these_ids = texts_list_df[texts_list_df["label"] == "-"]["id"].values
        keep_indices = [corpus_text_ids.index(x) for x in label_only_these_ids]
        number_to_label = len(keep_indices)

    if number_to_label > 0:
        if label_only_unlabeled:
            sparse_vectorized_corpus_alt = sparse_vectorized_corpus[keep_indices, :]
            predictions = fitted_classifier.predict(sparse_vectorized_corpus_alt)
            predictions_df = pd.DataFrame(predictions)
            predictions_df["id"] = label_only_these_ids
            labeled_text_ids = label_only_these_ids

        predictions_df = predictions_df.rename(columns={0: "label"})
        predictions_df = predictions_df.merge(texts_list_df[["id", "text"]], left_on="id", right_on="id",
                                              how="left")
        predictions_df = predictions_df[["id", "text", "label"]]
        update_objects = predictions_df.to_dict("records")

        text_list_full, texts_list_list = \
            update_texts_list_by_id_sql(update_objs=update_objects,
                                        selected_label=None,
                                        update_ids=None,
                                        sub_list_limit=sub_list_limit,
                                        update_in_place=update_in_place)
    else:
        text_list_full = get_text_list(table_name="texts")
        texts_list_list = create_text_list_list(text_list_full_sql=text_list_full, sub_list_limit=sub_list_limit)
        labeled_text_ids = []

    return text_list_full, texts_list_list, labeled_text_ids

def generate_summary_sql(text_lists):
    labels = [text_obj["label"] for text_obj in text_lists]
    label_counter = Counter(labels)
    total_texts = len(text_lists)
    number_unlabeled = label_counter["-"]
    number_labeled = total_texts - number_unlabeled

    set_total_summary(text_lists=text_lists)
    set_label_summary(text_lists=text_lists)
    set_variable(name="NUMBER_UNLABELED_TEXTS", value=number_unlabeled)

    summary_headline = \
        "Total Labeled : {:,} / {:,} {:.1%}".format(number_labeled, total_texts, number_labeled / total_texts)
    set_variable(name="LABEL_SUMMARY_STRING", value=summary_headline)

    total_summary_sql = get_total_summary_sql()
    label_summary_sql = get_label_summary_sql()
    number_unlabeled_texts_sql = number_unlabeled
    label_summary_string_sql = summary_headline
    return total_summary_sql, label_summary_sql, number_unlabeled_texts_sql, label_summary_string_sql

def set_variable(name, value):
    conn = get_db_connection()
    cur = conn.cursor()

    test_query = cur.execute('SELECT * FROM variables WHERE name = ?', (name,)).fetchall()
    if len(test_query) > 0:
        cur.execute('DELETE FROM variables WHERE name = ?', (name,))
        query = """INSERT INTO variables (name, value) VALUES (?, ?)
                    """
    else:
        query = """INSERT INTO variables (name, value) VALUES (?, ?)
                    """
    cur.execute(query, (name, value))
    conn.commit()
    conn.close()
    return 1

def get_variable_value(name):
    conn = get_db_connection()
    cur = conn.cursor()
    query = cur.execute('SELECT value FROM variables WHERE name = ?', (name,)).fetchall()
    value = [dict(row)["value"] for row in query]
    value = value[0]

    if name in ["TOTAL_PAGES", "NUMBER_UNLABELED_TEXTS", "MAX_CONTENT_PATH", "TEXTS_LIMIT", "TABLE_LIMIT",
                "MAX_FEATURES", "RND_STATE", "PREDICTIONS_NUMBER", "SEARCH_RESULTS_LENGTH", "GROUP_1_KEEP_TOP",
                "GROUP_3_KEEP_TOP", "CONFIRM_LABEL_ALL_TEXTS_COUNTS", "SEARCH_TOTAL_PAGES", "LABEL_ALL_BATCH_NO",
                "LABEL_ALL_TOTAL_BATCHES", "NUMBER_AUTO_LABELED", "LABEL_ALL_BATCH_SIZE"]:
        value = int(value)

    if name in ["KEEP_ORIGINAL", "GROUP_1_EXCLUDE_ALREADY_LABELED", "GROUP_2_EXCLUDE_ALREADY_LABELED",
                "PREDICTIONS_VERBOSE", "SIMILAR_TEXT_VERBOSE", "FIT_CLASSIFIER_VERBOSE", "FIRST_LABELING_FLAG",
                "FULL_FIT_IF_LABELS_GOT_OVERRIDDEN", "FORCE_FULL_FIT_FOR_DIFFICULT_TEXTS",
                "LABELS_GOT_OVERRIDDEN_FLAG", "UPDATE_TEXTS_IN_PLACE"]:
        if value == "True":
            value = True
        else:
            value = False

    if name in ["PREDICTIONS_PROBABILITY"]:
        value = float(value)

    conn.commit()
    conn.close()
    return value

def get_difficult_texts_sql():
    try:
        conn = get_db_connection()
        select_diff_texts_query = get_variable_value(name="SELECT_DIFF_TEXTS_QUERY")
        sql_cols_list_y_classes = get_pkl(name="SQL_COLS_LIST_Y_CLASSES")
        sql_table = conn.execute(select_diff_texts_query).fetchall()

        total_summary = list()
        for row in sql_table:
            temp_row = {col: dict(row)[col] for col in sql_cols_list_y_classes}
            total_summary.append(temp_row)

        conn.close()
        return total_summary
    except:
        return []

def reset_difficult_texts_sql():
    try:
        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute('DELETE FROM difficultTexts')
        conn.commit()
        conn.close()
        return None
    except:
        return None

def get_available_datasets():
    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute("DELETE FROM availableDatasets;")
    cur.execute("INSERT INTO availableDatasets SELECT * FROM fixedDatasets;")
    conn.commit()
    conn.close()
    dataset_name, dataset_url, date_time, y_classes, total_summary = has_save_data()

    if dataset_name and date_time and y_classes and total_summary:
        date_at_end_check = re.findall(r"(.*)\-[0-9]{4}\-[0-9]{2}\-[0-9]{2}\-\-[0-9]{2}\-[0-9]{2}\-[0-9]{2}", dataset_name)
        if len(date_at_end_check) > 0:
            dataset_name_alt = date_at_end_check[0]
        else:
            dataset_name_alt = dataset_name

        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute("INSERT INTO availableDatasets (name, description, url) VALUES (?, ?, ?)",
                    (dataset_name_alt + "-" + date_time,
                     "A partially labeled dataset having " + total_summary[2]["percentage"] +
                     " of " + total_summary[0]["number"] + " texts labeled.",
                     dataset_url))
        conn.commit()
        conn.close()

    conn = get_db_connection()
    available_datasets_sql = conn.execute('SELECT * FROM availableDatasets').fetchall()
    conn.close()
    return available_datasets_sql

def has_save_data():
    try:
        dataset_name = get_pkl(name="DATASET_NAME")
        dataset_url = get_pkl(name="DATASET_URL")
        date_time = get_pkl(name="DATE_TIME")
        y_classes = get_pkl(name="Y_CLASSES")
        total_summary = get_pkl(name="TOTAL_SUMMARY")
        return dataset_name, dataset_url, date_time, y_classes, total_summary
    except:
        return None, None, None, None, None

def get_all_predictions_sql(fitted_classifier, sparse_vectorized_corpus, corpus_text_ids, texts_list,
                            top=5,
                            y_classes=["earthquake", "fire", "flood", "hurricane"],
                            verbose=False,
                            round_to=2,
                            format_as_percentage=False):

    predictions = fitted_classifier.predict_proba(sparse_vectorized_corpus)

    predictions_df = pd.DataFrame(predictions)
    y_classes = [x.replace(" ", "_") for x in y_classes]
    predictions_df.columns = y_classes

    predictions_summary = predictions_df.replace(0.0, np.NaN).mean(axis=0)

    predictions_df["id"] = corpus_text_ids

    texts_list_df = pd.DataFrame(texts_list)
    predictions_df = predictions_df.merge(texts_list_df, left_on="id", right_on="id")

    keep_cols = ["id", "text"]
    keep_cols.extend(y_classes)
    predictions_df = predictions_df[keep_cols]

    pred_scores = score_predictions(predictions_df[y_classes], use_entropy=True, num_labels=len(y_classes))
    overall_quality = np.mean(pred_scores)
    overall_quality_score_decimal_sql = overall_quality
    predictions_df["pred_scores"] = pred_scores

    if round_to and not format_as_percentage:
        predictions_df[y_classes] = predictions_df[y_classes].round(round_to)
        predictions_summary = predictions_summary.round(round_to)
        overall_quality = overall_quality.round(round_to)

    if format_as_percentage:
        if verbose:
            print(">> get_all_predictions > predictions_df.head() :")
            print(predictions_df.head(top))
        predictions_df[y_classes] = predictions_df[y_classes]\
            .astype(float)\
            .applymap(lambda x: "{0:.0%}".format(x))

        # predictions_summary = (predictions_summary.astype(float) * 100).round(1).astype(str) + "%"
        overall_quality = (overall_quality.astype(float) * 100).round(1).astype(str) + "%"

    predictions_df = predictions_df.sort_values(["pred_scores"], ascending=[True])

    if verbose:
        print(">> get_all_predictions > predictions_df.head() :")
        print(predictions_df.head(top))
        print(">> get_all_predictions > predictions_df.tail() :")
        print(predictions_df.tail(top))

    keep_cols = ["id", "text"]
    keep_cols.extend(y_classes)
    sql_cols_list = [x + ' TEXT NOT NULL' for x in keep_cols]
    sql_cols = ", ".join(sql_cols_list)
    top_texts = predictions_df.head(top)[keep_cols].to_dict("records")

    sql_query_1 = """
                    DROP TABLE IF EXISTS difficultTexts;
                """
    sql_query_2 = """
                    CREATE TABLE difficultTexts (
                """ + sql_cols + """
                    );
                """

    conn = get_db_connection()
    cur = conn.cursor()
    cur.execute(sql_query_1)
    conn.commit()
    cur.execute(sql_query_2)
    conn.commit()

    parameters = ", ".join(["?"] * len(keep_cols))
    query = "INSERT INTO difficultTexts (" + ", ".join(keep_cols) + ") VALUES (%s)" % parameters

    for record in top_texts:
        insert_values = [value for key, value in record.items()]
        cur.execute(query, (insert_values))
    conn.commit()
    conn.close()

    conn = get_db_connection()
    select_diff_texts_query = "SELECT " + ", ".join(keep_cols) + " FROM difficultTexts;"
    set_variable(name="SELECT_DIFF_TEXTS_QUERY", value=select_diff_texts_query)
    set_pkl(name="SQL_COLS_LIST_Y_CLASSES", pkl_data=keep_cols, reset=None)
    sql_table = conn.execute(select_diff_texts_query).fetchall()
    texts_group_3_sql = []
    for row in sql_table:
        texts_group_3_sql.append({key: value for key, value in dict(row).items()})
    conn.close()

    update_overall_quality_scores(value=overall_quality_score_decimal_sql)
    set_variable(name="OVERALL_QUALITY_SCORE", value=overall_quality)
    overall_quality_score_sql = overall_quality
    overall_quality_score_decimal_previous_sql = get_decimal_value(name="OVERALL_QUALITY_SCORE_DECIMAL_PREVIOUS")
    return texts_group_3_sql, overall_quality_score_sql, \
           overall_quality_score_decimal_sql, overall_quality_score_decimal_previous_sql

def get_top_predictions_sql(selected_class, fitted_classifier, sparse_vectorized_corpus, corpus_text_ids,
                            texts_list,
                            top=5,
                            cutoff_proba=0.95,
                            y_classes=["earthquake", "fire", "flood", "hurricane"],
                            verbose=False,
                            exclude_already_labeled=True):

    predictions = fitted_classifier.predict_proba(sparse_vectorized_corpus)

    predictions_df = pd.DataFrame(predictions)
    predictions_df.columns = y_classes
    predictions_df["id"] = corpus_text_ids

    keep_cols = ["id"]
    keep_cols.extend([selected_class])
    predictions_df = predictions_df[keep_cols]

    predictions_df = predictions_df[predictions_df[selected_class] > cutoff_proba]
    predictions_df = predictions_df.sort_values([selected_class], ascending=False)

    if exclude_already_labeled:
        texts_list_df = pd.DataFrame.from_dict(texts_list)
        predictions_df = predictions_df.merge(texts_list_df, left_on="id", right_on="id", how="left")
        predictions_df = predictions_df[predictions_df["label"].isin(["-"])]

    if verbose:
        print(">> get_top_predictions > predictions_df :")
        print(predictions_df.head(top))

    filter_list = predictions_df.head(top)["id"].values
    top_texts = filter_all_texts(texts_list, filter_list, exclude_already_labeled=False)

    set_texts_group_x(top_texts=top_texts, table_name="group2Texts")
    texts_group_3_sql = get_texts_group_x(table_name="group2Texts")

    return texts_group_3_sql

def fit_classifier_sql(sparse_vectorized_corpus, corpus_text_ids, texts_list, texts_list_labeled,
                       y_classes=["earthquake", "fire", "flood", "hurricane"],
                       verbose=False,
                       random_state=2584,
                       n_jobs=-1,
                       labels_got_overridden_flag=False,
                       full_fit_if_labels_got_overridden=False):
    texts_list_labeled_df = pd.DataFrame.from_dict(texts_list_labeled)

    if verbose:
        print("texts_list_labeled_df :")
        print(texts_list_labeled_df.head())

    ids = texts_list_labeled_df["id"].values
    y_train = texts_list_labeled_df["label"].values
    indices = [corpus_text_ids.index(x) for x in ids]
    X_train = sparse_vectorized_corpus[indices, :]

    classifier_sql = get_pkl(name="CLASSIFIER")
    if classifier_sql:
        clf = classifier_sql
    else:
        # clf = make_pipeline(StandardScaler(), SGDClassifier(max_iter=1000, tol=1e-3, random_state=2584))
        clf = SGDClassifier(loss="modified_huber", max_iter=1000, tol=1e-3, random_state=random_state, n_jobs=n_jobs)

    if labels_got_overridden_flag:
        if full_fit_if_labels_got_overridden:
            all_texts_list_labeled_df = pd.DataFrame.from_dict(texts_list)
            all_texts_list_labeled_df = all_texts_list_labeled_df[~all_texts_list_labeled_df["label"].isin(["-"])]

            y_classes_labeled = list(set(all_texts_list_labeled_df["label"].values))
            all_classes_present = all(label in y_classes_labeled for label in y_classes)
            clf = SGDClassifier(loss="modified_huber", max_iter=1000, tol=1e-3, random_state=random_state,
                                n_jobs=n_jobs)

            ids_all = all_texts_list_labeled_df["id"].values
            y_train_all = all_texts_list_labeled_df["label"].values
            indices_all = [corpus_text_ids.index(x) for x in ids_all]
            X_train_all = sparse_vectorized_corpus[indices_all, :]

            if all_classes_present:
                clf.fit(X_train_all, y_train_all)
            else:
                clf.partial_fit(X_train_all, y_train_all, classes=y_classes)
        else:
            clf.partial_fit(X_train, y_train, classes=y_classes)
    else:
        clf.partial_fit(X_train, y_train, classes=y_classes)

    set_pkl(name="CLASSIFIER", pkl_data=clf, reset=False)
    return clf

def load_new_data_sql(source_file,
                      text_id_col,
                      text_value_col,
                      source_folder="./output/upload/",
                      shuffle_by="kmeans",
                      table_limit=50, texts_limit=1000, max_features=100,
                      y_classes=["Label 1", "Label 2", "Label 3", "Label 4"], rnd_state=258):

    data_df = get_new_data(source_file=source_file,
                           source_folder=source_folder,
                           number_samples=None,
                           random_state=rnd_state)

    corpus_text_ids = [str(x) for x in data_df[text_id_col].values]

    vectorizer = TfidfVectorizer(ngram_range=(1, 2), stop_words="english", max_features=max_features)
    vectorized_corpus = vectorizer.fit_transform(data_df[text_value_col].values)

    if shuffle_by == "kmeans":
        kmeans = KMeans(n_clusters=len(y_classes), random_state=rnd_state).fit(vectorized_corpus)
        kmeans_labels = kmeans.labels_
        texts_list, adj_text_ids = convert_new_data_into_list_json(data_df,
                                                                   limit=texts_limit,
                                                                   shuffle_list=kmeans_labels,
                                                                   random_shuffle=False,
                                                                   random_state=rnd_state,
                                                                   id_col=text_id_col,
                                                                   text_col=text_value_col,
                                                                   label_col="label")
    else:
        texts_list, adj_text_ids = convert_new_data_into_list_json(data_df,
                                                                   limit=texts_limit,
                                                                   shuffle_list=[],
                                                                   random_shuffle=True,
                                                                   random_state=rnd_state,
                                                                   id_col=text_id_col,
                                                                   text_col=text_value_col,
                                                                   label_col="label")
    populate_texts_table_sql(texts_list=texts_list, table_name="texts")

    set_pkl(name="DATASET_NAME", pkl_data=None, reset=True)
    set_pkl(name="DATASET_NAME", pkl_data=source_file, reset=False)

    set_pkl(name="DATASET_URL", pkl_data=None, reset=True)
    set_pkl(name="DATASET_URL", pkl_data="-", reset=False)

    set_pkl(name="CORPUS_TEXT_IDS", pkl_data=None, reset=True)
    set_pkl(name="CORPUS_TEXT_IDS", pkl_data=adj_text_ids, reset=False)

    texts_list_list = [texts_list[i:i + table_limit] for i in range(0, len(texts_list), table_limit)]
    total_pages = len(texts_list_list)
    set_variable(name="TOTAL_PAGES", value=total_pages)

    set_pkl(name="TEXTS_LIST_LIST", pkl_data=None, reset=True)
    set_pkl(name="TEXTS_LIST_LIST", pkl_data=texts_list_list, reset=False)

    set_pkl(name="VECTORIZED_CORPUS", pkl_data=None, reset=True)
    set_pkl(name="VECTORIZED_CORPUS", pkl_data=vectorized_corpus, reset=False)

    set_pkl(name="VECTORIZER", pkl_data=None, reset=True)
    set_pkl(name="VECTORIZER", pkl_data=vectorizer, reset=False)

    return texts_list, texts_list_list, adj_text_ids, total_pages, vectorized_corpus, vectorizer, corpus_text_ids

def load_demo_data_sql(dataset_name="Disaster Tweets Dataset", shuffle_by="kmeans",
                       table_limit=50, texts_limit=1000, max_features=100,
                       y_classes=["Earthquake", "Fire", "Flood", "Hurricane"], rnd_state=258):
    if dataset_name == "Disaster Tweets Dataset":
        consolidated_disaster_tweet_data_df = get_disaster_tweet_demo_data(number_samples=None,
                                                                           filter_data_types=["train"],
                                                                           random_state=rnd_state)
        corpus_text_ids = [str(x) for x in consolidated_disaster_tweet_data_df["tweet_id"].values]
        set_pkl(name="CORPUS_TEXT_IDS", pkl_data=corpus_text_ids, reset=False)

        vectorizer = TfidfVectorizer(ngram_range=(1, 2), stop_words="english", max_features=max_features)

        # https://stackoverflow.com/questions/69326639/sklearn-warnings-in-version-1-0
        vectorized_corpus = \
            vectorizer.fit_transform(consolidated_disaster_tweet_data_df["tweet_text"].values)

        set_pkl(name="VECTORIZER", pkl_data=vectorizer, reset=False)
        set_pkl(name="VECTORIZED_CORPUS", pkl_data=vectorized_corpus, reset=False)

        if shuffle_by == "kmeans":
            kmeans = KMeans(n_clusters=len(y_classes), random_state=rnd_state).fit(vectorized_corpus)
            kmeans_labels = kmeans.labels_
            texts_list, adj_text_ids = convert_demo_data_into_list_json(consolidated_disaster_tweet_data_df,
                                                                        limit=texts_limit,
                                                                        keep_labels=False,
                                                                        shuffle_list=kmeans_labels,
                                                                        random_shuffle=False,
                                                                        random_state=rnd_state)
        else:
            texts_list, adj_text_ids = convert_demo_data_into_list_json(consolidated_disaster_tweet_data_df,
                                                                        limit=texts_limit,
                                                                        keep_labels=False,
                                                                        shuffle_list=[],
                                                                        random_shuffle=True,
                                                                        random_state=rnd_state)

    texts_list_list = [texts_list[i:i + table_limit] for i in range(0, len(texts_list), table_limit)]

    total_pages = len(texts_list_list)
    set_variable(name="TOTAL_PAGES", value=total_pages)

    return texts_list, texts_list_list, adj_text_ids, total_pages, vectorized_corpus, vectorizer, corpus_text_ids

def generate_all_predictions_if_appropriate(n_jobs=-1,
                                            labels_got_overridden_flag=True,
                                            full_fit_if_labels_got_overridden=True,
                                            round_to=1,
                                            format_as_percentage=True):

    classifier_sql = get_pkl(name="CLASSIFIER")

    try:
        if classifier_sql:
            label_summary_sql = get_label_summary_sql()
            label_summary_df = pd.DataFrame.from_dict(label_summary_sql)
            y_classes_labeled = label_summary_df["name"].values
            y_classes_sql = get_y_classes()
            all_classes_present = all(label in y_classes_labeled for label in y_classes_sql)
            if all_classes_present:
                force_full_fit_for_difficult_texts_sql = get_variable_value(name="FORCE_FULL_FIT_FOR_DIFFICULT_TEXTS")
                random_state_sql = get_variable_value(name="RND_STATE")
                corpus_text_ids_sql = get_pkl(name="CORPUS_TEXT_IDS")
                vectorized_corpus_sql = get_pkl(name="VECTORIZED_CORPUS")
                predictions_verbose_sql = get_variable_value(name="PREDICTIONS_VERBOSE")
                fit_classifier_verbose_sql = get_variable_value(name="FIT_CLASSIFIER_VERBOSE")
                text_list_full_sql = get_text_list(table_name="texts")

                if force_full_fit_for_difficult_texts_sql:
                    fit_classifier_sql(sparse_vectorized_corpus=vectorized_corpus_sql,
                                       corpus_text_ids=corpus_text_ids_sql,
                                       texts_list=text_list_full_sql,
                                       texts_list_labeled=text_list_full_sql,
                                       y_classes=y_classes_sql,
                                       verbose=fit_classifier_verbose_sql,
                                       random_state=random_state_sql,
                                       n_jobs=n_jobs,
                                       labels_got_overridden_flag=labels_got_overridden_flag,
                                       full_fit_if_labels_got_overridden=full_fit_if_labels_got_overridden)

                top_sql = get_variable_value(name="GROUP_3_KEEP_TOP")
                texts_group_3_sql, overall_quality_score_sql, \
                    overall_quality_score_decimal_sql, overall_quality_score_decimal_previous_sql = \
                        get_all_predictions_sql(fitted_classifier=classifier_sql,
                                                sparse_vectorized_corpus=vectorized_corpus_sql,
                                                corpus_text_ids=corpus_text_ids_sql,
                                                texts_list=text_list_full_sql,
                                                top=top_sql,
                                                y_classes=y_classes_sql,
                                                verbose=predictions_verbose_sql,
                                                round_to=round_to,
                                                format_as_percentage=format_as_percentage)
                return 1, "The difficult texts list has been generated.", \
                        texts_group_3_sql, overall_quality_score_sql, \
                            overall_quality_score_decimal_sql, overall_quality_score_decimal_previous_sql
            else:
                return 0, """Examples of all labels are not present.
                                  Label more texts then try generating the difficult text list.""", [], "-", 0.0, 0.0
        else:
                return 0, "Label more texts then try generating the difficult text list.", [], "-", 0.0, 0.0

    except:
        return -1, "An error occurred when trying to generate the difficult texts.", [], "-", 0.0, 0.0

def get_top_similar_texts_sql(all_texts_json, similarities_series, top=5, exclude_already_labeled=False, verbose=True):
    if exclude_already_labeled:
        all_texts_df = pd.DataFrame.from_dict(all_texts_json)
        similarities_df = pd.DataFrame(similarities_series).reset_index().rename(columns={0: "proba", "index": "id"})

        if verbose:
            print(">> get_top_similar_texts > similarities_df :")
            print(similarities_df.head())

            print(">> all_texts_df > all_texts_df :")
            print(all_texts_df.head())

        similarities_df = similarities_df.merge(all_texts_df, left_on="id", right_on="id")
        similarities_df = similarities_df[similarities_df["label"].isin(["-"])]

        filter_list = similarities_df.head(top)["id"].values
    else:
        filter_list = similarities_series.head(top).index.values

    top_texts = filter_all_texts(all_texts_json, filter_list, exclude_already_labeled=False)

    if verbose:
        print(">> all_texts_df > len(top_texts) :", len(top_texts))

    set_texts_group_x(top_texts=top_texts, table_name="group1Texts")

    return top_texts

def read_new_dataset(source_file_name, text_id_col, text_value_col, source_dir="./output/upload"):
    try:
        dataset_df = pd.read_csv(os.path.join(source_dir, source_file_name))
        dataset_df = dataset_df[[text_id_col, text_value_col]]
        dataset_df.rename(columns={text_id_col: "id", text_value_col: "text"}, inplace=True)
        dataset_df["label"] = "-"

        return 1, dataset_df
    except:
        return 0, None

def load_save_state_sql():
    try:
        save_state_json = get_pkl(name="SAVE_STATE")

        set_variable(name="TOTAL_PAGES", value=save_state_json["TOTAL_PAGES"])
        add_y_classes(y_classses_list=save_state_json["Y_CLASSES"], begin_fresh=True)
        set_variable(name="SHUFFLE_BY", value=save_state_json["SHUFFLE_BY"])
        set_variable(name="HTML_CONFIG_TEMPLATE", value=save_state_json["HTML_CONFIG_TEMPLATE"])
        set_variable(name="DATASET_NAME", value=save_state_json["DATASET_NAME"])

        set_variable(name="LABEL_ALL_BATCH_NO", value=-99)
        set_variable(name="LABEL_ALL_TOTAL_BATCHES", value=0)
        set_variable(name="NUMBER_AUTO_LABELED", value=0)

        set_variable(name="SEARCH_MESSAGE", value=save_state_json["SEARCH_MESSAGE"])
        set_variable(name="NUMBER_UNLABELED_TEXTS", value=save_state_json["NUMBER_UNLABELED_TEXTS"])
        set_variable(name="LABEL_SUMMARY_STRING", value=save_state_json["LABEL_SUMMARY_STRING"])
        set_variable(name="OVERALL_QUALITY_SCORE", value=save_state_json["OVERALL_QUALITY_SCORE"])
        set_variable(name="OVERALL_QUALITY_SCORE_DECIMAL", value=save_state_json["OVERALL_QUALITY_SCORE_DECIMAL"])
        set_variable(name="OVERALL_QUALITY_SCORE_DECIMAL_PREVIOUS",
                     value=save_state_json["OVERALL_QUALITY_SCORE_DECIMAL_PREVIOUS"])
        set_variable(name="ALLOW_SEARCH_TO_OVERRIDE_EXISTING_LABELS",
                     value=save_state_json["ALLOW_SEARCH_TO_OVERRIDE_EXISTING_LABELS"])
        texts_list_sql = get_pkl(name="TEXTS_LIST")
        populate_texts_table_sql(texts_list=texts_list_sql, table_name="texts", reset_labels=False)
        generate_summary_sql(text_lists=texts_list_sql)
        return 1
    except:
        return 0

def get_disaster_tweet_demo_data(number_samples=None,
                                 filter_data_types=None,
                                 random_state=1144,
                                 source_file="./data/consolidated_disaster_tweet_data.tsv",
                                 retry_count=5):

    data_df = pd.read_csv(source_file, sep="\t")
    for count in range(0, retry_count):
        if len(data_df) == 0:
            data_df = pd.read_csv(source_file, sep="\t")
        else:
            print("len(data_df) :", len(data_df))
            break

    if filter_data_types:
        data_df = data_df[data_df["data_type"].isin(filter_data_types)]

    if number_samples:
        data_df = data_df.sample(number_samples, random_state=random_state)

    if filter_data_types or number_samples:
        data_df = data_df.reset_index(drop=True)

    return data_df

def get_new_data(source_file, source_folder="./output/upload/", number_samples=None, random_state=1144):
    full_source_file_name = os.path.join(source_folder, source_file)
    data_df = pd.read_csv(full_source_file_name, sep=",")
    if number_samples:
        data_df = data_df.sample(number_samples, random_state=random_state)
        data_df = data_df.reset_index(drop=True)
    return data_df

def convert_demo_data_into_list_json(data_df, limit=50, keep_labels=False,
                                     shuffle_list=[], random_shuffle=False, random_state=21524,
                                     new_label_col_name="assigned_label",
                                     old_label_col_name="event_type",
                                     id_col="tweet_id",
                                     text_col="tweet_text",
                                     label_col="assigned_label"):

    if keep_labels:
        data_df[new_label_col_name] = data_df[old_label_col_name]
    else:
        data_df[new_label_col_name] = "-"

    data_df[id_col] = data_df[id_col].values.astype(str)

    if len(shuffle_list) > 0:
        sort_index = data_df.index.values

        group_dict = {}
        for group in set(shuffle_list):
            group_dict[group] = []

        for (group, index) in zip(shuffle_list, sort_index):
            group_dict[group].append(index)

        dictionaries = list(group_dict.values())

        sort_indices = []
        for sort_indices_tuple in itertools.zip_longest(*dictionaries):
            temp_list = [x for x in [*sort_indices_tuple] if x is not None]
            sort_indices.extend(temp_list)

        data_df = data_df.iloc[sort_indices, :]

    if len(shuffle_list) == 0 and random_shuffle:
        data_df = data_df.sample(frac=1, random_state=random_state).reset_index(drop=True)

    all_texts = data_df[[id_col, text_col, label_col]].values.tolist()

    max_length = len(all_texts)
    if limit < max_length:
        all_texts_adj = random.sample(all_texts, limit)
    else:
        all_texts_adj = all_texts

    all_texts_json = [{"id": text[0], "text": text[1], "label": text[2]} for text in all_texts_adj]
    adj_text_ids = [text[0] for text in all_texts_adj]
    return all_texts_json, adj_text_ids

def convert_new_data_into_list_json(data_df,
                                    limit=50,
                                    shuffle_list=[],
                                    random_shuffle=False,
                                    random_state=21524,
                                    id_col="id",
                                    text_col="text",
                                    label_col="label"):

    data_df[id_col] = data_df[id_col].values.astype(str)
    data_df[label_col] = "-"

    if len(shuffle_list) > 0:
        sort_index = data_df.index.values

        group_dict = {}
        for group in set(shuffle_list):
            group_dict[group] = []

        for (group, index) in zip(shuffle_list, sort_index):
            group_dict[group].append(index)

        dictionaries = list(group_dict.values())

        sort_indices = []
        for sort_indices_tuple in itertools.zip_longest(*dictionaries):
            temp_list = [x for x in [*sort_indices_tuple] if x is not None]
            sort_indices.extend(temp_list)

        data_df = data_df.iloc[sort_indices, :]

    if len(shuffle_list) == 0 and random_shuffle:
        data_df = data_df.sample(frac=1, random_state=random_state).reset_index(drop=True)

    all_texts = data_df[[id_col, text_col, label_col]].values.tolist()

    max_length = len(all_texts)
    if limit < max_length:
        all_texts_adj = random.sample(all_texts, limit)
    else:
        all_texts_adj = all_texts

    all_texts_json = [{"id": text[0], "text": text[1], "label": text[2]} for text in all_texts_adj]
    adj_text_ids = [text[0] for text in all_texts_adj]
    return all_texts_json, adj_text_ids

def update_all_texts(all_texts, text_id, label):
    all_texts_df = pd.DataFrame(all_texts, columns=["tweet_id", "tweet_text", "assigned_label"])
    all_texts_df.loc[all_texts_df["tweet_id"] == str(text_id), "assigned_label"] = label

    all_texts_updated = all_texts_df.values

    return all_texts_updated

def filter_all_texts(all_text, filter_list, exclude_already_labeled=False):
    filtered_all_text = []

    # Slow - 10,000 records - duration 0:00:31.719903
    # for filter_id in filter_list:
    #     for text in all_text:
    #         if text["id"] == filter_id:
    #             filtered_all_text.append(text)

    # Faster - 10,000 records - duration 0:00:07.619622
    # [filtered_all_text.append(text) for text in all_text if text["id"] in filter_list]

    # Fastest - 10,000 records - duration 0:00:00.102955
    all_text_df = pd.DataFrame(all_text)
    filtered_all_text_df = all_text_df[all_text_df["id"].isin(filter_list)]

    if exclude_already_labeled:
        filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["label"].isin(["-"])]

    filtered_all_text = filtered_all_text_df.to_dict("records")

    return filtered_all_text

def search_all_texts_sql(all_text, include_search_term, exclude_search_term,
                         allow_search_to_override_existing_labels="No",
                         include_behavior="conjunction",
                         exclude_behavior="disjunction",
                         all_upper=True):

    all_text_df = pd.DataFrame(all_text)

    if allow_search_to_override_existing_labels == "No":
        filtered_all_text_df = all_text_df[all_text_df["label"].isin(["-"])]
    else:
        filtered_all_text_df = all_text_df

    if include_search_term and len(include_search_term) > 0:
        include_search_terms = include_search_term.split()
        if all_upper:
            if include_behavior == "disjunction":
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).str.upper().apply(
                    lambda text: any(word.upper() in text for word in include_search_terms))]
            else:
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).str.upper().apply(
                    lambda text: all(word.upper() in text for word in include_search_terms))]
        else:
            if include_behavior == "disjunction":
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).apply(
                    lambda text: any(word in text for word in include_search_terms))]
            else:
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).apply(
                    lambda text: all(word in text for word in include_search_terms))]

    if exclude_search_term and len(exclude_search_term) > 0:
        exclude_search_terms = exclude_search_term.split()
        if all_upper:
            if exclude_behavior == "disjunction":
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).str.upper().apply(
                    lambda text: any(word.upper() not in text for word in exclude_search_terms))]
            else:
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).str.upper().apply(
                    lambda text: all(word.upper() not in text for word in exclude_search_terms))]
        else:
            if exclude_behavior == "disjunction":
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).apply(
                    lambda text: any(word not in text for word in exclude_search_terms))]
            else:
                filtered_all_text_df = filtered_all_text_df[filtered_all_text_df["text"].astype(str).apply(
                    lambda text: all(word not in text for word in exclude_search_terms))]

    filtered_all_text = filtered_all_text_df.to_dict("records")

    search_results_sql = filtered_all_text
    search_results_length_sql = len(filtered_all_text)
    populate_texts_table_sql(filtered_all_text, table_name="searchResults")
    set_variable(name="SEARCH_RESULT_LENGTH", value=search_results_length_sql)
    return search_results_sql, search_results_length_sql

def update_texts_list(texts_list, sub_list_limit, old_obj_lst=[], new_obj_lst=[], texts_list_list=[]):
    updated_texts_list = texts_list

    if len(old_obj_lst) > 0 or len(new_obj_lst) > 0:
        if len(old_obj_lst) > 0:
            for old_obj in old_obj_lst:
                updated_texts_list.remove(old_obj)

        if len(new_obj_lst) > 0:
            for new_obj in new_obj_lst:
                updated_texts_list.append(new_obj)

    texts_list_list.clear()
    updated_texts_list_list = \
        [updated_texts_list[i:i + sub_list_limit] for i in range(0, len(updated_texts_list), sub_list_limit)]
    texts_list_list.extend(updated_texts_list_list)
    return updated_texts_list, updated_texts_list_list

def cosine_similarities(mat):
    # https://stackoverflow.com/questions/17627219/whats-the-fastest-way-in-python-to-calculate-cosine-similarity-given-sparse-mat
    col_normed_mat = pp.normalize(mat.tocsc(), axis=1)
    return col_normed_mat * col_normed_mat.T

def get_all_similarities(sparse_vectorized_corpus, corpus_text_ids):
    # Slow - vectorized_corpus.shape : (76484, 10) - Unable to allocate 43.6 GiB for an array with shape (76484, 76484) and data type float64
    # similarities = cosine_similarity(sparse_vectorized_corpus)
    # similarities_df = pd.DataFrame(similarities, columns=corpus_text_ids)

    # Faster - vectorized_corpus.shape : (76484, 10) - duration 0:01:43.129781
    # similarities = cosine_similarity(sparse_vectorized_corpus, dense_output=False)
    # similarities_df = pd.DataFrame.sparse.from_spmatrix(similarities, columns=corpus_text_ids)

    # Faster - vectorized_corpus.shape : (76484, 10) - duration 0:02:03.657139
    # similarities = np.dot(sparse_vectorized_corpus, sparse_vectorized_corpus.T)
    # similarities_df = pd.DataFrame.sparse.from_spmatrix(similarities, columns=corpus_text_ids)

    # Fastest - vectorized_corpus.shape : (76484, 10) - duration 0:01:59.331099
    similarities = cosine_similarities(sparse_vectorized_corpus)
    similarities_df = pd.DataFrame.sparse.from_spmatrix(similarities, columns=corpus_text_ids)

    similarities_df["id"] = corpus_text_ids
    similarities_df = similarities_df.set_index(["id"])
    return similarities_df

def get_all_similarities_one_at_a_time(sparse_vectorized_corpus, corpus_text_ids, text_id, keep_original=False):
    text_id_index = corpus_text_ids.index(text_id)
    # Fastest - vectorized_corpus.shape : (76484, 10) - duration 0:01:59.331099
    single_vectorized_record = sparse_vectorized_corpus[text_id_index, :]
    similarities = np.dot(single_vectorized_record, sparse_vectorized_corpus.T).toarray().ravel()
    similarities_series = pd.Series(similarities, index=corpus_text_ids)
    corpus_text_ids_adj = corpus_text_ids.copy()

    corpus_text_ids_adj.remove(text_id)

    similarities_series = similarities_series.filter(corpus_text_ids_adj)
    similarities_series.index.name = "id"
    similarities_series = similarities_series.sort_values(ascending=False)

    if keep_original:
        similarities_series = pd.concat([pd.Series(99.0, index=[text_id]), similarities_series])

    return similarities_series

def score_predictions(predictions_df, use_entropy=True, num_labels=5):
    prediction_values = predictions_df.values

    if use_entropy:
        all_score_combined = []
        for prediction in prediction_values:
            qk = [1/num_labels] * num_labels
            all_score_combined.append(entropy(prediction, base=num_labels, qk=qk))
    else:
        all_scores_num_non_zero = []
        all_scores_std = []
        for prediction in prediction_values:
            score_std = np.std(prediction)
            all_scores_std.append(score_std)
            score_num_non_zero = len([x for x in prediction if x > 0.0])
            all_scores_num_non_zero.append(score_num_non_zero)

        all_score_combined = np.array(all_scores_std) / np.array(all_scores_num_non_zero)

    return all_score_combined

def get_similarities_single_record(similarities_df, corpus_text_id):
    keep_indices = [x for x in similarities_df.index.values if x not in [corpus_text_id]]
    similarities = similarities_df[corpus_text_id]
    similarities = similarities.filter(keep_indices)
    similarities = similarities.sort_values(ascending=False)

    return similarities

def generate_click_record(click_location, click_type, click_object, guid=None):
    time_stamp = datetime.now()

    if not guid:
        guid = uuid.uuid4()

    click_record = {"click_id": str(guid),
                     "click_location": click_location,
                     "click_type": click_type,
                     "click_object": click_object,
                     "click_date_time": time_stamp}
    return click_record, guid

def generate_value_record(guid, value_type, value):
    value_record = {"click_id": str(guid),
                     "value_type": value_type,
                     "value": value}
    return value_record

def add_log_record(record, log=[]):
    log.extend(record)
    return None

def get_alert_message(label_summary_sql, overall_quality_score_decimal_sql, overall_quality_score_decimal_previous_sql,
                      texts_group_3_sql):
    # **** Future development ************************************************************************
    # This section will use these variables for future development.
    # The purpose is to give the user more sophisticated feedback while the user labels texts.
    # print("texts_group_3_sql :")
    # print(texts_group_3_sql)
    #
    # print(">> get_alert_message >> label_summary_sql :")
    # print(label_summary_sql)
    # **** Future development ************************************************************************

    if not overall_quality_score_decimal_previous_sql and not overall_quality_score_decimal_sql:
        alert_message = ""
    elif not overall_quality_score_decimal_previous_sql and overall_quality_score_decimal_sql:
        if overall_quality_score_decimal_sql < 0.60:
            alert_message = "More labels are required to improve the overall quality score."
        elif overall_quality_score_decimal_sql < 0.80:
            alert_message = "This is a fairly good start. Keep labeling to try to get the quality score close to 100%"
        else:
            alert_message = "This is a reasonable quality score. Click on the score above to go to the difficult text section."
    elif overall_quality_score_decimal_previous_sql < overall_quality_score_decimal_sql:
        alert_message = "The quality score is improving. Keep labeling."
    elif overall_quality_score_decimal_previous_sql > overall_quality_score_decimal_sql:
        alert_message = "The quality score has dropped. Examine the texts more carefully before assigning a label."
    else:
        alert_message = ""
    return alert_message

if __name__ == "__main__":
    start_time = datetime.now()
    print(">> Start time :", start_time.strftime("%m/%d/%Y %H:%M:%S"), "*"*100)

    # ***** Add code for testing the functions here ********************************************************************

    # ******************************************************************************************************************
    end_time = datetime.now()
    duration = end_time - start_time

    print(">> End time :", end_time.strftime("%m/%d/%Y @ %H:%M:%S"), "*"*100)
    print(">> Duration :", duration, "*"*100)


import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib import colors
from scipy import misc
import argparse
import os

parser = argparse.ArgumentParser(description='Display recorded Kinect data', fromfile_prefix_chars="@")
parser.add_argument('-p', dest='recording_path', action='store', type=str, default="", help="Which recording do you want to display?")
parser.add_argument('-f', dest='frame', action='store', type=int, default=0, help='Which frame do you want to display')

args = parser.parse_args()
frame = args.frame

recording_path = args.recording_path.strip(os.sep)

#COLOR
colorFile = os.path.join(recording_path, 'COLOR', '%d.uint8' % frame)
if (os.path.isfile(colorFile)):
    color = np.fromfile(colorFile, dtype=np.uint8)
    color.shape = (1080, 1920, 4)
    plt.subplot(2, 2, 1)
    plt.imshow(color)
    plt.title('Original color 1920x1080')

#DEPTH
depthFile = os.path.join(recording_path, 'DEPTH', '%d.uint16' % frame)
if (os.path.isfile(depthFile)):
    min_depth = 500  # Min reliable depth
    max_depth = 4500 # Max reliable depth

    max_value = np.iinfo(np.uint16).max

    depth = np.fromfile(depthFile, dtype=np.uint16)
    depth.shape = (424, 512)

    plt.subplot(2, 2, 2)
    plt.imshow(depth, interpolation='nearest', cmap=cm.gist_heat)
    plt.title('Depth values 512x424')

#INDEX
indexFile = os.path.join(recording_path, 'INDEX', '%d.uint8' % frame)
if (os.path.isfile(indexFile)):
    index = np.fromfile(indexFile, dtype=np.uint8)
    index.shape = (424, 512)
    index[index == 255] = 8

    index_color_map = colors.ListedColormap(['red', 'green', 'blue', 'cyan', 'magenta', 'yellow', '#ff7700', 'white', 'black'])

    plt.subplot(2, 2, 3)
    plt.imshow(index, interpolation='nearest', cmap=index_color_map)
    plt.title('Segmentation data 512x424')

#TRACKEDCOLOR
trackedFile = os.path.join(recording_path, 'TRACKEDCOLOR', '%d.uint8' % frame)
if (os.path.isfile(trackedFile)):
    tracked = np.fromfile(trackedFile, dtype=np.uint8);
    tracked.shape = (424, 512, 4)
    plt.subplot(2, 2, 4)
    plt.imshow(tracked)
    plt.title('Color mapped to tracked depth space 512x424')

plt.show()


import os
import math
import torch
import torch.distributed as dist
from torch.nn import Module
from torch.utils.data.distributed import DistributedSampler
import multiprocessing as mp
from torch.utils.data.sampler import Sampler
import numpy as np

class DistModule(Module):
    def __init__(self, module):
        super(DistModule, self).__init__()
        self.module = module
        broadcast_params(self.module)
        #dist._clear_group_cache()
    def forward(self, *inputs, **kwargs):
        return self.module(*inputs, **kwargs)
    def train(self, mode=True):
        #dist._clear_group_cache()
        super(DistModule, self).train(mode)
        self.module.train(mode)

def average_gradients(model):
    """ average gradients """
    """ test network """
    """
    for name, param in model.named_parameters():
        print(name,type(param.grad))
        dist.all_reduce(param.grad.data)
    """
    for param in model.parameters():
        if param.requires_grad:
            dist.all_reduce(param.grad.data)

def broadcast_params(model):
    """ broadcast model parameters """
    for p in model.state_dict().values():
        dist.broadcast(p, 0)

def dist_init(port):
    if mp.get_start_method(allow_none=True) != 'spawn':
        mp.get_context('spawn')
        # mp.set_start_method('spawn')
    proc_id = int(os.environ['SLURM_PROCID'])
    ntasks = int(os.environ['SLURM_NTASKS'])
    node_list = os.environ['SLURM_NODELIST']
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(proc_id%num_gpus)

    if '[' in node_list:
        beg = node_list.find('[')
        pos1 = node_list.find('-', beg)
        if pos1 < 0:
            pos1 = 1000
        pos2 = node_list.find(',', beg)
        if pos2 < 0:
            pos2 = 1000
        node_list = node_list[:min(pos1,pos2)].replace('[', '')
    addr = node_list[10:].replace('-', '.')
    print(addr)
    os.environ['MASTER_PORT'] = port
    os.environ['MASTER_ADDR'] = addr
    os.environ['WORLD_SIZE'] = str(ntasks)
    os.environ['RANK'] = str(proc_id)
    dist.init_process_group(backend='nccl')

    rank = dist.get_rank()
    world_size = dist.get_world_size()
    return rank, world_size

def dist_init_pytorch(port):
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend='nccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    return rank, world_size

class CustomSampler(DistributedSampler):
    def __init__(self, dataset, num_replicas=None, rank=None):
        if num_replicas is None:
            num_replicas = dist.get_world_size()
        if rank is None:
            rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas
    def __iter__(self):
        # deterministically shuffle based on epoch
        g = torch.Generator()
        g.manual_seed(self.epoch)
        indices = list(torch.randperm(len(self.dataset), generator=g))
        offset = self.num_samples * self.rank
        indices = indices[offset:min(offset + self.num_samples, len(indices))]
        return iter(indices)

    def __len__(self):
        return self.num_samples

class DistributedSampler(Sampler):
    """Sampler that restricts data loading to a subset of the dataset.

    It is especially useful in conjunction with
    :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
    process can pass a DistributedSampler instance as a DataLoader sampler,
    and load a subset of the original dataset that is exclusive to it.

    .. note::
        Dataset is assumed to be of constant size.

    Arguments:
        dataset: Dataset used for sampling.
        world_size (optional): Number of processes participating in
            distributed training.
        rank (optional): Rank of the current process within world_size.
    """

    def __init__(self, dataset, world_size=None, rank=None, round_up=False):
        if world_size is None:
            world_size = dist.get_world_size()
        if rank is None:
            rank = dist.get_rank()
        self.dataset = dataset
        self.world_size = world_size
        self.rank = rank
        self.round_up = round_up
        self.epoch = 0

        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.world_size))
        if self.round_up:
            self.total_size = self.num_samples * self.world_size
        else:
            self.total_size = len(self.dataset)

    def __iter__(self):
        # deterministically shuffle based on epoch
        g = torch.Generator()
        g.manual_seed(self.epoch)
        indices = list(torch.randperm(len(self.dataset), generator=g))
        # indices = list(range(len(self.dataset)))
        # indices = self.dataset
        # print ('all: ', indices)
        # print (indices[:(self.total_size - len(indices))])
        # add extra samples to make it evenly divisible
        if self.round_up:
            indices += indices[:(self.total_size - len(indices))]
        assert len(indices) == self.total_size

        # subsample
        offset = self.num_samples * self.rank
        indices = indices[offset:min(offset + self.num_samples, self.total_size)]
        if self.round_up or (not self.round_up and self.rank < self.world_size - 1):
            assert len(indices) == self.num_samples

        # print (self.rank, len(indices))

        # print (self.rank, indices)
        return iter(indices)

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):
        self.epoch = epoch

class DistributedGivenIterationSampler(Sampler):
    def __init__(self, dataset, total_iter, batch_size, world_size=None, rank=None, last_iter=-1, shuffle_strategy=0):
        if world_size is None:
            world_size = dist.get_world_size()
        if rank is None:
            rank = dist.get_rank()
        assert rank < world_size
        self.dataset = dataset
        self.total_iter = total_iter
        self.batch_size = batch_size
        self.world_size = world_size
        self.rank = rank
        self.last_iter = last_iter
        self.shuffle_strategy = shuffle_strategy

        self.total_size = self.total_iter * self.batch_size

        self.call = 0

    def gen_s2(self):
        length = len(self.dataset)
        print('using shuffle strategy 2, initializing index...')

        for i in range((self.last_iter + 1) * self.batch_size, self.total_size):
            yield np.random.randint(0, length)

    def __iter__(self):
        if self.call == 0:
            self.call = 1
            if self.shuffle_strategy == 2:
                np.random.seed(self.rank)  # set different random seed
                return self.gen_s2()
            else:
                self.indices = self.gen_new_list()
                return iter(self.indices[(self.last_iter + 1) * self.batch_size:])
        else:
            raise RuntimeError("this sampler is not designed to be called more than once!!")

    def gen_new_list(self):

        # each process shuffle independently
        if self.shuffle_strategy == 0:

            np.random.seed(self.rank)

            indices = np.arange(len(self.dataset))
            indices = indices[:self.total_size]
            num_repeat = (self.total_size - 1) // indices.shape[0] + 1
            indices = np.tile(indices, num_repeat)
            indices = indices[:self.total_size]

            for beg in range(0, self.total_size, len(self.dataset)):
                end = min(beg + len(self.dataset), self.total_size)
                np.random.shuffle(indices[beg:end])

        # each process shuffle all list with same seed, and pick one piece according to rank
        elif self.shuffle_strategy == 1:

            np.random.seed(0)

            all_size = self.total_size * self.world_size
            indices = np.arange(len(self.dataset))
            indices = indices[:all_size]
            num_repeat = (all_size - 1) // indices.shape[0] + 1
            indices = np.tile(indices, num_repeat)
            indices = indices[:all_size]

            np.random.shuffle(indices)
            beg = self.total_size * self.rank
            indices = indices[beg:beg + self.total_size]

        else:
            raise RuntimeError('Unknown shuffle strategy')

        assert len(indices) == self.total_size

        return indices

    def __len__(self):
        return self.total_size - (self.last_iter + 1) * self.batch_size


import os
import datetime
import json
import random
import numpy as np
import cv2
from PIL import Image
from helpers.common import *
from helpers.input import press_keyboard_key, release_keyboard_key
from helpers.models.inception_v3 import inception_v3
from networks.abstract import AbstractNetwork

'''
The typical network for driving games - WASD (forward, left, backward, right).
'''

class WASDNetwork(AbstractNetwork):

    def get_xy(self, iteration):
        X = []
        Y = []

        processed_data_file = open(self.processed_data_file_path, 'r+').read()
        processed_data = processed_data_file.split("\n")
        processed_data = list(filter(None, processed_data))

        if self.train_shuffle_data:
            random.shuffle(processed_data)

        from_index, to_index = get_from_and_to_index(iteration=iteration,
                                                     batch_size=self.train_batch_size,
                                                     total_size=len(processed_data))
        processed_data = processed_data[from_index:to_index]

        for row in processed_data:
            row_data = json.loads(row)

            X.append(np.array(Image.open(row_data['image_path'])))
            Y.append(self.convert_controls_to_array(row_data['controls']))

        return X, Y

    # Model related stuff
    def prepare_model(self, load_existing=True):
        # Prepare data
        input2_size, input1_size = self.size
        output_size = len(self.convert_controls_to_array({}))

        # Return the model
        self.model = inception_v3(
            input1_size=input1_size,
            input2_size=input2_size,
            output_size=output_size,
            tensorboard_dir=self.network_logs_dir,
            checkpoint_path=self.network_checkpoint_path,
            best_checkpoint_path=self.network_checkpoint_path
        )

        model_file_exists = os.path.isfile(self.network_model_path + '.index')

        if load_existing and not self.force_new_model and model_file_exists:
            self.model.load(self.network_model_path)

        return self.model

    def trigger_action(self, action):
        if action == 'forward' or action == 'forward+left' or action == 'forward+right':
            press_keyboard_key('w')
        else:
            release_keyboard_key('w')

        if action == 'backward' or action == 'backward+left' or action == 'backward+right':
            press_keyboard_key('s')
        else:
            release_keyboard_key('s')

        if action == 'left' or action == 'forward+left' or action == 'backward+left':
            press_keyboard_key('a')
        else:
            release_keyboard_key('a')

        if action == 'right' or action == 'forward+right' or action == 'backward+right':
            press_keyboard_key('a')
        else:
            release_keyboard_key('a')

        if action == 'none':
            release_keyboard_key('w')
            release_keyboard_key('s')
            release_keyboard_key('a')
            release_keyboard_key('d')

    # Helpers
    def process_image(self, image, return_array=False):
        if isinstance(image, str):
            image = Image.open(image)

        # Do processing
        processed_image = image.resize(self.size, Image.ANTIALIAS)
        processed_image_array = np.array(processed_image)
        processed_image_array = cv2.cvtColor(
            processed_image_array, cv2.COLOR_BGR2RGB)

        if return_array:
            return processed_image_array

        return Image.fromarray(processed_image_array)

    def get_controls_from_inputs(self, inputs):
        forward = inputs['keyboard']['w'] or inputs['gamepad']['axes']['right_trigger'] > 0
        backward = inputs['keyboard']['s'] or inputs['gamepad']['axes']['left_trigger'] > 0
        left = inputs['keyboard']['a'] or inputs['gamepad']['axes']['left']['x'] < 0
        right = inputs['keyboard']['d'] or inputs['gamepad']['axes']['left']['x'] > 0

        return {
            'forward': forward,
            'backward': backward,
            'left': left,
            'right': right,
        }

    def get_controls_map(self):
        return {
            'forward': [1, 0, 0, 0, 0, 0, 0, 0, 0],
            'backward': [0, 1, 0, 0, 0, 0, 0, 0, 0],
            'left': [0, 0, 1, 0, 0, 0, 0, 0, 0],
            'right': [0, 0, 0, 1, 0, 0, 0, 0, 0],
            'forward+left': [0, 0, 0, 0, 1, 0, 0, 0, 0],
            'forward+right': [0, 0, 0, 0, 0, 1, 0, 0, 0],
            'backward+left': [0, 0, 0, 0, 0, 0, 1, 0, 0],
            'backward+right': [0, 0, 0, 0, 0, 0, 0, 1, 0],
            'none': [0, 0, 0, 0, 0, 0, 0, 0, 1],
        }

    def convert_controls_to_array(self, controls):
        controls_map = self.get_controls_map()

        left = 'left' in controls and controls['left']
        right = 'right' in controls and controls['right']
        forward = 'forward' in controls and controls['forward']
        backward = 'backward' in controls and controls['backward']
        forwardLeft = (forward and left)
        forwardRight = (forward and right)
        backwardLeft = (backward and left)
        backwardRight = (backward and right)

        output = controls_map['none']

        if forwardLeft:
            output = controls_map['forward+left']
        elif forwardRight:
            output = controls_map['forward+right']
        elif backwardLeft:
            output = controls_map['backward+left']
        elif backwardRight:
            output = controls_map['backward+right']
        elif forward:
            output = controls_map['forward']
        elif backward:
            output = controls_map['backward']
        elif left:
            output = controls_map['left']
        elif right:
            output = controls_map['right']

        return output


import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
import cv2
###########################################################################################
sys.path.append("..")

import label_map_util

import visualization_utils as vis_util

MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

# 학습이 완료된(Frozen) Object Detection을 위한 그래프 pb파일 경로를 지정한다.
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

NUM_CLASSES = 90

opener = urllib.request.URLopener()
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
  file_name = os.path.basename(file.name)
  if 'frozen_inference_graph.pb' in file_name:
    tar_file.extract(file, os.getcwd())

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='')

label_map = label_map_util.load_labelmap(os.path.join('./models/research/object_detection/data/', 'mscoco_label_map.pbtxt'))
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)

def load_image_into_numpy_array(image):
  (im_width, im_height,_) = image.shape
  return np.array(image.getdata()).reshape(
      (im_height, im_width, 3)).astype(np.uint8)

# inch로 표현한 출력 이미지의 크기

with detection_graph.as_default():
  with tf.Session(graph=detection_graph) as sess:

    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')

    detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')

    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')
    img = cv2.imread('../test_images/image3.jpg')
    image_np = np.array(img)
      # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
    image_np_expanded = np.expand_dims(image_np, axis=0)
      # Actual detection.
    (boxes, scores, classes, num) = sess.run(
          [detection_boxes, detection_scores, detection_classes, num_detections],
          feed_dict={image_tensor: image_np_expanded})
      # Visualization of the results of a detection.
    vis_util.visualize_boxes_and_labels_on_image_array(
          image_np,
          np.squeeze(boxes),
          np.squeeze(classes).astype(np.int32),
          np.squeeze(scores),
          category_index,
          use_normalized_coordinates=True,
          line_thickness=2)
    cv2.imshow('object detection', cv2.resize(image_np, (800,600)))

    cv2.waitKey(0)
    cv2.destroyAllWindows()