"""
K-Means Playground — interactive parameter exploration.
"""

import numpy as np

META = {
    "title": "K-Means Clustering",
    "icon": "🔵",
    "description": (
        "Watch K-Means partition data into clusters. Change K, the number "
        "of data points, and cluster spread to see how the algorithm adapts."
    ),
}

PARAMS = [
    {
        "name": "k",
        "label": "K (Clusters)",
        "type": "slider",
        "min": 2,
        "max": 8,
        "step": 1,
        "default": 3,
    },
    {
        "name": "n_points",
        "label": "Points per Cluster",
        "type": "slider",
        "min": 20,
        "max": 150,
        "step": 10,
        "default": 60,
    },
    {
        "name": "spread",
        "label": "Cluster Spread",
        "type": "slider",
        "min": 0.3,
        "max": 2.5,
        "step": 0.1,
        "default": 0.9,
    },
]


def compute(params: dict) -> dict:
    """Run K-Means and visualize clusters + centroids."""
    k = int(params.get("k", 3))
    n = int(params.get("n_points", 60))
    spread = float(params.get("spread", 0.9))

    np.random.seed(42)

    # Generate true cluster centers spread around a circle
    true_k = min(k, 8)
    angles = np.linspace(0, 2 * np.pi, true_k, endpoint=False)
    centers = np.column_stack([4 * np.cos(angles), 4 * np.sin(angles)])

    # Generate data
    X = np.vstack([np.random.randn(n, 2) * spread + c for c in centers])

    # Simple K-Means implementation
    # Initialize centroids via K-Means++
    rng = np.random.RandomState(42)
    centroids = [X[rng.randint(len(X))]]
    for _ in range(k - 1):
        dists = np.min([np.sum((X - c) ** 2, axis=1) for c in centroids], axis=0)
        probs = dists / dists.sum()
        centroids.append(X[rng.choice(len(X), p=probs)])
    centroids = np.array(centroids)

    # Iterate
    for _ in range(30):
        dists = np.array([np.sum((X - c) ** 2, axis=1) for c in centroids])
        labels = np.argmin(dists, axis=0)
        new_centroids = np.array([X[labels == i].mean(axis=0) if np.any(labels == i) else centroids[i] for i in range(k)])
        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids

    cluster_colors = ["#6C63FF", "#FF6584", "#34d399", "#fbbf24", "#a78bfa", "#f472b6", "#38bdf8", "#fb923c"]

    data = []
    for i in range(k):
        mask = labels == i
        data.append({
            "x": X[mask, 0].tolist(),
            "y": X[mask, 1].tolist(),
            "mode": "markers",
            "type": "scatter",
            "name": f"Cluster {i}",
            "marker": {"size": 7, "color": cluster_colors[i % len(cluster_colors)], "opacity": 0.75,
                        "line": {"width": 1, "color": "#fff"}},
        })

    data.append({
        "x": centroids[:, 0].tolist(),
        "y": centroids[:, 1].tolist(),
        "mode": "markers",
        "type": "scatter",
        "name": "Centroids",
        "marker": {"size": 18, "color": "#fff", "symbol": "x",
                    "line": {"width": 3, "color": "#fff"}},
    })

    layout = {
        "template": "plotly_dark",
        "paper_bgcolor": "rgba(0,0,0,0)",
        "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": f"K-Means Clustering (K={k})", "font": {"size": 20}},
        "xaxis": {"title": "Feature 1", "gridcolor": "rgba(255,255,255,0.08)"},
        "yaxis": {"title": "Feature 2", "gridcolor": "rgba(255,255,255,0.08)", "scaleanchor": "x"},
        "legend": {"x": 0.02, "y": 0.98},
        "margin": {"l": 60, "r": 30, "t": 60, "b": 50},
    }

    return {"data": data, "layout": layout}
