"""
作用:
    模型重构
时间:
    2021年6月19日22:59:46
命名规则:
    模块名: 小写+下划线
    类名: 驼峰式 MyClass
    常量: 全大写
    函数名,普通变量名: 小写+下划线
    私有函数: 下划线开头
"""
import numpy as np
import matplotlib.pyplot as plt
import pylab as pl

import calculate as cal
import matplotlib
import matplotlib.animation as animation
import pprint
import warnings
import pandas as pd
import copy

class Vehicle:
    def __init__(self, **kwargs):
        # self._settings_expected = ['name','vehicle_mass','vehicle_width','vehicle_longth','vehicle_moment_of_inertia','tire_mass','tire_radius','tire_moment_of_inertia',]
        self.name = None
        self.vehicle_width = 0
        self.vehicle_longth = 0
        self.vehicle_moment_of_inertia = 0
        self.tire_mass = 0
        self.tire_radius = 0
        self.tire_moment_of_inertia = 0
        self.vehicle_mass = 0
        self.vehicle_barycenter_velocity = np.zeros(3)  # (x,y,z)
        self.vehicle_angular_velocity = np.zeros(3)
        self.vehicle_euler_angles = np.zeros(3) # (ZXZ)
        self.vehicle_barycenter_position = np.zeros(3) # (x,y,z)  # position of barycenter
        self.vehicle_direction_vector = np.zeros(3)
        self.tire_angular_velocity = np.zeros((4,3)) # 车身参考系  (a,b,c)  # omegas
        self.__dict__ = {**self.__dict__, **kwargs}  # 导入
        self.vehicle_moment_of_inertia = cal.Momentum_of_Inertia(
            self.vehicle_longth,
            self.vehicle_width,
            self.vehicle_mass
        )
        self.tire_moment_of_inertia = 30 * self.tire_radius ** 2

        print("Success for %s" % self.name)
    """
        def check(self):
        settings_expected = self._settings_expected
        number_errors = 0
        information_missed = []
        for i in settings_expected:
            if i not in self.__dict__.keys():
                information_missed.append(i)
        if len(information_missed) == 0:
            print("Load information success")
        else:
            print("信息不全, 缺失%d个信息" % len(information_missed))
            print(information_missed)
    """


    def info(self):
        print("车辆信息:")
        pprint.pprint(self.__dict__)

    def info_template(self):
        """
        导出一个模板, 便于输入参数
        :return:
        """
        pass
    def load_infomations(self,**kwargs):
        """
        导入数据, 完成各种形式
        :param kwargs:
        :return:
        """
        pass



class Garage:
    """
    鸡肋鸡肋, 基类? 鸡类,
    """
    def __init__(self):
        self.name = None
        self.vehicles = []
    def add(self,vehicle):
        self.vehicles.append(vehicle)

    def info(self):
        for vehicle in self.vehicles:
            print(vehicle.name)

    def query_by_vehicle_name(self,name):
        for vehicle in self.vehicles:
            if vehicle.name == name:
                vehicle.info()
                return True
        print("没找着%s" % name)


class Environment():
    def __init__(self,**kwargs):
        self.duration = 0 # duration after hit (second)
        self.interval = 0 # time of every tick
        self.gravity = 0
        self.road_coefficient_of_sliding_friction = 0
        self.road_longth = 0
        self.road_width = 0
        self.road_tire_rolling_friction_coefficient = 0
        self.coefficient_of_air_resistance = 0
        self.__dict__ = {**self.__dict__, **kwargs}

    def info(self):
        print("环境信息:")
        pprint.pprint(self.__dict__)

class Recorder:
    def __init__(self,car):
        self.object = car

    def start_record(self,N,t):
        record_tmp = []
        self.times = N
        self.interval = t
        print("开始记录:", "次数%d 时间间隔%5fs, 总时长%.1fs" % (N,t,N*t))
        for _ in range(N):
            self.object.run(t=t)
            tmp = copy.copy(self.object)
            record_tmp.append(tmp.__dict__)  # 将类转换为对应属性名与属性值的字典
        self.records = pd.DataFrame(record_tmp)

    def plot_barycenter_velocity(self):
        N = len(self.records.V)
        Vx = np.zeros(N)
        Vy = np.zeros(N)
        Vnorm = np.zeros(N)
        for i in range(N):
            v = self.records.V[i]
            Vx[i] = v[0]
            Vy[i] = v[1]
            Vnorm[i] = np.linalg.norm(v)
        fig, axes = plt.subplots(3,1, sharex='col')
        plt.style.use('seaborn')
        fig.suptitle('Velocity')
        axes[0].plot(range(1, N + 1), Vx)
        axes[1].plot(range(1, N + 1), Vy)
        axes[2].plot(range(1, N + 1), Vnorm)
        axes[2].set_xlabel('time')
        axes[0].set_ylabel('X')
        axes[1].set_ylabel('Y')
        axes[2].set_ylabel('Norm')


    def plot_car_angular_velocity(self):
        N = len(self.records.OMEGA)
        Y = np.zeros(N)
        for i in range(N):
            Y[i] = self.records.OMEGA[i]
        plt.figure()

        plt.title("Car's angular_velocity")
        plt.plot(range(1,N+1), Y)
        plt.xlabel('time(s)')
        plt.ylabel('angular_velocity (rad/s)')
        plt.style.use('seaborn')
        plt.show()

class Collision:
    def __init__(self):
        self.name = None
        self.vehicles = []
        self.environment = None
        pass

    def add_vehicle(self,vehicle):
        """
        导入车辆
        :param vehicle:
        :return:
        """
        self.vehicles.append(vehicle)

    def set_environment(self,environment):
        """
        设置碰撞时的环境
        :param environment:
        :return:
        """
        self.environment = environment

    def info(self):
        pprint.pprint(self.vehicles)
        pprint.pprint(self.environment)

    def calculate_velocity(self,vehicle):
        # 这个矩阵左乘一个向量x在车身坐标系的坐标, 可得到这个向量在地面坐标系的坐标
        rotate_matrix_bary2terre = cal.euler2mat(vehicle.vehicle_euler_angles).T  #??
        V = vehicle.vehicle_barycenter_velocity  # 速度(Vx,Vy,Vz)
        # 计算四个轮胎与地面接触点的相对速度, 地面参考系
        O = vehicle.vehicle_barycenter_position # O 为 质心在地面参考系中的坐标
        # M 为轮胎中心, 从右前轮开始绕车逆时针编号0,1,2,3
        L = vehicle.vehicle_longth
        W = vehicle.vehicle_width
        OM = np.zeros((4,3))
        OM[0] = np.array([W / 2, L / 2, 0])
        OM[1] = np.array([-W / 2, L / 2, 0])
        OM[2] = np.array([-W / 2, -L / 2, 0])
        OM[3] = np.array([W / 2, -L / 2, 0])

        OM[0] = np.dot(rotate_matrix_bary2terre, OM[0])
        OM[1] = np.dot(rotate_matrix_bary2terre, OM[1])
        OM[2] = np.dot(rotate_matrix_bary2terre, OM[2])
        OM[3] = np.dot(rotate_matrix_bary2terre, OM[3])

        velocity_M = np.zeros((4,3))
        velocity_contact_point = np.zeros((4,3))
        "Varignon  première fois  V_M = V_O + OMEGA ∧ OM"
        for k in range(4):
            velocity_M[k] = vehicle.vehicle_barycenter_velocity + \
            np.cross(vehicle.vehicle_angular_velocity,OM[k])

        "Varignon deuxièmre fois V_contact = V_M + omega ∧ M-pointContact"
        for k in range(4):
            tire_omega = np.dot(rotate_matrix_bary2terre, vehicle.tire_angular_velocity[k])
            M_contact = vehicle.tire_radius * np.array([0,0,-1])
            velocity_contact_point[k] = velocity_M[k] + np.cross(tire_omega,M_contact)
        # 返回四个轮胎与地面接触点在地面参考系中的速度(x,y,z) * 4
        vehicle.OM = OM
        vehicle.velocity_contact_point =  velocity_contact_point
        return velocity_contact_point

    def calculate_force(self,vehicle):
        tire_downforce = 1/4 * vehicle.vehicle_mass * self.environment.gravity
        sliding_friction_module = tire_downforce * self.environment.road_coefficient_of_sliding_friction
        # 与接触点在地面参考系中的速度方向相反
        sliding_friction = np.zeros((4,3))
        for k in range(4):
            V = vehicle.velocity_contact_point[k]
            direction = -1 * V / np.linalg.norm(V)
            sliding_friction[k] = sliding_friction_module * direction
        rolling_friction_module = vehicle.vehicle_mass * self.environment.gravity * self.environment.road_tire_rolling_friction_coefficient
        direction = -1 * vehicle.vehicle_barycenter_velocity / np.linalg.norm(vehicle.vehicle_barycenter_velocity)
        rolling_friction = rolling_friction_module * direction  # 给质心运动方向加一个阻力, 暂且称为滚动摩擦力
        vehicle.sliding_friction = sliding_friction
        vehicle.rolling_friction = rolling_friction
        return sliding_friction, rolling_friction

    def calculate_moment(self,vehicle):
        sliding_friction_moment = np.zeros(4) # 最终都要投影到 z 轴上
        for k in range(4):
            OM = vehicle.OM[k]
            F = vehicle.sliding_friction[k]
            sliding_friction_moment[k] = np.dot(np.cross(OM,F),np.array([0,0,1]))
        vehicle.sliding_friction_moment = sliding_friction_moment
        return sliding_friction_moment

    def tick(self,vehicle):
        """
        迭代开始
        :return:
        """
        self.calculate_velocity(vehicle)
        self.calculate_force(vehicle)
        self.calculate_moment(vehicle)
        rotate_matrix_bary2terre = cal.euler2mat(vehicle.vehicle_euler_angles).T

        a = (np.sum(vehicle.sliding_friction,axis=0) + vehicle.rolling_friction) / vehicle.vehicle_mass
        a_angle = np.sum(vehicle.sliding_friction_moment) / vehicle.vehicle_moment_of_inertia
        direction_angle = vehicle.vehicle_euler_angles[0] + np.pi/2
        interval = self.environment.interval
        for k in range(4):
            omega = vehicle.tire_angular_velocity[k]
            sliding_friction = vehicle.sliding_friction[k]
            r = np.array([0, 0, -1]) * vehicle.tire_radius  # 对于目前模型在两个坐标系中一样
            normal_vector_tire = np.dot(rotate_matrix_bary2terre, np.array([1, 0, 0]))
            sliding_friction_moment_to_tire = np.dot(np.cross(r, sliding_friction), normal_vector_tire)
            vehicle.tire_angular_velocity[k] = omega  + sliding_friction_moment_to_tire / vehicle.tire_moment_of_inertia * interval
        euler_angle = vehicle.vehicle_euler_angles[0]
        # 以下的vehicle.vehicle_angular_velocity[2] 是临时措施, 不能通用, 需要修改 !!!!!
        # 以下的vehicle.vehicle_angular_velocity[2] 是临时措施, 不能通用, 需要修改 !!!!!
        # 以下的vehicle.vehicle_angular_velocity[2] 是临时措施, 不能通用, 需要修改 !!!!!
        vehicle.vehicle_euler_angles[0] = (euler_angle + 1/2 * a_angle * interval**2 + vehicle.vehicle_angular_velocity[2] * interval) % (2 * np.pi)
        vehicle.vehicle_barycenter_position = vehicle.vehicle_barycenter_position + 1/2 * a * interval**2 + vehicle.vehicle_barycenter_velocity * interval

        vehicle.vehicle_barycenter_velocity = vehicle.vehicle_barycenter_velocity + a * interval
        vehicle.vehicle_angular_velocity[2] = vehicle.vehicle_angular_velocity[2] + a_angle * interval

    def start(self):
        N = int(self.environment.duration // self.environment.interval + 1)
        vehicle0 = self.vehicles[0]
        X = np.zeros(N)
        Y = np.zeros(N)
        X2 = np.zeros(N)
        Y2 = np.zeros(N)
        step = 0
        for i in range(N):
            step += 1
            if step % 1000 == 0:
                print(step*100/N,"%")
            self.tick(vehicle0)
            print(vehicle0.vehicle_barycenter_velocity)
            X[i] = vehicle0.vehicle_barycenter_position[0]
            Y[i] = vehicle0.vehicle_barycenter_position[1]
            direction = cal.angle2vector(vehicle0.vehicle_euler_angles[0]+np.pi/2)
            l = vehicle0.vehicle_longth
            X2[i] = X[i] + direction[0] * l / 2
            Y2[i] = Y[i] + direction[1] * l / 2
        #     X.append(vehicle0.vehicle_barycenter_position[0])
        #     Y.append(vehicle0.vehicle_barycenter_position[1])
        # plt.figure()
        # plt.plot(X,Y)
        # plt.axis('equal')
        # plt.show()
        def update_points(num):
            '''
            更新数据点
            '''
            point_ani.set_data(X[num], Y[num])
            point_ani2.set_data(X2[num], Y2[num])
            return point_ani,point_ani2

        fig = plt.figure(tight_layout=True)
        plt.plot(X, Y, )
        plt.axis('equal')
        point_ani, = plt.plot(X[0], Y[0], "ro")
        point_ani2, = plt.plot(X2[0], Y2[0], "bo")
        plt.grid(ls="--")
        # 开始制作动画
        ani = animation.FuncAnimation(fig, update_points, np.arange(0, 10000, 10), interval=5, blit=True)

        # ani.save('sin_test2.gif', writer='imagemagick', fps=10)
        plt.show()


if __name__ == '__main__':
    print("start test")