import os
import pandas as pd
from datetime import datetime
import csv
import datetime
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
import matplotlib.dates as mdates

train_headers = ['timestamp', 'eps_train', 'eps_trained_session', 'sum', 'mean']
eval_headers  = ['timestamp', 'method', 'eps_train', 'eval_eps_used', 'sum', 'mean']

model_path = 'models'


def dataframes(model_name):
    def df_timestamp_to_datetime(df):
        df['timestamp'] = df['timestamp'].map(lambda t: datetime.datetime.fromtimestamp(t))
        return df
    
    log_path = os.path.join(model_path, model_name, 'logs')
    raw_dfs = [ pd.read_csv(os.path.join(log_path, 'eval.log'), sep=';', names=eval_headers),
                pd.read_csv(os.path.join(log_path, 'train.log'), sep=';', names=train_headers) ]
    dfs = [ df_timestamp_to_datetime(df) for df in raw_dfs ]
    dataframes = {
        'eval': dfs[0],
        'train': dfs[1] 
    }
    return dataframes


if __name__ == '__main__':
    fig, ax = plt.subplots(1, 1)

    plt.ion()
    plt.title('Mean over episodes')
    plt.xlabel('Episodes trained')
    plt.ylabel('Mean')
    plt.grid(True)

    #ax.set_xlim(left=0)
    ax.set_ylim([-2, 2])

    plt.show()

    while True:
        df = dataframes('default')['eval']

        print(df)
        
        x = df['eps_train']
        y = df['mean']
    
        plt.scatter(x, y, c=[[1, 0.5, 0]])
        #fig.canvas.draw()
        plt.pause(2)