import cv2
import time
import numpy as np
from math import atan2, pi, cos, sin
import serial
from threading import Thread

from ICM20948 import ICM20948, ICM20948DataStream


def remap(value, maxInput, minInput, maxOutput, minOutput):
	value = maxInput if value > maxInput else value
	value = minInput if value < minInput else value
	inputSpan = maxInput - minInput
	outputSpan = maxOutput - minOutput
	scaledThrust = float(value - minInput) / float(inputSpan)
	return minOutput + (scaledThrust * outputSpan)


class Compass_BNO055:
    def __init__(self):
        self.ser = serial.Serial('/dev/cu.SLAB_USBtoUART', baudrate=115200)
        self.buffer = {'heading_bno': 0, 'heading_calc': 0, 'heading_calc_smooth': 0, 'mag_vals': {'x': 0, 'y': 0, 'z': 0}}
        self.running = False
        time.sleep(1)


    def start(self):
        t = Thread(target=self.update, args=())
        t.daemon = True
        self.running = True
        t.start()
        self.start_time = time.time()
        return self

    def calculate_heading(self, mag_x, mag_y):
        heading = 180 * atan2(mag_y, mag_x)/pi - 90
        if heading < 0:
            heading += 360
        return heading


    def update(self):
        while self.running:
            try:
                self.ser.flushInput()
                ser_bytes = str(self.ser.readline())
                useful = ser_bytes.split('\\')[0].split('\'')[1].split('|')
                heading = float(useful[1])
                mag_data_raw = useful[3].split(';')
                mag_raw_x = float(mag_data_raw[0])
                mag_raw_y = float(mag_data_raw[1])
                mag_raw_z = float(mag_data_raw[2])

                mag_raw_x_infl = int(remap(mag_raw_x, 30, -30, 0, 400))
                mag_raw_y_infl = int(remap(mag_raw_y, 30, -30, 0, 400))
                mag_raw_z_infl = int(remap(mag_raw_z, 30, -30, 0, 400))

                if abs(mag_raw_x) > 30 or abs(mag_raw_y) > 30:
                    print(f"Got {mag_raw_x} | {mag_raw_y}, skipping")
                else:
                    heading_bno = heading
                    heading_calc = self.calculate_heading(mag_raw_x, mag_raw_y)
                    # print()
                    # print(heading_calc_smooth)
                    # print(delta)
                    # print()

                    self.buffer = {'heading_bno': heading_bno, 'heading_calc': heading_calc, 'heading_calc_smooth': heading_calc_smooth, 'mag_vals': {'x': mag_raw_x_infl, 'y': mag_raw_y_infl, 'z': mag_raw_z_infl}}
            except Exception as e:
                print(f"Error: {e}")

    def read(self):
        return self.buffer

    def stop(self):
        self.running = False



points = []

def draw_frame(heading_bno, heading_calc, heading_calc_smooth, mag_x, mag_y, draw_bno=True, draw_smooth=True):
    global points
    print(f"heading_bno: {heading_bno} | heading_calc: {heading_calc} | heading_calc_smooth: {heading_calc_smooth} | mag_x: {mag_x} | mag_y: {mag_y} | ")
    frame = np.full(shape=[400, 400, 3], fill_value=255, dtype=np.uint8)
    length = 150

    if draw_bno:
        bno_P1 = (200, 200)
        bno_P2_0 = int(bno_P1[0] + length * cos(heading_bno * pi / 180.0))
        bno_P2_1 = int(bno_P1[1] + length * sin(heading_bno * pi / 180.0))
        bno_P2 = (bno_P2_0, bno_P2_1)

    calc_P1 = (200, 200)
    calc_P2_0 = int(calc_P1[0] + length * cos(heading_calc * pi / 180.0))
    calc_P2_1 = int(calc_P1[1] + length * sin(heading_calc * pi / 180.0))
    calc_P2 = (calc_P2_0, calc_P2_1)

    if draw_smooth:
        s_calc_P1 = (200, 200)
        s_calc_P2_0 = int(s_calc_P1[0] + length * cos(heading_calc_smooth * pi / 180.0))
        s_calc_P2_1 = int(s_calc_P1[1] + length * sin(heading_calc_smooth * pi / 180.0))
        s_calc_P2 = (s_calc_P2_0, s_calc_P2_1)

    if draw_bno:
        frame = cv2.line(frame,bno_P1,bno_P2,(255,0,0),2)
    frame = cv2.line(frame,calc_P1,calc_P2,(0,255,0),2)
    if draw_smooth:
        frame = cv2.line(frame,s_calc_P1,s_calc_P2,(0,90,0),2)
    print(mag_x, mag_y)
    frame = cv2.circle(frame, (mag_x, mag_y), radius=1, color=(0, 255, 0), thickness=1)

    for p in points:
        frame = cv2.circle(frame, (p[0], p[1]), radius=1, color=(0, 0, 255), thickness=1)
    points.append((mag_x, mag_y))
    frame = cv2.rotate(frame, cv2.ROTATE_90_COUNTERCLOCKWISE)

    font = cv2.FONT_HERSHEY_SIMPLEX
    org_n = (200, 30)
    org_s = (200, 370)
    org_w = (30, 200)
    org_e = (370, 200)
    fontScale = 1
    color = (255, 0, 0)
    thickness = 2
    frame = cv2.putText(frame, 'N', org_n, font, fontScale, color, thickness, cv2.LINE_AA)
    frame = cv2.putText(frame, 'S', org_s, font, fontScale, color, thickness, cv2.LINE_AA)
    frame = cv2.putText(frame, 'W', org_w, font, fontScale, color, thickness, cv2.LINE_AA)
    frame = cv2.putText(frame, 'E', org_e, font, fontScale, color, thickness, cv2.LINE_AA)

    cv2.imshow('direction', frame)



# compass = Compass_BNO055().start()
sensor = ICM20948DataStream().start()
compass = ICM20948(sensor)
compass.calibrate()

while True:
    heading, data = compass.get_heading_and_data()
    smooth_heading = compass.get_smooth_heading()
    mag_raw_x_infl = int(remap(data[0], 1, -1, 400, 0))
    mag_raw_y_infl = int(remap(data[1], 1, -1, 400, 0))
    print(f"Heading: {heading}, Smooth: {smooth_heading}, X: {mag_raw_x_infl}, Y: {mag_raw_y_infl}")
    draw_frame(0, heading, smooth_heading, mag_raw_x_infl, mag_raw_y_infl, False, True)
    cv2.waitKey(1)
    time.sleep(0.05)


# while True:
#     new_data = compass.read()
#     # heading_bno = new_data['heading_bno']
#     heading_calc = new_data['heading_calc']
#     heading_calc_smooth = new_data['heading_calc_smooth']
#     mag_x = new_data['mag_vals']['x']
#     mag_y = new_data['mag_vals']['y']
#     # draw_frame(heading_bno, heading_calc, heading_calc_smooth, mag_x, mag_y)
#     draw_frame(heading_bno, heading_calc, heading_calc_smooth, mag_x, mag_y, False, False)
#     cv2.waitKey(1)
#     time.sleep(0.05)