#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep  2 21:36:32 2024

@author: amriksen
"""

import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from my_Root_Finding_ComplexPlane_Module import *

# Define the symbolic variable and function
z_sym = sp.symbols('z')
func_sym = z_sym**2 * (z_sym**3 - 1)
#func_sym = z_sym * (z_sym + 0.01) * (z_sym**3 - 1)
mul = 2.0 # for z = 0

# Compute the derivative symbolically
dfunc_sym = sp.diff(func_sym, z_sym)
d2func_sym = sp.diff(dfunc_sym, z_sym)

# Convert the symbolic expressions to numerical functions
func = sp.lambdify(z_sym, func_sym, "numpy")
dfunc = sp.lambdify(z_sym, dfunc_sym, "numpy")
d2func = sp.lambdify(z_sym, d2func_sym, "numpy")

# Define the grid of points in the complex plane
resolution = 500
x = np.linspace(-2, 2, resolution)
y = np.linspace(-2, 2, resolution)
X, Y = np.meshgrid(x, y)
Z = X + 1j * Y

# Roots of the function (written in exp(i*theta) form of a complex number)
roots = [1, np.exp(2j * np.pi / 3), np.exp(-2j * np.pi / 3), 0] 


# Initialize the basin of attraction
basin = np.zeros(Z.shape, dtype=int)

# Iterate over each point in the complex plane
for i in range(Z.shape[0]):
    for j in range(Z.shape[1]):
        z = Z[i, j]
        final_z = my_Newton_Method_Multiple_Roots(z, func, dfunc, d2func)
        #final_z = my_Relaxed_Newton_Method_ComplexPlane(z, func, dfunc, mul)
        root_index = find_root_index(final_z, roots)
        basin[i, j] = root_index

# Plotting the basin of attraction
plt.figure(figsize=(8, 8))
plt.imshow(basin, extent=(-2, 2, -2, 2), cmap='jet')
plt.colorbar()
plt.title('Basin of Attraction for f(z) = z^2(z^3 - 1)')
plt.xlabel('Re(z)')
plt.ylabel('Im(z)')
plt.show()