from math import pi
import numpy as np
from manim import *
from scipy.integrate import solve_ivp


def lorenz_system(t, state, sigma=10, rho=28, beta=8 / 3):
    x, y, z = state
    dxdt = sigma * (y - x)
    dydt = x * (rho - z) - y
    dzdt = x * y - beta * z
    return [dxdt, dydt, dzdt]


def ode_solution_points(function, state0, time, dt=0.01):
    solution = solve_ivp(
        function, t_span=(0, time), y0=state0, t_eval=np.arange(0, time, dt)
    )
    return solution.y.T


class LorenzAttractor(ThreeDScene):
    def construct(self):
        revolutions = 1.0
        camera_rotation_rate = 0.2
        warm_up_time = 0.2
        dt = 0.002
        epsilon = 0.001
        states = [[5.75, 10.0, 10.1 + n * epsilon] for n in range(2)]
        colors = color_gradient([BLUE, GREEN], len(states))

        self.set_camera_orientation(phi=5 * PI / 12, theta=-6 * PI / 12)
        self.begin_ambient_camera_rotation(rate=camera_rotation_rate)
        axes = ThreeDAxes(
            x_range=(-50, 50, 5),
            y_range=(-50, 50, 5),
            z_range=(0, 50, 5),
            x_length=16,
            y_length=16,
            z_length=8,
        )
        axes.center()
        self.add(axes)

        run_time = 2.0 * pi * revolutions / camera_rotation_rate
        warm_up_len = int(warm_up_time / dt)

        curves = VGroup()
        for state, color in zip(states, colors):
            points = ode_solution_points(
                lorenz_system, state, time=warm_up_time + run_time, dt=dt
            )
            curve = VMobject().set_points_as_corners(axes.c2p(points[warm_up_len:]))
            curve.set_stroke(color)
            curves.add(curve)

        dots = Group(*(Dot3D(color=color, radius=0.1) for color in colors))

        globals().update(locals())
        def update_dots(dots):
            for dot, curve in zip(dots, curves):
                dot.move_to(curve.get_end())

        dots.add_updater(update_dots)

        self.add(dots)
        self.play(*(
            Create(curve, run_time=run_time, rate_func=linear)
            for curve in curves
        ))