controller

#!/usr/bin/env python3
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QVBoxLayout, QWidget
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QTimer, QThread, pyqtSignal
import numpy as np
import time
import sys
import zmq
from VideoThread import VideoThread
import socket as sock

running = True
rate_limit = 20.0
steering_rate = 40.0

# Client (e.g., on your laptop)
context = zmq.Context()
socket = context.socket(zmq.REQ)
socket.connect("tcp://tank:5555")
socket.setsockopt(zmq.RCVTIMEO, 1000)

# Get our IP address
def get_our_ip():
    with sock.socket(sock.AF_INET, sock.SOCK_DGRAM) as s:
        s.connect(('tank', 5000))
        return s.getsockname()[0]

our_ip = get_our_ip()

def send_message(message):
    socket.send_string(message)
    try:
        reply = socket.recv_string()
    except zmq.error.Again:
        print("Timeout")

controls = [Qt.Key.Key_Up, Qt.Key.Key_Down, Qt.Key.Key_Left, Qt.Key.Key_Right, Qt.Key.Key_Space]
keymap = {Qt.Key.Key_Space : 1,
          Qt.Key.Key_Left  : 2,
          Qt.Key.Key_Right : 4,
          Qt.Key.Key_Up    : 8,
          Qt.Key.Key_Down  : 16 }



class CameraControlWindow(QMainWindow):
    def __init__(self, url):
        super().__init__()
        print(f'Attempting to connect to tank from {our_ip}')
        send_message(f'connect:{our_ip}')
        print('Connected to tank')
        self.video_thread = VideoThread(url)
        self.video_thread.update_frame_signal.connect(self.display_frame)

        self.init_ui()
        self.video_thread.start()

        # Timer to run the control loop more frequently
        self.control_timer = QTimer()
        self.control_timer.timeout.connect(self.control_update)
        self.control_timer.start(10)  # 100 Hz control loop

        self.target = np.array([0.0, 0.0])
        self.speed = np.array([0.0, 0.0])
        self.acceleration = 50  # Example acceleration value
        self.keystate = 0
        self.lastcommand = None

        # Initialize time
        self.previous_time = time.perf_counter()

    def keyPressEvent(self, event):
        # Remove after debug
        if event.key() == Qt.Key.Key_Q:
            sys.exit(0)

        if event.key() in controls and not event.isAutoRepeat():
            self.handle_inputs('pressed', event.key())

    def keyReleaseEvent(self, event):
        if event.key() in controls and not event.isAutoRepeat():
            self.handle_inputs('released', event.key())

    def handle_inputs(self, action, key):
        if action == "pressed":
            self.keystate |= keymap[key]
        elif action == "released":
            self.keystate &= (0xFF & ~keymap[key])

        # coasting is allowed but any keys pressed sets a new goal
        self.target = np.array([0.0,0.0])
        if self.keystate & keymap[Qt.Key.Key_Up]:
            self.target += np.array([100.0,100.0])
        if self.keystate & keymap[Qt.Key.Key_Down]:
            self.target += np.array([-100.0,-100.0])

        # Motor drive is correct but left/right are swapped
        # at the motor driver - swap them here
        if self.keystate & keymap[Qt.Key.Key_Left]:
            self.target += np.array([50.0,-50.0])
        if self.keystate & keymap[Qt.Key.Key_Right]:
            self.target += np.array([-50.0,50.0])

        # braking is immediate
        if self.keystate & keymap[Qt.Key.Key_Space]:
            self.target = self.speed = np.array([0.0,0.0])

        scale = np.max(np.abs(self.target))
        # handle additive overrange
        if scale > 100:
            self.target = self.target * 100.0 / scale

    def control_update(self):
        now = time.perf_counter()
        delta_time = now - self.previous_time
        self.previous_time = now

        # Calculate the desired change in speed
        delta = self.target - self.speed
        # Calculate the change in speed allowed by the rate limit
        allowed_delta = np.clip(delta, -rate_limit * delta_time, rate_limit * delta_time)

        # Update speed by applying the allowed change
        self.speed += allowed_delta

        command = 'drive:' + ','.join(np.round(self.speed).astype(int).astype(str))
        if self.lastcommand != command:
            send_message(command)
        self.lastcommand = command

    def init_ui(self):
        # Get screen resolution
        screen = QApplication.primaryScreen()
        screen_rect = screen.availableGeometry()

        self.setWindowTitle("Front Camera")
        self.showFullScreen()  # Set the window to fullscreen

        self.image_label = QLabel(self)
        self.image_label.setGeometry(0, 0, screen_rect.width(), screen_rect.height())
        self.image_label.setAlignment(Qt.AlignCenter)

        central_widget = QWidget()
        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        central_widget.setLayout(layout)

        self.setCentralWidget(central_widget)

    def display_frame(self, qt_image):
        self.image_label.setPixmap(QPixmap.fromImage(qt_image))

    def closeEvent(self, event):
        send_message(f'disconnect:{our_ip}')
        self.video_thread.stop()
        event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CameraControlWindow("udp://0.0.0.0:5000")
    window.show()
    sys.exit(app.exec_())