import os
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

"""
def tabulate_events(dpath):
    files = os.listdir(dpath)[0]
    summary_iterators = [EventAccumulator(os.path.join(dpath, files)).Reload()]

    tags = summary_iterators[0].Tags()['scalars']

    for it in summary_iterators:
        assert it.Tags()['scalars'] == tags

    out = {t: [] for t in tags}
    steps = []

    for tag in tags:
        steps = [e.step for e in summary_iterators[0].Scalars(tag)]

        for events in zip(*[acc.Scalars(tag) for acc in summary_iterators]):
            assert len(set(e.step for e in events)) == 1

            out[tag].append([e.value for e in events])
    
    return out, steps

events, steps = tabulate_events('/home/vedant/pytorch_sac_ae/log/runs')

data = []

for tag, values in events.items():
    for run_idx, run_values in enumerate(values):
        for step_idx, value in enumerate(run_values):
            data.append({
                'tag': tag,
                'run': run_idx,
                'step': steps[step_idx],
                'value': value,
            })

df = pd.DataFrame(data)
print(df.head())
exit()

plt.figure(figsize=(10,6))
sns.lineplot(data=df, x='step', y='value', hue='tag', ci='sd')
plt.show()
"""


from tensorboard.backend.event_processing import event_accumulator



def data_from_tb(files):
    all_steps, all_rewards = [], []
    for file in files:
        ea = event_accumulator.EventAccumulator(file, size_guidance={'scalars': 0})
        ea.Reload()

        episode_rewards = ea.Scalars('train/episode_reward')
        steps = [event.step for event in episode_rewards][:990000]
        rewards = [event.value for event in episode_rewards][:990000]
        all_steps.append(steps)
        all_rewards.append(rewards)
    return all_steps, all_rewards


files = ['/home/vedant/pytorch_sac_ae/log/runs/tb_21_05_2023-13_19_36/events.out.tfevents.1684667976.cpswkstn6-nvidia4090.1749060.0',
         '/home/vedant/pytorch_sac_ae/log/runs/tb_22_05_2023-09_56_30/events.out.tfevents.1684742190.cpswkstn6-nvidia4090.1976229.0']

all_steps, all_rewards = data_from_tb(files)
mean_rewards = np.mean(all_rewards, axis=0)
std_rewards = np.std(all_rewards, axis=0)
mean_steps = np.mean(all_steps, axis=0) 

df = pd.DataFrame({'Steps': mean_steps,'Rewards': mean_rewards,'Standard Deviation': std_rewards})

sns.relplot(x='Steps', y='Rewards', kind='line', data=df, ci="sd")
plt.fill_between(df['Steps'], df['Rewards'] - df['Standard Deviation'], df['Rewards'] + df['Standard Deviation'], color='b', alpha=.1)
plt.title("Mean Rewards vs Steps with Standard Deviation")
plt.show()