from traceback import print_exc

from rlbot.agents.base_script import BaseScript
from rlbot.utils.structures.game_data_struct import GameTickPacket

import rl_ball_sym as rlbs


class rl_ball_sym(BaseScript):
    def __init__(self):
        super().__init__("rl_ball_sym")

    def main(self):
        rlbs.load_soccar()

        while 1:
            try:
                self.packet: GameTickPacket = self.wait_game_tick_packet()
                current_location = self.packet.game_ball.physics.location
                current_velocity = self.packet.game_ball.physics.velocity
                current_angular_velocity = self.packet.game_ball.physics.angular_velocity

                rlbs.set_ball({
                    "time": self.packet.game_info.seconds_elapsed,
                    "location": [current_location.x, current_location.y, current_location.z],
                    "velocity": [current_velocity.x, current_velocity.y, current_velocity.z],
                    "angular_velocity": [current_angular_velocity.x, current_angular_velocity.y, current_angular_velocity.z],
                })

                path_prediction = rlbs.get_ball_prediction_struct()

                self.renderer.begin_rendering()
                self.renderer.draw_polyline_3d(tuple((path_prediction["slices"][i]["location"][0], path_prediction["slices"][i]["location"][1], path_prediction["slices"][i]["location"][2]) for i in range(0, path_prediction["num_slices"], 4)), self.renderer.red())
                self.renderer.end_rendering()
            except Exception:
                print_exc()


if __name__ == "__main__":
    rl_ball_sym = rl_ball_sym()
    rl_ball_sym.main()

[
    {
        "time": 0.008333,
        "location": [
            -2283.9,
            1683.8,
            323.4,
        ],
        "velocity": [
            1273.4,
            -39.7,
            757.6,
        ]
    },
    {
        "time": 0.025,
        "location": [
            -2262.6,
            1683.1,
            335.9,
        ],
        "velocity": [
            1272.7,
            -39.7,
            746.4,
        ]
    }
    ...
]

[
    {
        "time": 0.008333,
        "location": [
            -2283.9,
            1683.8,
            323.4,
        ],
        "velocity": [
            1273.4,
            -39.7,
            757.6,
        ]
        "angular_velocity": [
            2.3,
            -0.8,
            3.8,
        }
    },
    {
        "time": 0.016666,
        "location": [
            -2273.3,
            1683.4,
            329.7,
        ],
        "velocity": [
            1273.1,
            -39.7,
            752.0,
        ],
        "angular_velocity": [
            2.3,
            -0.8,
            3.8
        ]
    }
    ...
]

{
    "time": 0.008333,
    "location": [
        -2283.9,
        1683.8,
        323.4,
    ],
    "velocity": [
        1273.4,
        -39.7,
        757.6,
    ]
    "angular_velocity": [
        2.3,
        -0.8,
        3.8,
    }
}