import numpy as np
import pandas as pd
import pyarrow.parquet as pq

from loguru import logger
from time import perf_counter
from pathlib import Path
from PySide6 import QtWidgets

from vispy.scene import SceneCanvas, AxisWidget, Line, GridLines, Grid, ViewBox, Label, PanZoomCamera
from vispy.app import use_app

DEFAULT_COLORS = ["#790cff", "#ff0000", "#0000ff", "#a5ee30", "#f8cb14", "#ef00ff", "#51a9ff", "#119e40"]

class MainWindow(QtWidgets.QMainWindow):

    def __init__(self):
        super().__init__()
        self.setWindowTitle("Plotter")

        # Save last filepath
        self.last_filepath = ""

        # Gui Layout
        self.central_widget = QtWidgets.QWidget()
        self.main_layout = QtWidgets.QGridLayout()
        self.central_widget.setLayout(self.main_layout)

        self._canvas = CanvasWrapper()
        self.main_layout.addWidget(self._canvas.canvas.native, 0, 0)
        self.setCentralWidget(self.central_widget)

        # Toolbar
        self.upload_button = QtWidgets.QPushButton("Upload files")
        self.upload_button.clicked.connect(self.filepicker)

        self.toolbar = self.addToolBar("Controls")
        self.toolbar.setMovable(False)
        self.toolbar.setStyleSheet("QToolBar{spacing: 5px; margin: 5px; border: none; background: transparent;}")
        self.toolbar.addWidget(self.upload_button)

        # Uploaded data
        self.filepaths = []
        self.plot_widgets = []
        self.data = []

    def filepicker(self):
        files, _ = QtWidgets.QFileDialog.getOpenFileNames(
            self,
            "Select files",
            self.last_filepath,
            "Parquet Files (*.parquet)",
        )
        if len(files) > 0:
            self.last_filepath = str(Path(files[0]).parent)
            self.filepaths.extend(files)
            self.load_data(files)

    def load_data(self, new_files):
        start = perf_counter()
        for filepath in new_files:
            start2 = perf_counter()
            data_table = pq.read_table(filepath).to_pandas()
            logger.info(f"Took {perf_counter() - start2} to load pandas dataframe")
            filename = Path(filepath).stem.replace("_picolog", "")
            self._canvas.generate_new_plotview(filename, data_table)
        logger.success(f"Took {perf_counter() - start} to load and plot data")


class CanvasWrapper:

    def __init__(self):
        self.PLOT_BG_COLOR = "#FFFFFF"
        self.PLOT_TXT_COLOR = "#000000"

        self.canvas = SceneCanvas()
        self.grid: Grid = self.canvas.central_widget.add_grid(bgcolor=self.PLOT_BG_COLOR)
        self.plots = []

    def generate_new_plotview(self, file_name: str, data_table: pd.DataFrame):
        plot_grid: Grid = self.grid.add_grid()

        title = Label(str(file_name), font_size=8, color='black')
        title.height_max = 35
        plot_grid.add_widget(title, row=0, col=1)

        plot_view: ViewBox = plot_grid.add_view(row=1, col=1)
        plot_view.camera = PanZoomCamera()
        plot_view.camera.set_range(x=(0, 200), y=(-50, 50))

        x_axis = AxisWidget(orientation='bottom', axis_label='Time (ms)', axis_font_size=8, tick_font_size=5, axis_label_margin=35, tick_label_margin=15, text_color=self.PLOT_TXT_COLOR)
        x_axis.height_max = 60
        plot_grid.add_widget(x_axis, row=2, col=1)
        x_axis.link_view(plot_view)

        y_axis1 = AxisWidget(orientation='left', axis_label='Voltage (V)', axis_font_size=8, tick_font_size=5, axis_label_margin=60, tick_label_margin=5, text_color=self.PLOT_TXT_COLOR)
        y_axis1.width_max = 85
        plot_grid.add_widget(y_axis1, row=1, col=0)
        y_axis1.link_view(plot_view)

        y_axis2 = AxisWidget(orientation='right', axis_label='Current (A)', axis_font_size=8, tick_font_size=5, axis_label_margin=60, tick_label_margin=5, text_color=self.PLOT_TXT_COLOR)
        y_axis2.width_max = 85
        plot_grid.add_widget(y_axis2, row=1, col=2)
        y_axis2.link_view(plot_view)

        legend: Grid = plot_grid.add_grid(row=1, col=3)
        legend.width_max = 170

        colors = iter(DEFAULT_COLORS)
        legend_row = 0
        for column in data_table:
            line_color = next(colors)
            data = np.array([data_table[column].index, data_table[column]]).transpose()
            Line(data, parent=plot_view.scene, color=line_color, width=2)
            legend_color = Label("——", color=line_color, bold=True, font_size=10, anchor_x='center', face="Arial")
            legend_color.height_max = 20
            legend_color.width_max = 50
            legend_text = Label(data_table[column].name, color=self.PLOT_TXT_COLOR, font_size=8, anchor_x='left')
            legend_text.height_max = 20
            legend_text.width_max = 10
            legend.add_widget(legend_color, row=legend_row, col=0)
            legend.add_widget(legend_text, row=legend_row, col=1)
            legend_row += 1

        gridlines = GridLines(scale=(0.1, 0.1), color=(0, 0, 0, 0.5))
        gridlines.set_gl_state('translucent')
        plot_view.add(gridlines)

        self.plots.append(plot_grid)
        self.grid.next_row()


if __name__ == "__main__":
    app = use_app("pyside6")
    app.create()
    window = MainWindow()
    window.setMinimumSize(1200, 600)
    window.show()
    app.run()