import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from sklearn.mixture import GaussianMixture

def plot_em_results(df, cluster_column, title, x_column='umap_x', y_column='umap_y', 
                    show_ellipses=True, covariance_type='full'):
    """
    Jeśli parametr show_ellipses=True funkcja aproksymuje elipsę odpowiadającą macierzy kowariancji.
    - covariance_type: 'full' - na podstawie rzeczywistych etykiet
    - covariance_type: ''tied', 'diag', 'spherical' - buduje model GMM na danych dwuwymiarowych
    """
    
    # 2. Podstawowy wykres punktowy (z nowymi etykietami klastrów z przestrzeni 2D)
    fig = px.scatter(
        df, x=x_column, y=y_column, 
        color=df[cluster_column].astype(str),
        labels={cluster_column: "Klaster"},
        hover_name='title', title=f"{title} ({covariance_type})", template='plotly_white',
        width=1000, height=700
    )
    fig.update_traces(marker=dict(size=4, opacity=0.5))

    if show_ellipses:
        if covariance_type == 'full':
            clusters = sorted(df[cluster_column].unique(), key=int)
            t = np.linspace(0, 2*np.pi, 100) # Parametr do rysowania obwodu
            
            for cluster_id in clusters:
                points = df[df[cluster_column] == cluster_id][[x_column, y_column]].values
                if len(points) < 5: continue
    
                mean = points.mean(axis=0)
                cov = np.cov(points, rowvar=False)
                vals, vecs = np.linalg.eigh(cov)
                
                # Skalowanie dla 2 odchyleń standardowych (ok. 95% punktów)
                # Używamy pierwiastka z wartości własnych jako promieni
                width, height = 2 * np.sqrt(vals) 
                
                # Obliczanie punktów elipsy z uwzględnieniem obrotu
                # x = mean_x + a*cos(t)*cos(theta) - b*sin(t)*sin(theta)
                # y = mean_y + a*cos(t)*sin(theta) + b*sin(t)*cos(theta)
                theta = np.arctan2(vecs[1, 0], vecs[0, 0])
                
                x_ellipse = mean[0] + width * np.cos(t) * np.cos(theta) - height * np.sin(t) * np.sin(theta)
                y_ellipse = mean[1] + width * np.cos(t) * np.sin(theta) + height * np.sin(t) * np.cos(theta)
    
                # Dodanie elipsy jako ścieżki (go.Scatter z wypełnieniem lub linią)
                fig.add_trace(go.Scatter(
                    x=x_ellipse, y=y_ellipse,
                    mode='lines',
                    line=dict(color='rgba(0,0,0,0.4)', width=1.5, dash='dot'),
                    showlegend=False,
                    hoverinfo='skip'
                ))
    
                # Centroid
                fig.add_trace(go.Scatter(
                    x=[mean[0]], y=[mean[1]],
                    mode='markers',
                    marker=dict(symbol='x', size=10, color='black'),
                    showlegend=False, hoverinfo='skip'
                ))
        else:
            # Pobieramy dane 2D do trenowania EM
            X_2d = df[[x_column, y_column]].values
            
            # Dynamicznie sprawdzamy liczbę klastrów na podstawie unikalnych wartości w kolumnie
            # (Zakładamy, że kolumna zawiera poprawne identyfikatory, np. liczby)
            n_clusters = len(df[cluster_column].unique())
            
            # 1. Uruchamiamy Gaussian Mixture bezpośrednio na danych 2D
            gmm = GaussianMixture(n_components=n_clusters, covariance_type=covariance_type, random_state=42)
            gmm.fit_predict(X_2d)
            t = np.linspace(0, 2*np.pi, 100) # Parametr obwodu
            
            for cluster_id in range(n_clusters):
                # Pobieramy średnią dla danego klastra z modelu GMM
                mean = gmm.means_[cluster_id]
                
                # Ekstrakcja macierzy kowariancji w zależności od covariance_type
                if covariance_type == 'full':
                    cov = gmm.covariances_[cluster_id]
                elif covariance_type == 'tied':
                    cov = gmm.covariances_  # Wspólna macierz dla wszystkich
                elif covariance_type == 'diag':
                    cov = np.diag(gmm.covariances_[cluster_id]) # Tylko wariancje na przekątnej
                elif covariance_type == 'spherical':
                    cov = np.eye(2) * gmm.covariances_[cluster_id] # Taka sama wariancja w każdym kierunku
                
                # Geometria elipsy na podstawie macierzy kowariancji
                vals, vecs = np.linalg.eigh(cov)
                
                # Skalowanie dla 2 odchyleń standardowych (95% prawdopodobieństwa rozkładu)
                width, height = 2 * np.sqrt(vals) 
                theta = np.arctan2(vecs[1, 0], vecs[0, 0])
                
                x_ellipse = mean[0] + width * np.cos(t) * np.cos(theta) - height * np.sin(t) * np.sin(theta)
                y_ellipse = mean[1] + width * np.cos(t) * np.sin(theta) + height * np.sin(t) * np.cos(theta)
    
                # Rysowanie elipsy
                fig.add_trace(go.Scatter(
                    x=x_ellipse, y=y_ellipse,
                    mode='lines',
                    line=dict(color='rgba(0,0,0,0.5)', width=1.5, dash='dot'),
                    showlegend=False, hoverinfo='skip'
                ))
    
                # Rysowanie centroidu wyznaczonego przez EM
                fig.add_trace(go.Scatter(
                    x=[mean[0]], y=[mean[1]],
                    mode='markers',
                    marker=dict(symbol='x', size=10, color='black'),
                    showlegend=False, hoverinfo='skip'
                ))

    # Skalowanie osi (dodaje margines, żeby elipsy nie wychodziły poza wykres)
    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 = 0.15
    fig.update_layout(
        xaxis=dict(range=[x_min - (x_max-x_min)*margin, x_max + (x_max-x_min)*margin]),
        yaxis=dict(range=[y_min - (y_max-y_min)*margin, y_max + (y_max-y_min)*margin])
    )

    fig.show()