#!/usr/bin/env python3
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib
import os

DIR: str = os.path.dirname(__file__) 

matplotlib.use('pgf')
matplotlib.rcParams.update({
    'pgf.texsystem': 'pdflatex',
    'font.family': 'serif',
    'text.usetex': True,
    'pgf.rcfonts': False,
})

df = pd.read_csv(DIR + '/time.csv')
df['cpu'] = df['user'] + df['sys']
del df['real']
del df['user']
del df['sys']

df_cpu = df.loc[~df.cpu.isnull(), :]
del df_cpu['ltl']
df_cpu = df_cpu.melt(id_vars=['cpu'])
df_cpu['value'] = df_cpu['value'].astype(str)

df_timeout = df.loc[:, :].copy()
df_timeout['timeout'] = 0
df_timeout.loc[df_timeout['cpu'].isnull(), 'timeout'] = 1
del df_timeout['cpu']


def timeout_by(name: str) -> pd.DataFrame:
    df_return = df_timeout.loc[:, [name, 'timeout']] \
        .groupby(name) \
        .sum() \
        .reset_index()
    df_return[name] = df_return[name].astype(str)
    return df_return


def plot_by(df_cpu: pd.DataFrame, name: str, color: any): 
    # Initialize the matplotlib figure
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))

    sns.stripplot(ax=ax1, color=color, data=df_cpu.loc[df_cpu.variable == name],
                  x='cpu', y='value')
    sns.barplot(ax=ax2, color=color, data=timeout_by(name), y=name, x='timeout',
                width=0.75)

    ax1.set(ylabel='Value of ' + name.upper(), xlim=[0, 300],
            xlabel='CPU time (seconds)')
    ax2.set(ylabel='Value of ' + name.upper(), xlim=[0, 100], 
            xlabel='Timeouts (%)')
    
    sns.despine(left=True, bottom=True)
    plt.savefig(DIR + '/../plots/' + name + '.pgf')


def main():
    plot_by(df_cpu, 'n', 'teal')
    plot_by(df_cpu, 'length', 'red')
    plot_by(df_cpu, 'r', 'orange')


if __name__ == '__main__':
    main()