import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np

# import solarsystem as solar
import solarsystem_exercice as solar

fig = plt.figure(figsize=[6, 6])
ax = plt.axes([0.0, 0.0, 1.0, 1.0], xlim=(-1.8, 1.8), ylim=(-1.8, 1.8))
ax.set_aspect("equal")
ax.axis("off")


def perform_simulation(
    folder: str,
    num_frames: int,
    num_planets: int,
    colors: np.ndarray,
    radius: np.ndarray,
    distance_planet_sun: np.ndarray,
    test_system: solar.SolarSystem,
):
    """
    This is the visualization function
    """
    plt.style.use("dark_background")

    _, fig_height = plt.gcf().get_size_inches()

    track_old_x_positions = np.zeros(shape=(num_frames, num_planets))
    track_old_y_positions = np.zeros(shape=(num_frames, num_planets))

    def animate(i):
        print(f"\r{i}/{num_frames}", end="")
        global ax
        ax.clear()
        ax = plt.axes([0.0, 0.0, 1.0, 1.0], xlim=(-1.8, 1.8), ylim=(-1.8, 1.8))
        ax.set_aspect("equal")
        ax.axis("off")
        for _ in range(100):
            _ = test_system()
        new_positions = fig_height * test_system() / np.max(distance_planet_sun) / 4
        plots = []
        for cpt_planet, new_position in enumerate(new_positions):
            ax.scatter(
                new_position[0],
                new_position[1],
                color=colors[cpt_planet],
                s=radius[cpt_planet],
                edgecolors=None,
                zorder=10,
            )
            plt.plot(
                track_old_x_positions[:i, cpt_planet],
                track_old_y_positions[:i, cpt_planet],
                color=colors[cpt_planet],
                linestyle="-",
            )
            track_old_x_positions[i, cpt_planet] = new_position[0]
            track_old_y_positions[i, cpt_planet] = new_position[1]
        plt.text(-1.7, 1.65, s=f"{i}", color=[1, 1, 1], fontsize=22)
        plt.savefig(f"{folder}/simulation_frame{i}.png")
        return plots

    _ = animation.FuncAnimation(
        fig,
        animate,
        repeat=False,
        frames=num_frames,
        blit=True,
        interval=20,
        cache_frame_data=False,
    )
    plt.show()
    plt.close()