#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

import sys
import zmq
import time
import math
import signal
from collections import deque
from threading import Thread
import flatbuffers
#generated by flatc
import topics.nav.imu
import topics.motor.command

from minnow_low_level_control.HeadingControl import *   
from minnow_low_level_control.SurgeSpeedControl import *  

class Subscriber(Thread):
    def __init__(self, context, id, topics_callbacks):
        super().__init__()
        self.context = context
        self.topics_callbacks = topics_callbacks
        self.loop = False

    def run(self):
        print('starting subscriber, on topic(s) {}'.format(self.topics_callbacks.keys()))
        subscriber = self.context.socket(zmq.SUB)
        subscriber.connect("tcp://127.0.0.1:5555")
        for topic in self.topics_callbacks.keys():
            subscriber.setsockopt_string(zmq.SUBSCRIBE, topic)
        poller = zmq.Poller()
        poller.register(subscriber, zmq.POLLIN)
        self.loop = True
        while self.loop:
            evts = poller.poll(1000)
            if evts:
                message = subscriber.recv()
                for topic in self.topics_callbacks.keys():
                    topic_name_msg = message.split(None,1)
                    topic_name = topic_name_msg[0]
                    topic_msg = topic_name_msg[1]
                    if topic == topic_name.decode('utf-8'):
                        self.topics_callbacks[topic](topic_msg)

    def stop(self):
        self.loop = False

class Publisher:
    def __init__(self):
        signal.signal(signal.SIGINT, self.exit_signal)
        self.zmq_context = zmq.Context()
        self.setup_subscriber()
        self.nav_imu_msg = None
        #setup flatbuffers
        self.fb_builder = flatbuffers.Builder(1024)

        # This is for standalone troubleshooting of the python code ---------------------
        self.desired_speed = 0.0
        self.desired_heading = 270     # between 0 - 360
        self.desired_pitch = 0.0
        
        self.current_speed = 0.0
        # self.current_heading = 170
        self.current_pitch = -2.0
        self.contrl_desired_heading = 0
        if (self.desired_heading >= 0) and (self.desired_heading <= 180):
            self.contrl_desired_heading = self.desired_heading
        else:
            self.contrl_desired_heading = self.desired_heading - 360
        # -------------------------------------------------------------------------------
        self.speed_control_system = speed_controller()
        self.heading_control_system = heading_controller()

    def exit_signal(self, sig, frame):
        print('You pressed Ctrl+C!')
        self.subscriber.stop()
        time.sleep(1)
        self.subscriber.join()
        sys.exit(0)

    def setup_subscriber(self):
        self.subscriber = Subscriber(self.zmq_context, 0, {'nav.imu':self.nav_imu_callback})
        self.subscriber.start()

    def nav_imu_callback(self, msg):
        self.nav_imu_msg = topics.nav.imu.imu.GetRootAsimu(msg, 0)

    def run(self):
        socket = self.zmq_context.socket(zmq.PUB)
        socket.connect("tcp://127.0.0.1:5556")

        count = 0
        while True:
            count += 1
            try:
                current_heading = self.nav_imu_msg.Yaw()
            except:
                pass
            print('current hdg',current_heading)
            if (current_heading >= 0) and (current_heading <= 180):
                contrl_current_heading = current_heading
            else:
                contrl_current_heading = current_heading - 360
            print('cntrl desi hdg',self.contrl_desired_heading)
            self.heading_control_system.DesiredHeading(self.contrl_desired_heading)

            # Run speed controller
            (speed_contrl_thrust)=self.speed_control_system.update(self.desired_speed)
            print("Speed Control Thrust Output: %f" % (speed_contrl_thrust))
            # Run heading controller
            (hdg_differential_thrust,hdg_port_thrust,hdg_stbd_thrust)=self.heading_control_system.update(contrl_current_heading,speed_contrl_thrust)
            #print(hdg_differential_thrust)
            print("Heading control port thrust: %f" % hdg_port_thrust)
            print("Heading control stbd thrust: %f" % hdg_stbd_thrust)

            topics.motor.command.commandStart(self.fb_builder)
            topics.motor.command.commandAddTime(self.fb_builder, time.time())
            topics.motor.command.commandAddMotor1Command(self.fb_builder, hdg_port_thrust)
            topics.motor.command.commandAddMotor2Command(self.fb_builder, hdg_stbd_thrust)
            motor_msg = topics.motor.command.commandEnd(self.fb_builder)
            self.fb_builder.Finish(motor_msg)
            bin_motor_msg = self.fb_builder.Output()
            socket.send(b'motor.command' + b' ' + bin_motor_msg)

            time.sleep(0.2)

if __name__ == "__main__":
    pub = Publisher()
    pub.run()
