import wx
import threading
import time
import numpy as np
import matplotlib
matplotlib.use('WXAgg')
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas
from matplotlib.figure import Figure

class RandomWalkFrame(wx.Frame):
    def __init__(self, parent=None, title="wxPython + Matplotlib Random Walk (NumPy)"):
        super().__init__(parent, title=title, size=(800, 600))

        # Main panel
        panel = wx.Panel(self)
        vbox = wx.BoxSizer(wx.VERTICAL)

        # Matplotlib Figure
        self.figure = Figure()
        self.ax = self.figure.add_subplot(111)
        self.ax.set_title("Random Walk")
        self.line, = self.ax.plot([], [], lw=2)

        self.canvas = FigureCanvas(panel, -1, self.figure)
        vbox.Add(self.canvas, 1, wx.EXPAND)

        # Buttons
        hbox = wx.BoxSizer(wx.HORIZONTAL)
        self.start_btn = wx.Button(panel, label="Start")
        self.stop_btn = wx.Button(panel, label="Stop")
        self.stop_btn.Disable()
        hbox.Add(self.start_btn, 0, wx.ALL, 5)
        hbox.Add(self.stop_btn, 0, wx.ALL, 5)
        vbox.Add(hbox, 0, wx.ALIGN_CENTER)

        panel.SetSizer(vbox)

        # Bind buttons
        self.start_btn.Bind(wx.EVT_BUTTON, self.on_start)
        self.stop_btn.Bind(wx.EVT_BUTTON, self.on_stop)

        # Data for random walk
        self.xdata = np.array([], dtype=np.int32)
        self.ydata = np.array([], dtype=np.int32)
        self.running = False

        # Blitting setup
        self.background = None

        # FPS tracking
        self.last_time = time.time()
        self.frame_count = 0

    def on_start(self, event):
        if not self.running:
            self.running = True
            self.start_btn.Disable()
            self.stop_btn.Enable()
            self.prepare_blit()
            threading.Thread(target=self.random_walk, daemon=True).start()

    def on_stop(self, event):
        self.running = False
        self.start_btn.Enable()
        self.stop_btn.Disable()

    def prepare_blit(self):
        # Draw once and cache background
        self.canvas.draw()
        self.background = self.canvas.copy_from_bbox(self.ax.bbox)

    def random_walk(self):
        x = self.xdata[-1] if self.xdata.size > 0 else 0
        y = self.ydata[-1] if self.ydata.size > 0 else 0

        while self.running:
            steps = 100
            dx = np.arange(x + 1, x + steps + 1)
            dy = y + np.cumsum(np.random.choice([-1, 1], size=steps))

            # Append to arrays
            self.xdata = np.concatenate((self.xdata, dx))
            self.ydata = np.concatenate((self.ydata, dy))

            x = dx[-1]
            y = dy[-1]

            # Update GUI
            wx.CallAfter(self.update_plot)
            time.sleep(0.005)  # Control speed

    def update_plot(self):
        # Auto-scale axes
        if self.xdata.size > 0:
            self.ax.set_xlim(0, self.xdata[-1] + 10)
            ymin, ymax = self.ydata.min(), self.ydata.max()
            self.ax.set_ylim(ymin - 5, ymax + 5)

        # Restore background and update line
        self.canvas.restore_region(self.background)
        self.line.set_data(self.xdata, self.ydata)
        self.ax.draw_artist(self.line)
        self.canvas.blit(self.ax.bbox)

        # FPS calculation
        self.frame_count += 1
        now = time.time()
        if now - self.last_time >= 1.0:  # Every second
            fps = self.frame_count / (now - self.last_time)
            self.SetTitle(f"wxPython + Matplotlib Random Walk (NumPy) | FPS: {fps:.1f}")
            self.last_time = now
            self.frame_count = 0

if __name__ == "__main__":
    app = wx.App(False)
    frame = RandomWalkFrame()
    frame.Show()
    app.MainLoop()