# | fig-width: 15
# | fig-height: 10
# | output: asis
import matplotlib.pyplot as plt
from pyfragments import AnimatedFigure

with AnimatedFigure() as ani:
# Data for plotting
  train_loss = [1.0, 0.6, 0.3, 0.2, 0.1]
  val_loss = [1.2, 0.8, 0.5, 0.4, 0.6]
  iterations = [1, 2, 3, 4, 5]

  plt.xlim(0, 6)
  plt.ylim(0, 1.5)

# Labels and legend
  plt.xlabel('Trainingsiterationen', fontsize=14)
  plt.ylabel('Loss', fontsize=14)
  plt.legend(loc='upper right', fontsize=12)

# Plot
  for x in range(len(train_loss)+1):
    with ani.fragment():
      plt.plot(iterations[:x], train_loss[:x], 'bo-', label='Training', linewidth=2)
      plt.plot(iterations[:x], val_loss[:x], 'go-', label='Validation', linewidth=2)

# Add a dashed vertical line in the middle
  mid_iteration = (iterations[0] + iterations[-1]) / 2

  with ani.fragment():
    plt.axvline(x=mid_iteration, color='r', linestyle='--', linewidth=1.5)