"""Plot the actual epoch-average training losses recorded by TensorBoard."""
import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

ROOT = Path('/data/realstyle2')
OUT = ROOT / 'plots'
OUT.mkdir(exist_ok=True)
runs = []
for path in sorted((ROOT / 'logs/tensorboard').rglob('events*')):
    events = EventAccumulator(str(path), size_guidance={'scalars': 0})
    events.Reload()
    if 'loss/epoch_average' in events.Tags().get('scalars', []):
        runs.append((path, events.Scalars('loss/epoch_average')))
assert len(runs) == 1, 'Expected one completed training run'
source, values = runs[0]
epochs = [v.step for v in values]
losses = [v.value for v in values]
assert epochs == list(range(1, 41))
with (OUT / 'mean_loss.csv').open('w', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(['epoch', 'mean_loss'])
    writer.writerows(zip(epochs, losses))
minimum = min(range(len(losses)), key=losses.__getitem__)
summary = {'source': str(source), 'scalar': 'loss/epoch_average',
           'epochs': len(values), 'first_epoch_loss': losses[0],
           'final_epoch_loss': losses[-1], 'minimum_loss_epoch': epochs[minimum],
           'minimum_epoch_loss': losses[minimum], 'mean_of_epoch_losses': sum(losses)/len(losses)}
(OUT / 'loss_summary.json').write_text(json.dumps(summary, indent=2))

plt.rcParams.update({'font.family': 'DejaVu Sans', 'font.size': 12,
                     'axes.spines.top': False, 'axes.spines.right': False,
                     'axes.labelcolor': '#334155', 'xtick.color': '#64748b',
                     'ytick.color': '#64748b', 'svg.fonttype': 'none'})
fig, ax = plt.subplots(figsize=(12, 5.9), facecolor='white')
fig.subplots_adjust(left=.095, right=.975, top=.77, bottom=.17)
fig.text(.095, .92, 'REALSTYLE2 / ANIMA', fontsize=11, weight='bold', color='#64748b')
fig.text(.095, .852, 'Mean training loss by epoch', fontsize=24, weight='bold', color='#0f172a')
fig.text(.975, .86, '40 epochs  ·  4,320 steps', ha='right', fontsize=12, color='#475569')
ax.plot(epochs, losses, color='#2563eb', linewidth=2, marker='o', markersize=4.5,
        markerfacecolor='white', markeredgewidth=1.4, label='Epoch mean loss')
ax.scatter([epochs[minimum]], [losses[minimum]], color='#0f766e', s=65, zorder=5)
ax.annotate(f'Min: epoch {epochs[minimum]}\n{losses[minimum]:.5f}',
            xy=(epochs[minimum], losses[minimum]), xytext=(-48, -36),
            textcoords='offset points', fontsize=11, color='#0f766e', ha='right',
            arrowprops={'arrowstyle': '-', 'color': '#0f766e', 'lw': 1})
ax.annotate(f'Final: {losses[-1]:.5f}', xy=(epochs[-1], losses[-1]), xytext=(-8, 15),
            textcoords='offset points', ha='right', fontsize=11, color='#1d4ed8')
ax.set(xlim=(.5, 41), ylim=(min(losses)-.0038, max(losses)+.0018),
       xlabel='Epoch', ylabel='Mean loss')
ax.set_xticks([1, 5, 10, 15, 20, 25, 30, 35, 40])
ax.yaxis.set_major_locator(MultipleLocator(.002))
ax.yaxis.set_major_formatter(FormatStrFormatter('%.3f'))
ax.grid(axis='y', color='#e2e8f0', linewidth=.8)
ax.set_axisbelow(True)
ax.spines['left'].set_color('#cbd5e1')
ax.spines['bottom'].set_color('#cbd5e1')
fig.text(.095, .045, 'Source: TensorBoard loss/epoch_average · Raw epoch means, no smoothing · Y-axis does not start at zero',
         fontsize=10, color='#64748b')
fig.savefig(OUT / 'mean_loss.png', dpi=200, facecolor='white')
fig.savefig(OUT / 'mean_loss.svg', facecolor='white')
plt.close(fig)
print(json.dumps(summary, indent=2))