import pandas as pd
import sqlite3
import os
from matplotlib import pyplot as plt
import seaborn as sns

conn = sqlite3.connect('db.sqlite')


def check_folder():
    path = './charts'
    if not os.path.exists(path):
        os.mkdir(path)
        print('create charts directory...')

def time_series_chart(x, y, y_label, title):
    sns.set()
    plt.figure(figsize=(18,7))
    sns.lineplot(x, y)
    plt.xlabel('Round')
    plt.ylabel(y_label)
    plt.title(title)
    plt.savefig(os.path.join('charts', title + '.png'))


def main():
    check_folder()
    # loading records
    res = pd.read_sql_query(
        '''select timestamp, 
                  normal_user_count,
                  abnormal_user_count,
                  nyms_count
            from cluster_result''',
        con=conn)

    # plot time-series key indicator fig

    time_series_chart(res.index, res.loc[:, 'nyms_count'],
                      'Size of Pseudonym',
                      'Attack Anonymity Set')
    time_series_chart(res.index, res.loc[:, 'normal_user_count'],
                      'Forwarded User Counts',
                      'Normal Active User')
    time_series_chart(res.index, res.loc[:, 'abnormal_user_count'],
                      'Delayed User Counts',
                      'Abnormal Active User')
    print('save finished.')


if __name__ == '__main__':
    main()