import wx
import matplotlib.pyplot as plt
from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg, FigureCanvasWxAgg


class PlotFrame(wx.Frame):
    def __init__(self, parent, *args, **kwargs):
        wx.Frame.__init__(self, parent, *args, **kwargs)

        self.fig = plt.figure(figsize=(5, 4), dpi=60)
        self.subplot1 = self.fig.add_subplot(111)
        self.canvas = FigureCanvasWxAgg(self, wx.ID_ANY, self.fig)

        self.toolbar = NavigationToolbar2WxAgg(self.canvas)

        sizer = wx.BoxSizer(wx.VERTICAL)
        sizer.Add(self.canvas, 1, wx.EXPAND)

        # Add more wx controls to the same area as the toolbar
        export_data_btn = wx.Button(self, label="Export Data")
        some_other_control = wx.TextCtrl(self, value="TextCtrl")

        tool_sizer = wx.BoxSizer(wx.HORIZONTAL)

        # Option 1: Allowing the toolbar to expand to fill the empty space doesn't display correctly
        tool_sizer.Add(self.toolbar, 1, wx.EXPAND)
        tool_sizer.Add((10, -1), 0, wx.EXPAND)

        # Option 2: Fixed toolbar size and expandable empty space doesn't display correctly either
        # tool_sizer.Add(self.toolbar, 0, wx.EXPAND)
        # tool_sizer.Add((10, -1), 1, wx.EXPAND)

        # doesn't matter whether these are added or not, just wanted to highlight the use case
        tool_sizer.Add(export_data_btn, 0, wx.ALIGN_CENTER_VERTICAL | wx.RIGHT, 5)  
        tool_sizer.Add(some_other_control, 0, wx.ALIGN_CENTER_VERTICAL | wx.RIGHT, 5)

        sizer.Add(tool_sizer, 0, wx.EXPAND)

        self.SetSizer(sizer)
        self.Layout()
        self.CenterOnScreen()

        # self.toolbar.Realize()  # trying to realize the toolbar after layout is complete also doesn't work


if __name__ == '__main__':
    app = wx.App(False)
    frame = PlotFrame(None, title='Test Frame', size=(800, 600))
    app.SetTopWindow(frame)
    frame.Center()
    frame.Show()
    app.MainLoop()