import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from scipy.spatial import Voronoi


def plot_kmeans_results(df, cluster_column, title, x_column='umap_x', y_column='umap_y', show_voronoi=False):

    x_min, x_max = df[x_column].min(), df[x_column].max()
    y_min, y_max = df[y_column].min(), df[y_column].max()
    
    margin_x = (x_max - x_min) * 0.05
    margin_y = (y_max - y_min) * 0.05

    fig = px.scatter(
        df, x=x_column, y=y_column, 
        color=df[cluster_column].astype(str),
        labels={cluster_column: "Klaster"},
        hover_name='title', title=title, template='plotly_white',
        width=1000, height=700
    )
    fig.update_traces(marker=dict(size=4, opacity=0.6))


    if show_voronoi:
        centroids_2d = df.groupby(cluster_column)[[x_column, y_column]].mean().values
        vor = Voronoi(centroids_2d)
        center = centroids_2d.mean(axis=0)

        far_dist = (x_max - x_min) * 5 

        for pointidx, simplex in zip(vor.ridge_points, vor.ridge_vertices):
            simplex = np.asarray(simplex)
            if np.all(simplex >= 0):
                v1, v2 = vor.vertices[simplex]
                fig.add_shape(type="line", x0=v1[0], y0=v1[1], x1=v2[0], y1=v2[1],
                              line=dict(color="rgba(0,0,0,0.8)", width=1, dash="dot"))
            else:
                i = simplex[simplex >= 0][0]
                t = centroids_2d[pointidx[1]] - centroids_2d[pointidx[0]]
                t /= np.linalg.norm(t)
                n = np.array([-t[1], t[0]])
                midpoint = centroids_2d[pointidx].mean(axis=0)
                direction = np.sign(np.dot(midpoint - center, n)) * n
                far_point = vor.vertices[i] + direction * far_dist

                fig.add_shape(type="line", x0=vor.vertices[i][0], y0=vor.vertices[i][1],
                              x1=far_point[0], y1=far_point[1],
                              line=dict(color="rgba(0,0,0,0.8)", width=1, dash="dot"))
                
                
                
                fig.add_trace(go.Scatter(
                    x=centroids_2d[:, 0], 
                    y=centroids_2d[:, 1],
                    mode='markers', 
                    marker=dict(
                        symbol='x', 
                        size=15, 
                        color='white',      # Białe tło (obrys)
                        line=dict(width=4, color='white') 
                    ),
                    showlegend=False,
                    hoverinfo='skip'
                ))
                
                fig.add_trace(go.Scatter(
                    x=centroids_2d[:, 0], 
                    y=centroids_2d[:, 1],
                    mode='markers', 
                    marker=dict(
                        symbol='x', 
                        size=14, 
                        color='black',      
                        line=dict(width=2) 
                    ),
                    name='Centroidy',
                    showlegend=False,
                    hoverinfo='skip'
                ))
                
                fig.update_layout(
                    xaxis=dict(range=[x_min - margin_x, x_max + margin_x]),
                    yaxis=dict(range=[y_min - margin_y, y_max + margin_y]),
                    # To wymusza, by scatter był pod spodem warstw rysowanych później
                    scattermode='group' 
                )

    fig.update_layout(
        xaxis=dict(range=[x_min - margin_x, x_max + margin_x]),
        yaxis=dict(range=[y_min - margin_y, y_max + margin_y])
    )

    fig.show()