"""
KNN Playground — interactive parameter exploration.
"""

import numpy as np

META = {
    "title": "K-Nearest Neighbors",
    "icon": "🎯",
    "description": (
        "Visualize how KNN classifies 2D data. Adjust K, the number of neighbors, "
        "and see how the decision boundary changes — from jagged (small K) to smooth (large K)."
    ),
}

PARAMS = [
    {
        "name": "k",
        "label": "K (Neighbors)",
        "type": "slider",
        "min": 1,
        "max": 25,
        "step": 2,
        "default": 5,
    },
    {
        "name": "n_points",
        "label": "Points per Class",
        "type": "slider",
        "min": 20,
        "max": 150,
        "step": 10,
        "default": 60,
    },
    {
        "name": "noise",
        "label": "Noise / Spread",
        "type": "slider",
        "min": 0.3,
        "max": 2.0,
        "step": 0.1,
        "default": 0.8,
    },
]


def compute(params: dict) -> dict:
    """Compute KNN decision boundary heatmap."""
    k = int(params.get("k", 5))
    n = int(params.get("n_points", 60))
    noise = float(params.get("noise", 0.8))

    np.random.seed(42)
    # Two classes
    c0 = np.random.randn(n, 2) * noise + [1, 1]
    c1 = np.random.randn(n, 2) * noise + [3, 3]
    X = np.vstack([c0, c1])
    y = np.array([0] * n + [1] * n)

    # Build grid for decision boundary
    margin = 1.5
    x_min, x_max = X[:, 0].min() - margin, X[:, 0].max() + margin
    y_min, y_max = X[:, 1].min() - margin, X[:, 1].max() + margin
    res = 80
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, res),
                          np.linspace(y_min, y_max, res))
    grid = np.c_[xx.ravel(), yy.ravel()]

    # Manual KNN on grid
    from scipy.spatial.distance import cdist
    dists = cdist(grid, X)
    knn_idx = np.argsort(dists, axis=1)[:, :k]
    knn_labels = y[knn_idx]
    preds = (knn_labels.mean(axis=1) >= 0.5).astype(float)
    Z = preds.reshape(xx.shape)

    data = [
        {
            "z": Z.tolist(),
            "x": np.linspace(x_min, x_max, res).tolist(),
            "y": np.linspace(y_min, y_max, res).tolist(),
            "type": "heatmap",
            "colorscale": [[0, "rgba(108,99,255,0.25)"], [1, "rgba(255,101,132,0.25)"]],
            "showscale": False,
            "hoverinfo": "skip",
        },
        {
            "x": c0[:, 0].tolist(),
            "y": c0[:, 1].tolist(),
            "mode": "markers",
            "type": "scatter",
            "name": "Class 0",
            "marker": {"size": 8, "color": "#6C63FF", "line": {"width": 1, "color": "#fff"}},
        },
        {
            "x": c1[:, 0].tolist(),
            "y": c1[:, 1].tolist(),
            "mode": "markers",
            "type": "scatter",
            "name": "Class 1",
            "marker": {"size": 8, "color": "#FF6584", "line": {"width": 1, "color": "#fff"}},
        },
    ]

    layout = {
        "template": "plotly_dark",
        "paper_bgcolor": "rgba(0,0,0,0)",
        "plot_bgcolor": "rgba(0,0,0,0)",
        "title": {"text": f"KNN Decision Boundary (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}
