import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import ConvexHull
import cmath
import random
import matplotlib.animation as animation

# Generate 10 random complex numbers as roots of a polynomial
random.seed(42)
roots_P = [complex(random.uniform(-2, 2), random.uniform(-2, 2)) for _ in range(10)]

# Create polynomial from roots
poly_coeffs = np.poly(roots_P)

# Derivative of the polynomial
poly_derivative_coeffs = np.polyder(poly_coeffs)

# Compute the roots of the derivative polynomial P'
roots_P_prime = np.roots(poly_derivative_coeffs)

# Get the real and imaginary parts of the roots
real_P = [root.real for root in roots_P]
imag_P = [root.imag for root in roots_P]

real_P_prime = [root.real for root in roots_P_prime]
imag_P_prime = [root.imag for root in roots_P_prime]

# Smooth displacement of each root in a random direction for the animation

# Choose a random direction for each root
directions = [complex(random.uniform(-0.02, 0.02), random.uniform(-0.02, 0.02)) for _ in roots_P]

# Function to update plot smoothly along chosen directions for multiple derivatives with convex hulls for each
def update_smooth_multiple_hulls(frame):
    plt.clf()
    
    # Smoothly move the roots of P along the chosen directions
    moved_roots_P = [root + frame * direction for root, direction in zip(roots_P, directions)]

    # Update the polynomial and its derivatives
    poly_coeffs = np.poly(moved_roots_P)
    
    # Derivatives
    poly_derivative_1 = np.polyder(poly_coeffs, 1)  # P'
    poly_derivative_2 = np.polyder(poly_coeffs, 2)  # P''
    poly_derivative_3 = np.polyder(poly_coeffs, 3)  # P'''
    
    # Compute the new roots of P', P'', P'''
    moved_roots_P_prime = np.roots(poly_derivative_1)
    moved_roots_P_double_prime = np.roots(poly_derivative_2)
    moved_roots_P_triple_prime = np.roots(poly_derivative_3)

    # Compute the convex hull of the updated roots of P and its derivatives
    points_P = np.array([[root.real, root.imag] for root in moved_roots_P])
    points_P_prime = np.array([[root.real, root.imag] for root in moved_roots_P_prime])
    points_P_double_prime = np.array([[root.real, root.imag] for root in moved_roots_P_double_prime])
    points_P_triple_prime = np.array([[root.real, root.imag] for root in moved_roots_P_triple_prime])

    # Compute convex hulls
    hull_P = ConvexHull(points_P)
    hull_P_prime = ConvexHull(points_P_prime)
    hull_P_double_prime = ConvexHull(points_P_double_prime)
    hull_P_triple_prime = ConvexHull(points_P_triple_prime)

    # Plot updated roots of P
    plt.plot([root.real for root in moved_roots_P], [root.imag for root in moved_roots_P], 'bo', markersize=10, label='Roots of P')
    
    # Plot updated roots of P'
    plt.plot([root.real for root in moved_roots_P_prime], [root.imag for root in moved_roots_P_prime], 'rx', markersize=10, label="Roots of P'")
    
    # Plot updated roots of P''
    plt.plot([root.real for root in moved_roots_P_double_prime], [root.imag for root in moved_roots_P_double_prime], 'gs', markersize=10, label="Roots of P''")
    
    # Plot updated roots of P'''
    plt.plot([root.real for root in moved_roots_P_triple_prime], [root.imag for root in moved_roots_P_triple_prime], 'md', markersize=10, label="Roots of P'''")

    # Plot convex hulls
    for simplex in hull_P.simplices:
        plt.plot(points_P[simplex, 0], points_P[simplex, 1], 'b-', lw=2)
    
    for simplex in hull_P_prime.simplices:
        plt.plot(points_P_prime[simplex, 0], points_P_prime[simplex, 1], 'r--', lw=2)

    for simplex in hull_P_double_prime.simplices:
        plt.plot(points_P_double_prime[simplex, 0], points_P_double_prime[simplex, 1], 'g--', lw=2)

    for simplex in hull_P_triple_prime.simplices:
        plt.plot(points_P_triple_prime[simplex, 0], points_P_triple_prime[simplex, 1], 'm--', lw=2)

    # Set labels and legend
    plt.xticks([])
    plt.yticks([])
    plt.legend(loc='upper right')

# Create the animation with smooth displacement and convex hulls for each derivative
fig = plt.figure(figsize=(8, 8))
ani_hulls = animation.FuncAnimation(fig, update_smooth_multiple_hulls, frames=100, interval=100)

# Save the animation as a video file
if True:
    ani_hulls.save('gauss-lucas-iterated.mp4', writer='ffmpeg', fps=10)

# Show the animation
plt.show()