import matplotlib.pyplot as plt
import numpy as np

def draw_lifecycle_diagram():
    # Labels and colors for the lifecycle phases
    labels = [
        "1. Gathering Data", 
        "2. Data Preparation", 
        "3. Data Wrangling", 
        "4. Analyze Data", 
        "5. Train Model", 
        "6. Test Model", 
        "7. Deployment"
    ]
    colors = [
        "#3CB371", "#8FBC8F", "#00CED1", "#1E90FF", 
        "#6A5ACD", "#FF8C00", "#DC143C"
    ]

    # Create a figure and axis with equal aspect ratio
    fig, ax = plt.subplots(figsize=(9, 9), subplot_kw={"aspect": "equal"})
    size = 0.3  # Width of the pie sections

    # Create pie sections for the lifecycle phases
    wedges, _ = ax.pie(
        [1] * len(labels), 
        colors=colors, 
        radius=1, 
        startangle=90, 
        wedgeprops=dict(width=size, edgecolor='w')
    )

    # Add text labels around the circle
    for i, wedge in enumerate(wedges):
        # Calculate the angle for the label placement
        angle = (wedge.theta2 - wedge.theta1) / 2.0 + wedge.theta1
        x = np.cos(np.deg2rad(angle))
        y = np.sin(np.deg2rad(angle))

        # Add label text
        ax.text(
            1.2 * x, 1.2 * y, labels[i],
            ha="center", va="center", fontsize=10, weight="bold",
            bbox=dict(boxstyle="round,pad=0.3", facecolor=colors[i], edgecolor="w")
        )

    # Add center text with a descriptive title
    ax.text(
        0, 0, "Machine Learning\nLifecycle", 
        ha="center", va="center", fontsize=16, weight="bold", color="black",
        bbox=dict(boxstyle="round,pad=0.5", facecolor="white", edgecolor="black")
    )

    # Clean up the diagram style
    ax.set(aspect="equal", xticks=[], yticks=[], title="Machine Learning Lifecycle")
    return fig

# Save the diagram or display it in Streamlit
if __name__ == "__main__":
    # Display the diagram
    fig = draw_lifecycle_diagram()
    plt.show()