from PyQt5 import QtWidgets, QtCore
from pymeasure.display.widgets import TabWidget

from src.classes.stage import Stage

class ControlTab(TabWidget, QtWidgets.QWidget):

    stage = Stage()

    def __init__(self, name, parent=None):
        super().__init__(parent)
        self.name = name

        layout = QtWidgets.QGridLayout()

        self.x_label = QtWidgets.QLabel("X Position:")
        self.x_value = QtWidgets.QLabel(f"{self.stage.get_x_pos():.{3}f}mm")
        layout.addWidget(self.x_label, 0, 0)
        layout.addWidget(self.x_value, 0, 1)

        self.y_label = QtWidgets.QLabel("Y Position:")
        self.y_value = QtWidgets.QLabel(f"{self.stage.get_y_pos():.{3}f}mm")
        layout.addWidget(self.y_label, 1, 0)
        layout.addWidget(self.y_value, 1, 1)

        self.z_label = QtWidgets.QLabel("Z Position:")
        self.z_value = QtWidgets.QLabel(f"{self.stage.get_z_pos():.{3}f}mm")
        layout.addWidget(self.z_label, 2, 0)
        layout.addWidget(self.z_value, 2, 1)

        self.x_up_button = QtWidgets.QPushButton("X ↑")
        self.x_down_button = QtWidgets.QPushButton("X ↓")
        layout.addWidget(self.x_up_button, 0, 2)
        layout.addWidget(self.x_down_button, 0, 3)

        self.y_up_button = QtWidgets.QPushButton("Y ↑")
        self.y_down_button = QtWidgets.QPushButton("Y ↓")
        layout.addWidget(self.y_up_button, 1, 2)
        layout.addWidget(self.y_down_button, 1, 3)

        self.z_up_button = QtWidgets.QPushButton("Z ↑")
        self.z_down_button = QtWidgets.QPushButton("Z ↓")
        layout.addWidget(self.z_up_button, 2, 2)
        layout.addWidget(self.z_down_button, 2, 3)

        self.setLayout(layout)

        self.x_up_button.pressed.connect(self.jog_x_up)
        self.x_up_button.released.connect(self.stop_jogs)
        self.x_up_button.clicked.connect(self.increment_x_up)
        self.x_down_button.pressed.connect(self.jog_x_down)
        self.x_down_button.released.connect(self.stop_jogs)
        self.x_down_button.clicked.connect(self.increment_x_down)

        self.y_up_button.pressed.connect(self.jog_y_up)
        self.y_up_button.released.connect(self.stop_jogs)
        self.y_up_button.clicked.connect(self.increment_y_up)
        self.y_down_button.pressed.connect(self.jog_y_down)
        self.y_down_button.released.connect(self.stop_jogs)
        self.y_down_button.clicked.connect(self.increment_y_down)

        self.z_up_button.pressed.connect(self.jog_z_up)
        self.z_up_button.released.connect(self.stop_jogs)
        self.z_up_button.clicked.connect(self.increment_z_up)
        self.z_down_button.pressed.connect(self.jog_z_down)
        self.z_down_button.released.connect(self.stop_jogs)
        self.z_down_button.clicked.connect(self.increment_z_down)

    def jog_x_up(self):
        self.stage.motor_x.jog('+')

    def jog_x_down(self):
        self.stage.motor_x.jog('-')

    def increment_x_up(self):
        self.stage.move_x(1)
        self.update_positions()
    
    def increment_x_down(self):
        self.stage.move_x(-1)
        self.update_positions()

    def jog_y_up(self):
        self.stage.motor_y.jog('+')

    def jog_y_down(self):
       self.stage.motor_y.jog('-')

    def increment_y_up(self):
        self.stage.move_y(1)
        self.update_positions()
    
    def increment_y_down(self):
        self.stage.move_y(-1)
        self.update_positions()

    def jog_z_up(self):
        self.stage.motor_z.jog('+')

    def jog_z_down(self):
        self.stage.motor_z.jog('-')

    def increment_z_up(self):
        self.stage.move_z(1)
        self.update_positions()
    
    def increment_z_down(self):
        self.stage.move_z(-1)
        self.update_positions()

    def stop_jogs(self):
        self.stage.stop_all()
        self.update_positions()

    def update_positions(self):
        self.x_value.setText(f"{self.stage.get_x_pos():.{3}f}mm")
        self.y_value.setText(f"{self.stage.get_y_pos():.{3}f}mm")
        self.z_value.setText(f"{self.stage.get_z_pos():.{3}f}mm")