import tkinter as tk

from neural_net.epoch import Epoch

class EpochInformation(tk.LabelFrame):
    def __init__(self, parent, epoch: Epoch):
        super().__init__(parent, text="Last epoch info")
        self.epoch = epoch

        self.lbl_epoch_training_time = None
        self.lbl_last_loss = None
        self.create_ui()

    def create_ui(self):
        row = 0
        tk.Label(self, text="Duration:", anchor=tk.W).grid(column=0, row=row,
                                                                       sticky=tk.E,
                                                                       padx=(10, 20), pady=5)
        self.lbl_epoch_training_time = tk.Label(self, text=f"{self.epoch.duration:.2f}sec")
        self.lbl_epoch_training_time.grid(column=1, row=row, sticky=tk.E, padx=10, pady=5)
        row += 1
        tk.Label(self, text="Loss value:", anchor=tk.W).grid(column=0, row=row,
                                                                   sticky=tk.E, padx=(10, 20),
                                                                   pady=5)
        self.lbl_last_loss = tk.Label(self, text=f"{self.epoch.loss:.4f}")
        self.lbl_last_loss.grid(column=1, row=row, sticky=tk.E, padx=10, pady=5)

        row += 1
        tk.Label(self, text="Learning rate:", anchor=tk.W).grid(column=0, row=row,
                                                                   sticky=tk.E, padx=(10, 20),
                                                                   pady=5)
        self.lbl_learning_rate = tk.Label(self, text=f"{self.epoch.learning_rate:.4f}")
        self.lbl_learning_rate.grid(column=1, row=row, sticky=tk.E, padx=10, pady=5)

    def update(self):
        print(f"Updating training data for epoch {self.epoch}")
        self.lbl_epoch_training_time.config(text=f"{self.epoch.duration:.2f}sec")
        self.lbl_last_loss.config(text=f"{self.epoch.loss:.4f}")
        self.lbl_learning_rate.config(text=f"{self.epoch.learning_rate:.4f}")

    def set_epoch(self, epoch: Epoch):
        self.epoch = epoch
        self.update()