templates/clustering.html

{% extends "layout.html" %}

{% block content %}
<style>
    /* Custom styles for the K-Means visualization page */
    .kmeans-page-wrapper {
        font-family: "Inter", sans-serif;
        background-color: #f0f4f8; /* Light background */
        display: flex;
        justify-content: center;
        align-items: flex-start; /* Align items to start to prevent content being pushed up */
        min-height: 100vh; /* Ensure it takes full viewport height */
        padding: 20px;
        box-sizing: border-box;
        width: 100%; /* Ensure it takes full width of its parent (.main) */
        overflow-y: auto; /* Allow scrolling if content overflows vertically */
    }

    .container {
        max-width: 1000px; /* Increased max-width for better layout */
        width: 100%; /* Ensure it respects max-width and is responsive */
        background-color: #ffffff;
        padding: 2rem;
        border-radius: 12px;
        box-shadow: 0 10px 25px rgba(0, 0, 0, 0.1);
    }
    h1, h2, h3, h4 {
        color: #2c3e50;
    }
    .info-icon {
        cursor: help;
        margin-left: 5px;
        color: #6B7280; /* gray-500 */
        position: relative;
        display: inline-block;
    }
    .tooltip {
        visibility: hidden;
        width: 250px;
        background-color: #333;
        color: #fff;
        text-align: center;
        border-radius: 6px;
        padding: 8px 10px;
        position: absolute;
        z-index: 10;
        bottom: 125%;
        left: 50%;
        margin-left: -125px;
        opacity: 0;
        transition: opacity 0.3s;
        font-size: 0.85rem;
        line-height: 1.4;
    }
    .info-icon:hover .tooltip {
        visibility: visible;
        opacity: 1;
    }
    .tooltip::after {
        content: "";
        position: absolute;
        top: 100%;
        left: 50%;
        margin-left: -5px;
        border-width: 5px;
        border-style: solid;
        border-color: #333 transparent transparent transparent;
    }
    .highlight-blue { color: #2563EB; font-weight: 600; }
    .highlight-red { color: #DC2626; font-weight: 600; }
    .highlight-green { color: #16A34A; font-weight: 600; }
    .highlight-purple { color: #9333EA; font-weight: 600; }
    .highlight-bold { font-weight: 600; }
    .flow-box {
        background-color: #F3F4F6;
        border-radius: 0.5rem;
        padding: 1.5rem;
        text-align: center;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
        min-height: 120px;
        display: flex;
        flex-direction: column;
        justify-content: center;
        align-items: center;
    }
    .flow-arrow {
        font-size: 2.5rem;
        color: #9CA3AF;
        margin: 0 1rem;
        display: flex;
        align-items: center;
        justify-content: center;
    }
    .math {
        font-size: 1.1em;
        display: block;
        margin-top: 1em;
        margin-bottom: 1em;
        overflow-x: auto;
        padding: 0.5em;
        background-color: #f8fafc;
        border-left: 4px solid #6366f1;
        padding-left: 1rem;
    }
    #plotly-graph {
        border-radius: 0.5rem;
        box-shadow: inset 0 2px 4px 0 rgba(0, 0, 0, 0.06);
        padding: 1rem;
        background-color: #ffffff;
        min-height: 500px;
        height: 600px;
        width: 100%;
    }
</style>

<script src="https://cdn.tailwindcss.com"></script>
<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>


<div class="kmeans-page-wrapper">
    <div class="container">
        <h1 class="text-3xl font-bold mb-4 text-center">📊 Interactive Dynamic K-Means Clustering Visualization</h1>
        <p class="mb-6 text-center text-gray-600">
            Explore K-Means, an unsupervised learning algorithm that partitions data into K distinct clusters. For example, an online store uses K-Means to group customers based on purchase frequency and spending, creating segments like <span class="highlight-bold">Budget Shoppers</span>, <span class="highlight-bold">Frequent Buyers</span>, and <span class="highlight-bold">Big Spenders</span> for personalized marketing.
        </p>

        <div class="controls grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4 mb-6">
            <div class="flex flex-col">
                <label for="dimensions" class="text-gray-700 text-sm font-semibold mb-1">Dimensions:</label>
                <select id="dimensions" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
                    <option value="2D">2D</option>
                    <option value="3D">3D</option>
                </select>
            </div>

            <div class="flex flex-col">
                <label for="num-clusters" class="text-gray-700 text-sm font-semibold mb-1">Number of Clusters (K):</label>
                <input type="number" id="num-clusters" value="3" min="2" max="10" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>

            <div class="flex flex-col">
                <label for="max-iterations" class="text-gray-700 text-sm font-semibold mb-1">Max Iterations:</label>
                <input type="number" id="max-iterations" value="100" min="10" max="500" step="10" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>

            <div class="flex flex-col">
                <label for="data-points-total" class="text-gray-700 text-sm font-semibold mb-1">Total Data Points:</label>
                <input type="number" id="data-points-total" value="100" min="20" max="500" step="10" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>

            <div class="flex flex-col">
                <label for="new-point-x" class="text-gray-700 text-sm font-semibold mb-1">New Point X:</label>
                <input type="number" id="new-point-x" value="0" step="0.1" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>
            <div class="flex flex-col" id="new-point-y-wrapper">
                <label for="new-point-y" class="text-gray-700 text-sm font-semibold mb-1">New Point Y:</label>
                <input type="number" id="new-point-y" value="0" step="0.1" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>
            <div class="flex flex-col hidden" id="new-point-z-wrapper">
                <label for="new-point-z" class="text-gray-700 text-sm font-semibold mb-1">New Point Z:</label>
                <input type="number" id="new-point-z" value="0" step="0.1" class="p-2 border border-gray-300 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
            </div>

            <button id="add-point-btn" class="bg-blue-500 hover:bg-blue-600 text-white font-bold py-2 px-4 rounded-md focus:outline-none focus:ring-2 focus:ring-blue-500">
                Add New Point & Cluster
            </button>
            <button id="reset-data-btn" class="bg-red-500 hover:bg-red-600 text-white font-bold py-2 px-4 rounded-md focus:outline-none focus:ring-2 focus:ring-red-500">
                Reset Data & Recluster
            </button>
        </div>

        <div id="plotly-graph" class="w-full h-96 md:h-[500px] lg:h-[600px] border border-gray-300 rounded-lg"></div>
        <p id="prediction-result" class="mt-4 font-bold text-lg text-center text-gray-800 mb-8"></p>

        <a href="/kmeans-Dbscan-image">
  <button class="inline-block bg-gray-200 hover:bg-gray-300 centre  text-gray-800 px-4 py-2 rounded shadow">🖼️ KMeans + DBSCAN Image Clustering</button>
</a>

        <!-- New Section: How K-Means Clustering Works? -->
        <div class="mt-10 p-6 bg-green-50 rounded-xl border border-green-200">
            <h2 class="text-2xl font-bold mb-6 text-center text-green-700">How K-Means Clustering Works?</h2>
            <p class="mb-4 text-gray-700">
                We are given a data set of items with certain features and values for these features like a vector. The task is to categorize those items into groups. To achieve this we will use the K-means algorithm. <span class="highlight-bold">'K'</span> in the name of the algorithm represents the number of groups/clusters we want to classify our items into.
            </p>
            <p class="mb-4 text-gray-700">
                The algorithm works by first randomly picking some central points called <span class="highlight-bold">centroids</span> and each data point is then assigned to the closest centroid forming a cluster. After all the points are assigned to a cluster, the centroids are updated by finding the average position of the points in each cluster. This process repeats until the centroids stop changing, forming clusters. The goal of clustering is to divide the data points into clusters so that similar data points belong to the same group.
            </p>
        </div>
     
        <!-- New Section: Flow of Data Treated by K-Means Algorithm -->
        <div class="mt-8 p-6 bg-gray-50 rounded-lg border border-gray-200">
            <h2 class="text-2xl font-bold mb-4 text-center text-blue-700">Flow of Data Treated by K-Means Algorithm</h2>
            <div class="flex flex-wrap justify-center items-center gap-4">
                <div class="flow-box bg-blue-100">
                    <span class="text-5xl mb-2">✨</span> <p class="text-lg font-semibold text-blue-800">1. Initialize Centroids</p>
                    <p class="text-sm text-blue-600">Randomly pick K points</p>
                </div>
                <div class="flow-arrow">&rarr;</div>
                <div class="flow-box bg-blue-100">
                    <span class="text-5xl mb-2">📍</span> <p class="text-lg font-semibold text-blue-800">2. Assign Points to Clusters</p>
                    <p class="text-sm text-blue-600">To closest centroid</p>
                </div>
                <div class="flow-arrow">&rarr;</div>
                <div class="flow-box bg-blue-100">
                    <span class="text-5xl mb-2">🔄</span> <p class="text-lg font-semibold text-blue-800">3. Update Centroids</p>
                    <p class="text-sm text-blue-600">Calculate new means</p>
                </div>
                <div class="flow-arrow block md:hidden">&darr;</div>
                <div class="flow-arrow hidden md:block">&rarr;</div>
                <div class="flow-box bg-blue-100">
                    <span class="text-5xl mb-2">✅</span> <p class="text-lg font-semibold text-blue-800">4. Convergence Check</p>
                    <p class="text-sm text-blue-600">Centroids stabilized?</p>
                </div>
            </div>
            <p class="mt-6 text-center text-gray-600 text-sm">
                The K-Means algorithm iteratively refines clusters. It starts by randomly selecting initial centroids. Then, each data point is assigned to the closest centroid, forming preliminary clusters. Next, the centroids are updated to the mean position of all points within their assigned clusters. This assignment and update process repeats until the centroids no longer change significantly, indicating that the clusters have converged.
            </p>
        </div>

        <!-- Existing K-Means Specific Explanations Below -->
        <div class="mt-8 p-6 bg-blue-50 rounded-xl border border-blue-200">
            <h2 class="text-2xl font-bold mb-6 text-center text-blue-700">Understanding K-Means Clustering</h2>
            <p class="mb-4 text-gray-700">
                K-Means is an unsupervised learning algorithm that aims to partition $$n$$ observations into $$k$$ clusters in which each observation belongs to the cluster with the nearest mean (centroid), serving as a prototype of the cluster.
            </p>
            <h3 class="text-xl font-semibold mb-2">How K-Means Algorithm Works:</h3>
            <ol class="list-decimal list-inside text-gray-700 mb-4">
                <li class="mb-2">
                    <span class="highlight-bold">1. Initialization:</span> Randomly select $$K$$ data points from the dataset as initial centroids.
                </li>
                <li class="mb-2">
                    <span class="highlight-bold">2. Assignment Step:</span> Assign each data point to the cluster whose centroid is closest to it (based on Euclidean distance).
                </li>
                <li class="mb-2">
                    <span class="highlight-bold">3. Update Step:</span> Recalculate the centroids by taking the mean of all data points assigned to that cluster.
                </li>
                <li class="mb-2">
                    <span class="highlight-bold">4. Iteration:</span> Repeat steps 2 and 3 until the centroids no longer move significantly or a maximum number of iterations is reached.
                </li>
            </ol>
            <p class="mb-4 text-gray-700">
                The objective of K-Means is to minimize the sum of squared distances between data points and their assigned cluster's centroid, also known as the within-cluster sum of squares (WCSS).
            </p>
            <div class="math">
                $$ \min_{C} \sum_{i=1}^{k} \sum_{x \in C_i} \|x - \mu_i\|^2 $$
            </div>
            <p class="mb-4 text-gray-700">
                Where:
            </p>
            <ul class="list-disc list-inside text-gray-700 mb-4">
                <li>$$C$$ is the set of clusters.</li>
                <li>$$C_i$$ is the $$i$$-th cluster.</li>
                <li>$$x$$ is a data point.</li>
                <li>$$\mu_i$$ is the centroid of cluster $$C_i$$.</li>
            </ul>
            <h3 class="text-xl font-semibold mb-2">Key Concepts of K-Means:</h3>
            <ul class="list-disc list-inside text-gray-700 mb-4">
                <li class="mb-2">
                    <span class="highlight-bold">Centroid:</span> The mean position of all data points in a cluster.
                </li>
                <li class="mb-2">
                    <span class="highlight-bold">K:</span> The number of clusters to form. This is a hyperparameter that must be chosen before running the algorithm.
                </li>
                <li class="mb-2">
                    <span class="highlight-bold">Voronoi Diagram:</span> The partitioning of the plane into regions based on distance to points in a specific subset of the plane. In K-Means, these regions represent the areas where new points would be assigned to a particular cluster.
                </li>
            </ul>
        </div>
    </div>
</div>

<script>
    // --- KMeans Class (from your kmeans.js) ---
    class KMeans {
        constructor(k = 3, maxIterations = 100) {
            this.k = k;
            this.maxIterations = maxIterations;
            this.centroids = []; // Array of centroid coordinates
            this.clusterAssignments = []; // Array to store assigned cluster index for each data point
            this.dimensions = 0;
        }

        /**
         * Calculates Euclidean distance between two points.
         * @param {Array<number>} p1
         * @param {Array<number>} p2
         * @returns {number} Distance
         */
        _euclideanDistance(p1, p2) {
            let sum = 0;
            for (let i = 0; i < p1.length; i++) {
                sum += Math.pow(p1[i] - p2[i], 2);
            }
            return Math.sqrt(sum);
        }

        /**
         * Initializes centroids randomly from the data points.
         * @param {Array<Array<number>>} X - Array of feature vectors.
         */
        _initializeCentroids(X) {
            this.centroids = [];
            const dataIndices = Array.from({ length: X.length }, (_, i) => i);
            // Shuffle indices and pick k unique ones
            // Ensure k is not greater than X.length
            const actualK = Math.min(this.k, X.length);
            for (let i = 0; i < actualK; i++) {
                const randomIndex = Math.floor(Math.random() * dataIndices.length);
                this.centroids.push([...X[dataIndices[randomIndex]]]); // Deep copy
                dataIndices.splice(randomIndex, 1); // Remove chosen index to ensure uniqueness
            }
        }

        /**
         * Assigns each data point to the closest centroid.
         * @param {Array<Array<number>>} X - Array of feature vectors.
         * @returns {boolean} True if assignments changed, false otherwise.
         */
        _assignClusters(X) {
            let changed = false;
            const newAssignments = Array(X.length).fill(-1); // Initialize with -1

            for (let i = 0; i < X.length; i++) {
                let minDistance = Infinity;
                let closestCentroidIndex = -1;

                for (let j = 0; j < this.centroids.length; j++) { // Use this.centroids.length
                    const dist = this._euclideanDistance(X[i], this.centroids[j]);
                    if (dist < minDistance) {
                        minDistance = dist;
                        closestCentroidIndex = j;
                    }
                }
                newAssignments[i] = closestCentroidIndex;

                if (newAssignments[i] !== this.clusterAssignments[i]) {
                    changed = true;
                }
            }
            this.clusterAssignments = newAssignments;
            return changed;
        }

        /**
         * Updates centroid positions based on the mean of assigned points.
         * @param {Array<Array<number>>} X - Array of feature vectors.
         */
        _updateCentroids(X) {
            const newCentroidSums = Array(this.k).fill(0).map(() => Array(this.dimensions).fill(0));
            const newCentroidCounts = Array(this.k).fill(0);

            for (let i = 0; i < X.length; i++) {
                const cluster = this.clusterAssignments[i];
                if (cluster !== -1) { // Only include points that were assigned to a cluster
                    for (let j = 0; j < this.dimensions; j++) {
                        newCentroidSums[cluster][j] += X[i][j];
                    }
                    newCentroidCounts[cluster]++;
                }
            }

            for (let j = 0; j < this.k; j++) {
                if (newCentroidCounts[j] > 0) { // Avoid division by zero if a centroid has no points
                    for (let dim = 0; dim < this.dimensions; dim++) {
                        this.centroids[j][dim] = newCentroidSums[j][dim] / newCentroidCounts[j];
                    }
                } else {
                    // Handle empty cluster: re-initialize this centroid to a random data point
                    console.warn(`Centroid ${j} became empty. Re-initializing.`);
                    const randomIndex = Math.floor(Math.random() * X.length);
                    if (X.length > 0) {
                        this.centroids[j] = [...X[randomIndex]];
                    }
                }
            }
        }

        /**
         * Runs the K-Means clustering algorithm.
         * @param {Array<Array<number>>} X - Array of feature vectors.
         * @returns {Array<number>} The cluster assignments for each data point.
         */
        fit(X) {
            if (X.length === 0) {
                console.warn("No data for K-Means clustering.");
                this.centroids = [];
                this.clusterAssignments = [];
                return [];
            }

            this.dimensions = X[0].length;
            this._initializeCentroids(X);

            for (let i = 0; i < this.maxIterations; i++) {
                const assignmentsChanged = this._assignClusters(X);
                this._updateCentroids(X);
                if (!assignmentsChanged) {
                    console.log(`K-Means converged after ${i + 1} iterations.`);
                    break;
                }
            }
            return this.clusterAssignments;
        }

        /**
         * Predicts the cluster for a single new data point.
         * @param {Array<number>} observation - The feature vector to classify.
         * @returns {number} The predicted cluster index.
         */
        predict(observation) {
            if (this.centroids.length === 0) {
                console.warn("K-Means model not trained yet.");
                return null;
            }
            let minDistance = Infinity;
            let closestCentroidIndex = -1;
            for (let j = 0; j < this.centroids.length; j++) {
                const dist = this._euclideanDistance(observation, this.centroids[j]);
                if (dist < minDistance) {
                    minDistance = dist;
                    closestCentroidIndex = j;
                }
            }
            return closestCentroidIndex;
        }

        /**
         * Generates data for visualizing Voronoi tessellation (cluster regions) in 2D.
         * Done by classifying points on a grid.
         * @param {number} xMin
         * @param {number} xMax
         * @param {number} yMin
         * @param {number} yMax
         * @param {number} resolution
         * @returns {Array<Object>} Trace for the heatmap.
         */
        generateClusterRegions2D(xMin, xMax, yMin, yMax, resolution = 50) {
            if (this.centroids.length === 0 || this.dimensions < 2) return [];

            const x_values = Array.from({ length: resolution }, (_, i) => xMin + (xMax - xMin) * i / (resolution - 1));
            const y_values = Array.from({ length: resolution }, (_, i) => yMin + (yMax - yMin) * i / (resolution - 1));

            const z_values = Array(resolution).fill(0).map(() => Array(resolution).fill(0));

            for (let i = 0; i < resolution; i++) {
                for (let j = 0; j < resolution; j++) {
                    const x = x_values[j];
                    const y = y_values[i];
                    const predictedCluster = this.predict([x, y]);
                    z_values[i][j] = predictedCluster !== null ? predictedCluster : -1;
                }
            }

            return [{
                z: z_values,
                x: x_values,
                y: y_values,
                type: 'heatmap',
                colorscale: 'Plotly3', // Use a standard Plotly colorscale that works well with indices
                showscale: false,
                opacity: 0.3,
                hoverinfo: 'skip'
            }];
        }

        /**
         * Generates data for visualizing cluster regions in 3D.
         * Similar to 2D, we classify points on a 3D grid/volume.
         * @param {number} xMin
         * @param {number} xMax
         * @param {number} yMin
         * @param {number} yMax
         * @param {number} zMin
         * @param {number} zMax
         * @param {number} resolution
         * @returns {Array<Object>} Traces for cluster regions (as transparent scatter points).
         */
        generateClusterRegions3D(xMin, xMax, yMin, yMax, zMin, zMax, resolution = 15) {
            if (this.centroids.length === 0 || this.dimensions < 3) return [];

            const x_grid = Array.from({ length: resolution }, (_, i) => xMin + (xMax - xMin) * i / (resolution - 1));
            const y_grid = Array.from({ length: resolution }, (_, i) => yMin + (yMax - yMin) * i / (resolution - 1));
            const z_grid = Array.from({ length: resolution }, (_, i) => zMin + (zMax - zMin) * i / (resolution - 1));

            const classifiedPointsX = Array.from({ length: this.k }, () => []);
            const classifiedPointsY = Array.from({ length: this.k }, () => []);
            const classifiedPointsZ = Array.from({ length: this.k }, () => []);

            for (let i = 0; i < resolution; i++) {
                for (let j = 0; j < resolution; j++) {
                    for (let l = 0; l < resolution; l++) {
                        const x = x_grid[i];
                        const y = y_grid[j];
                        const z = z_grid[l];
                        const predictedCluster = this.predict([x, y, z]);

                        if (predictedCluster !== null) {
                            classifiedPointsX[predictedCluster].push(x);
                            classifiedPointsY[predictedCluster].push(y);
                            classifiedPointsZ[predictedCluster].push(z);
                        }
                    }
                }
            }

            const regionTraces = Array.from({ length: this.k }, (_, clusterIndex) => ({
                x: classifiedPointsX[clusterIndex],
                y: classifiedPointsY[clusterIndex],
                z: classifiedPointsZ[clusterIndex],
                mode: 'markers',
                type: 'scatter3d',
                marker: {
                    size: 2,
                    opacity: 0.05, // Very transparent to show the region
                    color: clusterIndex, // Use index for color mapping
                    colorscale: 'Plotly3', // Consistent colorscale with 2D heatmap
                    cmin: 0,
                    cmax: this.k - 1
                },
                name: `Cluster Region ${clusterIndex}`,
                hoverinfo: 'skip'
            }));
            return regionTraces;
        }
    }

    // --- Main K-Means Logic ---

    document.addEventListener('DOMContentLoaded', () => {
        const plotlyGraph = document.getElementById('plotly-graph');
        const dimensionsSelect = document.getElementById('dimensions');
        const numClustersInput = document.getElementById('num-clusters');
        const maxIterationsInput = document.getElementById('max-iterations');
        const dataPointsTotalInput = document.getElementById('data-points-total');
        const addPointBtn = document.getElementById('add-point-btn');
        const resetDataBtn = document.getElementById('reset-data-btn');
        const newPointXInput = document.getElementById('new-point-x');
        const newPointYInput = document.getElementById('new-point-y');
        const newPointZInput = document.getElementById('new-point-z');
        const newPointYWrapper = document.getElementById('new-point-y-wrapper');
        const newPointZWrapper = document.getElementById('new-point-z-wrapper');
        const predictionResultDisplay = document.getElementById('prediction-result');

        let currentData = []; // Stores { x, y, (z) } - no initial 'class' here
        let kmeansModel;
        let currentDimensions = dimensionsSelect.value; // "2D" or "3D"
        const clusterColorscale = 'Plotly3'; // Use a standard Plotly colorscale that works well with indices

        // --- Helper Functions ---

        function generateRandomData(totalPoints, dimensions, numClustersHint = 3) {
            const data = [];
            // Generate points around 'numClustersHint' distinct centers to make clustering visible
            const centers = Array.from({ length: numClustersHint }, () => ({
                x: (Math.random() - 0.5) * 10,
                y: (Math.random() - 0.5) * 10,
                z: (Math.random() - 0.5) * 10
            }));

            for (let i = 0; i < totalPoints; i++) {
                // Assign points to 'hint' clusters for initial visual separation
                const center = centers[Math.floor(Math.random() * numClustersHint)];
                const x = center.x + (Math.random() - 0.5) * 3; // Add some spread
                const y = center.y + (Math.random() - 0.5) * 3;

                if (dimensions === "2D") {
                    data.push({ x: x, y: y });
                } else { // 3D
                    const z = center.z + (Math.random() - 0.5) * 3;
                    data.push({ x: x, y: y, z: z });
                }
            }
            return data;
        }

        function prepareDataForModel(data) {
            return data.map(d => {
                if (currentDimensions === "2D") {
                    return [d.x, d.y];
                } else {
                    return [d.x, d.y, d.z];
                }
            });
        }

        function createPlotlyTraces(data, clusterAssignments, centroids, type) {
            const traces = [];
            const k = parseInt(numClustersInput.value);

            // 1. Data Points, colored by cluster assignment
            for (let cls = 0; cls < k; cls++) {
                const classData = data.filter((_, i) => clusterAssignments[i] === cls);
                if (type === "2D") {
                    traces.push({
                        x: classData.map(d => d.x),
                        y: classData.map(d => d.y),
                        mode: 'markers',
                        type: 'scatter',
                        name: `Cluster ${cls}`,
                        marker: {
                            color: cls, // Use cluster index for color
                            colorscale: clusterColorscale,
                            cmin: 0, cmax: k - 1, // Ensure colorscale maps correctly
                            size: 8,
                            line: {
                                color: 'white',
                                width: 1
                            }
                        },
                        legendgroup: `cluster${cls}`, // Group in legend
                        showlegend: true
                    });
                } else { // 3D
                    traces.push({
                        x: classData.map(d => d.x),
                        y: classData.map(d => d.y),
                        z: classData.map(d => d.z),
                        mode: 'markers',
                        type: 'scatter3d',
                        name: `Cluster ${cls}`,
                        marker: {
                            color: cls, // Use cluster index for color
                            colorscale: clusterColorscale,
                            cmin: 0, cmax: k - 1,
                            size: 6,
                            line: {
                                color: 'white',
                                width: 1
                            }
                        },
                        legendgroup: `cluster${cls}`,
                        showlegend: true
                    });
                }
            }

            // 2. Centroids
            if (centroids && centroids.length > 0) {
                if (type === "2D") {
                    traces.push({
                        x: centroids.map(c => c[0]),
                        y: centroids.map(c => c[1]),
                        mode: 'markers',
                        type: 'scatter',
                        name: 'Centroids',
                        marker: {
                            symbol: 'x', // 'x' marker for centroids
                            color: 'black',
                            size: 10,
                            line: {
                                color: 'white',
                                width: 2
                            }
                        },
                        hoverinfo: 'text',
                        text: centroids.map((c, i) => `Centroid ${i}: (${c[0].toFixed(2)}, ${c[1].toFixed(2)})`),
                        showlegend: true
                    });
                } else { // 3D
                    traces.push({
                        x: centroids.map(c => c[0]),
                        y: centroids.map(c => c[1]),
                        z: centroids.map(c => c[2]),
                        mode: 'markers',
                        type: 'scatter3d',
                        name: 'Centroids',
                        marker: {
                            symbol: 'x',
                            color: 'black',
                            size: 8,
                            line: {
                                color: 'white',
                                width: 1
                            }
                        },
                        hoverinfo: 'text',
                        text: centroids.map((c, i) => `Centroid ${i}: (${c[0].toFixed(2)}, ${c[1].toFixed(2)}, ${c[2].toFixed(2)})`),
                        showlegend: true
                    });
                }
            }
            return traces;
        }

        function getAxisRanges(data, dimensions) {
            if (data.length === 0) {
                return {
                    x: [-10, 10],
                    y: [-10, 10],
                    z: [-10, 10]
                };
            }

            const minX = Math.min(...data.map(d => d.x)) - 2;
            const maxX = Math.max(...data.map(d => d.x)) + 2;
            const minY = Math.min(...data.map(d => d.y)) - 2;
            const maxY = Math.max(...data.map(d => d.y)) + 2;

            if (dimensions === "2D") {
                return {
                    x: [minX, maxX],
                    y: [minY, maxY]
                };
            } else { // 3D
                const minZ = Math.min(...data.map(d => d.z)) - 2;
                const maxZ = Math.max(...data.map(d => d.z)) + 2;
                return {
                    x: [minX, maxX],
                    y: [minY, maxY],
                    z: [minZ, maxZ]
                };
            }
        }

        function updateGraph() {
            predictionResultDisplay.innerText = ''; // Clear previous prediction result

            const X = prepareDataForModel(currentData);
            if (X.length === 0) {
                Plotly.purge(plotlyGraph); // Clear graph if no data
                return;
            }

            // Initialize K-Means model with current parameters
            kmeansModel = new KMeans(
                parseInt(numClustersInput.value),
                parseInt(maxIterationsInput.value)
            );
            const clusterAssignments = kmeansModel.fit(X);

            // Update currentData with assigned clusters for plotting
            currentData.forEach((d, i) => {
                d.cluster = clusterAssignments[i];
            });

            const plotlyTraces = createPlotlyTraces(currentData, clusterAssignments, kmeansModel.centroids, currentDimensions);
            let layout;

            const ranges = getAxisRanges(currentData, currentDimensions);

            if (currentDimensions === "2D") {
                // Generate cluster regions (Voronoi) for 2D
                const regionTraces = kmeansModel.generateClusterRegions2D(
                    ranges.x[0], ranges.x[1],
                    ranges.y[0], ranges.y[1]
                );
                plotlyTraces.unshift(...regionTraces); // Add regions as background

                layout = {
                    title: 'K-Means Clustering (2D)',
                    xaxis: { title: 'Feature 1', range: ranges.x },
                    yaxis: { title: 'Feature 2', range: ranges.y },
                    hovermode: 'closest',
                    showlegend: true
                };
                Plotly.newPlot(plotlyGraph, plotlyTraces, layout);

            } else { // 3D
                // Generate cluster regions for 3D
                const regionTraces = kmeansModel.generateClusterRegions3D(
                    ranges.x[0], ranges.x[1],
                    ranges.y[0], ranges.y[1],
                    ranges.z[0], ranges.z[1]
                );
                plotlyTraces.unshift(...regionTraces); // Add regions as background

                layout = {
                    title: 'K-Means Clustering (3D)',
                    scene: {
                        xaxis: { title: 'Feature 1', range: ranges.x },
                        yaxis: { title: 'Feature 2', range: ranges.y },
                        zaxis: { title: 'Feature 3', range: ranges.z },
                        aspectmode: 'cube'
                    },
                    hovermode: 'closest',
                    showlegend: true
                };
                Plotly.newPlot(plotlyGraph, plotlyTraces, layout);
            }
        }

        // --- Event Listeners ---

        dimensionsSelect.addEventListener('change', (event) => {
            currentDimensions = event.target.value;
            if (currentDimensions === "2D") {
                newPointZWrapper.classList.add('hidden');
            } else {
                newPointZWrapper.classList.remove('hidden');
            }
            // Regenerate initial data for new dimensions
            currentData = generateRandomData(
                parseInt(dataPointsTotalInput.value),
                currentDimensions,
                parseInt(numClustersInput.value) // Pass K as a hint for data generation
            );
            updateGraph();
        });

        numClustersInput.addEventListener('change', () => {
            // When K changes, re-generate data (using new K as hint) and re-cluster
            currentData = generateRandomData(
                parseInt(dataPointsTotalInput.value),
                currentDimensions,
                parseInt(numClustersInput.value)
            );
            updateGraph();
        });

        maxIterationsInput.addEventListener('change', updateGraph);

        dataPointsTotalInput.addEventListener('change', () => {
            currentData = generateRandomData(
                parseInt(dataPointsTotalInput.value),
                currentDimensions,
                parseInt(numClustersInput.value) // Pass K as a hint for data generation
            );
            updateGraph();
        });


        addPointBtn.addEventListener('click', () => {
            const x = parseFloat(newPointXInput.value);
            const y = parseFloat(newPointYInput.value);
            let newPointFeatures;

            if (currentDimensions === "2D") {
                newPointFeatures = [x, y];
            } else { // 3D
                const z = parseFloat(newPointZInput.value);
                newPointFeatures = [x, y, z];
            }

            const predictedCluster = kmeansModel.predict(newPointFeatures);

            if (predictedCluster !== null) {
                const newPointData = { x: x, y: y, cluster: predictedCluster };
                if (currentDimensions === "3D") {
                    newPointData.z = newPointFeatures[2];
                }
                currentData.push(newPointData);
                updateGraph(); // This will re-run K-Means on the updated data
                predictionResultDisplay.innerText = `New point (${newPointFeatures.join(', ')}) assigned to: Cluster ${predictedCluster}`;
            } else {
                predictionResultDisplay.innerText = "K-Means model not trained yet or an error occurred. Try resetting data.";
            }
        });

        resetDataBtn.addEventListener('click', () => {
            currentData = generateRandomData(
                parseInt(dataPointsTotalInput.value),
                currentDimensions,
                parseInt(numClustersInput.value) // Pass K as a hint for data generation
            );
            updateGraph();
        });

        // Initial setup
        currentData = generateRandomData(
            parseInt(dataPointsTotalInput.value),
            currentDimensions,
            parseInt(numClustersInput.value)
        );
        updateGraph();
        MathJax.typeset(); // Ensure MathJax equations are rendered on initial load
    });
</script>
{% endblock %}