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()

        # X Input Field
        self.x_input = QtWidgets.QLineEdit()
        layout.addWidget(self.x_input, 0, 0)

        # Y Input Field
        self.y_input = QtWidgets.QLineEdit()
        layout.addWidget(self.y_input, 0, 1)

        # Z Input Field
        self.z_input = QtWidgets.QLineEdit()
        layout.addWidget(self.z_input, 0, 2)

        # Go Button
        self.go_button = QtWidgets.QPushButton("Go")
        layout.addWidget(self.go_button, 0, 3)

        self.x_label = QtWidgets.QLabel("X Position:")
        self.x_value = QtWidgets.QLabel(self.stage.get_x_pos_str())
        layout.addWidget(self.x_label, 1, 0)
        layout.addWidget(self.x_value, 1, 1)

        self.y_label = QtWidgets.QLabel("Y Position:")
        self.y_value = QtWidgets.QLabel(self.stage.get_y_pos_str())
        layout.addWidget(self.y_label, 2, 0)
        layout.addWidget(self.y_value, 2, 1)

        self.z_label = QtWidgets.QLabel("Focus Position:")
        self.z_value = QtWidgets.QLabel(self.stage.get_z_pos_str())
        layout.addWidget(self.z_label, 3, 0)
        layout.addWidget(self.z_value, 3, 1)

        self.x_up_button = QtWidgets.QPushButton("X ←")
        self.x_down_button = QtWidgets.QPushButton("X →")
        layout.addWidget(self.x_up_button, 1, 2)
        layout.addWidget(self.x_down_button, 1, 3)

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

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

        # Enable Connection Button
        self.enable_button = QtWidgets.QPushButton("Enable Connection")
        layout.addWidget(self.enable_button, 4, 0)

        # Close Connection Button
        self.close_button = QtWidgets.QPushButton("Close Connection")
        layout.addWidget(self.close_button, 4, 1)

        # Close Connection Button
        self.homing = QtWidgets.QPushButton("Home Axis")
        layout.addWidget(self.homing, 4, 2)

        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)

        self.enable_button.clicked.connect(self.enable_connections)
        self.close_button.clicked.connect(self.close_connections)
        self.homing.clicked.connect(self.home_axis)
        self.go_button.clicked.connect(self.go_to_position)

    def go_to_position(self):
        x = float(self.x_input.text())
        y = float(self.y_input.text())
        z = float(self.z_input.text())

        self.stage.move_x_to(x)
        self.stage.move_y_to(y)
        self.stage.move_z_to(z)

        self.update_positions()

    def enable_connections(self):
        self.stage.enable_connections()
        self.update_positions()

    def close_connections(self):
        self.stage.close_connections()
        self.update_positions()

    def home_axis(self):
        self.stage.home_axis()
        self.update_positions()

    def jog_x_up(self):
        self.stage.jog_x('+')
        self.update_positions()

    def jog_x_down(self):
        self.stage.jog_x('-')
        self.update_positions()

    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.jog_y('+')
        self.update_positions()

    def jog_y_down(self):
       self.stage.jog_y('-')
       self.update_positions()

    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.jog_z('+')
        self.update_positions()

    def jog_z_down(self):
        self.stage.jog_z('-')
        self.update_positions()

    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(self.stage.get_x_pos_str())
        self.y_value.setText(self.stage.get_y_pos_str())
        self.z_value.setText(self.stage.get_z_pos_str())