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):
        self.set_camera_orientation(phi=5 * PI / 12, theta=-6 * PI / 12)
        self.begin_ambient_camera_rotation(rate=0.2)
        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)

        state0 = [10, 10, 10]
        points = ode_solution_points(
            lorenz_system, state0, time=32.4, dt=0.002
        )

        curve = VMobject().set_points_as_corners(axes.c2p(points[500:]))
        curve.set_stroke(GREEN)
        self.play(Create(curve, run_time=31.4, rate_func=linear))