import pygame
import serial
import re
import numpy as np
import time

from numpy import array
from pyquaternion import Quaternion
from math import cos, sin, pi

from spatz.sensors.imu.wsen_isds import WSEN_ISDS_GYRO
from spatz.simulation import Simulation
from spatz.simulations.rocketpy import RocketPyCSV

# Blatantly stolen from: https://stackoverflow.com/questions/21019471/how-can-i-draw-a-3d-shape-using-pygame-no-other-modules

X, Y, Z = 0, 1, 2


def rotation_matrix(a, b, by):
    """
    rotation matrix of a, b, by radians around x, y, z axes (respectively)
    """
    sa, ca = sin(a), cos(a)
    sb, cb = sin(b), cos(b)
    sby, cby = sin(by), cos(by)
    return (
        (cb*cby, -cb*sby, sb),
        (ca*sby + sa*sb*cby, ca*cby - sby*sa*sb, -cb*sa),
        (sby*sa - ca*sb*cby, ca*sby*sb + sa*cby, ca*cb)
    )


class Physical:
    def __init__(self, vertices, edges, colors):
        """
        a 3D object that can rotate around the three axes
        :param vertices: a tuple of points (each has 3 coordinates)
        :param edges: a tuple of pairs (each pair is a set containing 2 vertices' indexes)
        """
        self.__vertices = array(vertices)
        self.__edges = tuple(edges)
        self.__colors = tuple(colors)
        self.__rotation = Quaternion(axis=(1, 0, 0), angle=0) # radians around each axis

    def rotate(self, quaternion):
        self.__rotation = quaternion

    @property
    def lines(self):
        location = array([self.__rotation.rotate(vertex) for vertex in self.__vertices])  # an index->location mapping
        return (((location[v1], location[v2]), color) for (v1, v2), color in zip(self.__edges, self.__colors))
    

BLACK, RED, GREEN, BLUE = (0, 0, 0), (255, 128, 128), (128, 255, 128), (128, 128, 255)

LIGHTRED, LIGHTGREEN, LIGHTBLUE = (128, 64, 64), (64, 128, 64), (64, 64, 128)


class Paint:
    def __init__(self, shape, shape2):
        self.__shape = shape
        self.__shape2 = shape2
        self.__size = 900, 450
        self.__clock = pygame.time.Clock()
        self.__screen = pygame.display.set_mode(self.__size)

    def __fit(self, vec):
        """
        ignore the z-element (creating a very cheap projection), and scale x, y to the coordinates of the screen
        """
        # notice that len(self.__size) is 2, hence zip(vec, self.__size) ignores the vector's last coordinate
        return [round(70 * coordinate + frame / 2) for coordinate, frame in zip(vec, self.__size)]

    def __draw_shape(self, thickness=4):
        for (start, end), color in self.__shape.lines:
            pygame.draw.line(self.__screen, color, self.__fit((start[0]-2, start[1], start[2])), self.__fit((end[0]-2, end[1], end[2])), thickness)

        for (start, end), color in self.__shape2.lines:
            pygame.draw.line(self.__screen, color, self.__fit((start[0]+2, start[1], start[2])), self.__fit((end[0]+2, end[1], end[2])), thickness)

    def draw(self):
        self.__screen.fill(BLACK)
        self.__draw_shape()
        pygame.display.flip()
        self.__clock.tick(40)          


def main():
    from pygame import K_q, K_w, K_a, K_s, K_z, K_x

    rotation = Quaternion(axis=[1, 0, 0], angle=pi/2)

    axes = Physical(
        vertices=((0, 0, 0), (0, 0, 2), (0, 2, 0), (2, 0, 0)),
        edges=({0, 1}, {0, 2}, {0, 3}),
        colors=(BLUE, GREEN, RED)
    ) 

    truth = Physical(
        vertices=((0, 0, 0), (0, 0, 2), (0, 2, 0), (2, 0, 0)),
        edges=({0, 1}, {0, 2}, {0, 3}),
        colors=(LIGHTBLUE, LIGHTGREEN, LIGHTRED)
    ) 
        
    pygame.init()
    pygame.display.set_caption("Simulation")
    renderer = Paint(axes, truth)

    simulation = Simulation().load(RocketPyCSV('nominal_wind.csv'))

    gyro = simulation.add_sensor(WSEN_ISDS_GYRO)
    offset = gyro.calibrate(100)
    att = simulation.add_observer([' e0', ' e1', ' e2', ' e3'])
    dt = 0.01

    quat = np.array([-0.706864,0.018510,0.018510,-0.706864])

    while True:
        time.sleep(0.01)
        simulation.advance(dt)
        omegas = gyro() - offset
        true = att()

        omegas = (omegas / 1000) * np.pi / 180

        matrix = np.array([
            [1, -dt/2*omegas[0], -dt/2*omegas[1], -dt/2*omegas[2]],
            [dt/2*omegas[0], 1, dt/2*omegas[2], -dt/2*omegas[1]],
            [dt/2*omegas[1], -dt/2*omegas[2], 1, dt/2*omegas[0]],
            [dt/2*omegas[2], dt/2*omegas[1], -dt/2*omegas[0], 1]
        ])

        quat = matrix @ quat
        quat /= np.linalg.norm(quat)


        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                exit()
            
        axes.rotate(Quaternion(x=quat[0], y=quat[1], z=quat[2], w=quat[3]))
        truth.rotate(Quaternion(x=true[0], y=true[1], z=true[2], w=true[3]))

        renderer.draw()

if __name__ == '__main__':
    main()