import tkinter as tk

from ui.app_state import AppState
from ui.components.digit_drawer import DigitDrawer
from ui.front_page.plots.predictions import PredictionsPlot


class TestModelSection(tk.LabelFrame):
    def __init__(self, parent, app_state: AppState):
        super().__init__(parent, text="Model testing")
        self.app_state = app_state
        self.update_ui()

    def clear_ui(self):
        for widget in self.winfo_children():
            widget.destroy()

    def update_ui(self):
        self.clear_ui()

        self.digit_drawer = DigitDrawer(self, 100, 100)
        self.digit_drawer.pack(fill=tk.BOTH, expand=True)

        # Predict Button (converts drawing to 28x28 and shows the array)
        self.predict_button = tk.Button(self, text="Predict", command=self.predict_number)
        self.predict_button.pack(fill=tk.BOTH, expand=True)

        frame_prediction = tk.Frame(self, height=200)
        frame_prediction.pack(fill=tk.BOTH, expand=True)
        (tk.Label(frame_prediction, text="Prediction: ")
            .pack(side=tk.LEFT))
        self.lbl_prediction = tk.Label(frame_prediction, text="/")
        self.lbl_prediction.pack(side=tk.LEFT)
        self.prediction_plot = PredictionsPlot(self)
        self.prediction_plot.pack(side=tk.BOTTOM, anchor=tk.S, fill=tk.X, expand=True)

    def predict_number(self):
        inputs = self.digit_drawer.convert_to_array()
        raw_predictions, predictions = self.app_state.neural_net.predict([inputs])
        print(predictions)
        self.lbl_prediction.config(text=f"{predictions[0]}")
        self.prediction_plot.update_data(raw_predictions)