<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Apple Clustering with TensorFlow.js</title>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@latest"></script>
    <style>
        body { font-family: Arial, sans-serif; margin: 20px; background-color: #f0f0f0; }
        #canvas { border: 2px solid #333; background-color: white; margin: 20px 0; }
        .controls { margin: 10px 0; }
        .controls label { display: inline-block; width: 150px; }
        .controls input { width: 60px; margin-right: 10px; }
        button { padding: 10px 20px; margin: 5px; cursor: pointer; background-color: #4CAF50; color: white; border: none; border-radius: 4px; }
        button:hover { background-color: #45a049; }
        #info { margin-top: 10px; padding: 10px; background-color: #e0e0e0; border-radius: 4px; }
    </style>
</head>
<body>
    <h1>Apple Clustering with Gradient Descent</h1>
    <div class="controls">
        <label>Number of dots:</label>
        <input type="number" id="numDots" value="100" min="4" max="500">
        <label>Number of apples:</label>
        <input type="number" id="numApples" value="4" min="1" max="10">
    </div>

    <div class="controls">
        <label>Learning rate:</label>
        <input type="number" id="learningRate" value="0.01" step="0.001" min="0.001" max="0.1">
        <label>Iterations:</label>
        <input type="number" id="iterations" value="100" min="10" max="1000">
    </div>

    <button onclick="initialize()">Initialize</button>
    <button onclick="runOptimization()">Run Optimization</button>
    <button onclick="stopOptimization()">Stop</button>

    <canvas id="canvas" width="600" height="600"></canvas>

    <div id="info">
        <div>Current Loss: <span id="currentLoss">-</span></div>
        <div>Iteration: <span id="currentIteration">-</span></div>
        <div>Reassignments: <span id="reassignments">-</span></div>
    </div>

    <script>
        const canvas = document.getElementById('canvas');
        const ctx = canvas.getContext('2d');
        const canvasSize = 600;
        
        let dots = [];
        let apples = [];
        let assignments = [];
        let isRunning = false;
        let animationId = null;
        
        function initialize() {
            const numDots = parseInt(document.getElementById('numDots').value);
            const numApples = parseInt(document.getElementById('numApples').value);
            
            if (numDots % numApples !== 0) {
                alert('Number of dots must be divisible by number of apples!');
                return;
            }
            
            // Generate random dots
            dots = [];
            for (let i = 0; i < numDots; i++) {
                dots.push({
                    x: Math.random(),
                    y: Math.random()
                });
            }
            
            // Generate random apple positions
            apples = [];
            for (let i = 0; i < numApples; i++) {
                apples.push({
                    x: Math.random(),
                    y: Math.random(),
                    vx: 0,
                    vy: 0
                });
            }
            
            // Random initial assignments ensuring equal distribution
            assignments = [];
            const dotsPerApple = numDots / numApples;
            for (let i = 0; i < numApples; i++) {
                for (let j = 0; j < dotsPerApple; j++) {
                    assignments.push(i);
                }
            }
            // Shuffle assignments
            for (let i = assignments.length - 1; i > 0; i--) {
                const j = Math.floor(Math.random() * (i + 1));
                [assignments[i], assignments[j]] = [assignments[j], assignments[i]];
            }
            
            draw();
            updateInfo();
        }
        
        function distance(p1, p2) {
            const dx = p1.x - p2.x;
            const dy = p1.y - p2.y;
            return Math.sqrt(dx * dx + dy * dy);
        }
        
        function calculateLoss() {
            let totalLoss = 0;
            const numApples = apples.length;
            for (let i = 0; i < dots.length; i++) {
                const assignedAppleIdx = assignments[i];
                const dot = dots[i];
                const d1 = distance(dot, apples[assignedAppleIdx]);
                let sumOtherDistances = 0;
                let count = 0;
                
                for (let j = 0; j < numApples; j++) {
                    if (j !== assignedAppleIdx) {
                        sumOtherDistances += distance(dot, apples[j]);
                        count++;
                    }
                }
                
                const d2 = count > 0 ? sumOtherDistances / count : 0;
                totalLoss += (d1 - d2);
            }
            return totalLoss / dots.length;
        }
        
        function calculateGradients() {
            const numApples = apples.length;
            const gradients = new Array(numApples).fill(null).map(() => ({ x: 0, y: 0 }));
            
            for (let j = 0; j < numApples; j++) {
                const appleJ = apples[j];
                for (let i = 0; i < dots.length; i++) {
                    const assignedAppleIdx = assignments[i];
                    const dot = dots[i];
                    
                    if (j === assignedAppleIdx) {
                        const d1 = distance(dot, appleJ);
                        if (d1 > 0) {
                            gradients[j].x += (appleJ.x - dot.x) / d1;
                            gradients[j].y += (appleJ.y - dot.y) / d1;
                        }
                    } else {
                        const d_ik = distance(dot, appleJ);
                        if (d_ik > 0) {
                            const factor = -1 / (numApples - 1);
                            gradients[j].x += factor * (appleJ.x - dot.x) / d_ik;
                            gradients[j].y += factor * (appleJ.y - dot.y) / d_ik;
                        }
                    }
                }
            }
            
            return gradients;
        }
        
        function reassignDots() {
            const numDots = dots.length;
            const numApples = apples.length;
            const dotsPerApple = numDots / numApples;
            
            // Calculate distance matrix
            const distances = [];
            for (let i = 0; i < numDots; i++) {
                distances[i] = [];
                for (let j = 0; j < numApples; j++) {
                    distances[i][j] = distance(dots[i], apples[j]);
                }
            }
            
            // Create preference lists for each dot (sorted by distance)
            const preferences = [];
            for (let i = 0; i < numDots; i++) {
                preferences[i] = [];
                for (let j = 0; j < numApples; j++) {
                    preferences[i].push({ apple: j, distance: distances[i][j] });
                }
                preferences[i].sort((a, b) => a.distance - b.distance);
            }
            
            // Hungarian algorithm-inspired assignment
            const newAssignments = new Array(numDots).fill(-1);
            const appleCounts = new Array(numApples).fill(0);
            
            // First pass: assign dots to their closest available apple
            const dotIndices = Array.from({length: numDots}, (_, i) => i);
            dotIndices.sort((a, b) => preferences[a][0].distance - preferences[b][0].distance);
            
            for (const dotIdx of dotIndices) {
                for (const pref of preferences[dotIdx]) {
                    if (appleCounts[pref.apple] < dotsPerApple) {
                        newAssignments[dotIdx] = pref.apple;
                        appleCounts[pref.apple]++;
                        break;
                    }
                }
            }
            
            // Second pass: optimize by swapping
            let improved = true;
            while (improved) {
                improved = false;
                
                for (let i = 0; i < numDots; i++) {
                    for (let j = i + 1; j < numDots; j++) {
                        if (newAssignments[i] === newAssignments[j]) continue;
                        
                        const currentCost = distances[i][newAssignments[i]] + distances[j][newAssignments[j]];
                        const swapCost = distances[i][newAssignments[j]] + distances[j][newAssignments[i]];
                        
                        if (swapCost < currentCost) {
                            [newAssignments[i], newAssignments[j]] = [newAssignments[j], newAssignments[i]];
                            improved = true;
                        }
                    }
                }
            }
            
            // Count reassignments
            let reassignmentCount = 0;
            for (let i = 0; i < numDots; i++) {
                if (assignments[i] !== newAssignments[i]) {
                    reassignmentCount++;
                }
            }
            
            assignments = newAssignments;
            return reassignmentCount;
        }
        
        async function runOptimization() {
            if (isRunning) return;
            isRunning = true;
            
            const learningRate = parseFloat(document.getElementById('learningRate').value);
            const iterations = parseInt(document.getElementById('iterations').value);
            
            for (let iter = 0; iter < iterations && isRunning; iter++) {
                // Calculate gradients
                const gradients = calculateGradients();
                
                // Update apple positions with momentum
                const momentum = 0.9;
                for (let i = 0; i < apples.length; i++) {
                    apples[i].vx = momentum * apples[i].vx - learningRate * gradients[i].x;
                    apples[i].vy = momentum * apples[i].vy - learningRate * gradients[i].y;
                    
                    apples[i].x += apples[i].vx;
                    apples[i].y += apples[i].vy;
                    
                    // Keep apples within bounds
                    apples[i].x = Math.max(0, Math.min(1, apples[i].x));
                    apples[i].y = Math.max(0, Math.min(1, apples[i].y));
                }
                
                // Reassign all dots optimally
                const reassignmentCount = reassignDots();
                
                // Update visualization
                draw();
                document.getElementById('currentIteration').textContent = iter + 1;
                document.getElementById('reassignments').textContent = reassignmentCount;
                updateInfo();
                
                // Allow UI to update
                await new Promise(resolve => setTimeout(resolve, 10));
            }
            
            isRunning = false;
        }
        
        function stopOptimization() {
            isRunning = false;
        }
        
        function draw() {
            ctx.clearRect(0, 0, canvasSize, canvasSize);
            
            // Create color palette for apples
            const colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FECA57', '#FF9FF3'];
            
            // Draw connections
            ctx.globalAlpha = 0.2;
            for (let i = 0; i < dots.length; i++) {
                const dot = dots[i];
                const appleIdx = assignments[i];
                const apple = apples[appleIdx];
                
                ctx.beginPath();
                ctx.moveTo(dot.x * canvasSize, dot.y * canvasSize);
                ctx.lineTo(apple.x * canvasSize, apple.y * canvasSize);
                ctx.strokeStyle = colors[appleIdx % colors.length];
                ctx.stroke();
            }
            ctx.globalAlpha = 1;
            
            // Draw dots
            for (let i = 0; i < dots.length; i++) {
                const dot = dots[i];
                const appleIdx = assignments[i];
                
                ctx.beginPath();
                ctx.arc(dot.x * canvasSize, dot.y * canvasSize, 3, 0, 2 * Math.PI);
                ctx.fillStyle = colors[appleIdx % colors.length];
                ctx.fill();
            }
            
            // Draw apples
            ctx.font = '30px Arial';
            for (let i = 0; i < apples.length; i++) {
                const apple = apples[i];
                ctx.fillText('🍎', apple.x * canvasSize - 15, apple.y * canvasSize + 10);
            }
        }
        
        function updateInfo() {
            const loss = calculateLoss();
            document.getElementById('currentLoss').textContent = loss.toFixed(4);
        }
        
        // Initialize on load
        window.onload = initialize;
    </script>
</body>
</html>

can you make it so that the loss is just mean for all apples of (average distance of the apple from each dot, but the dots are weighted based on the dots distance from its nearest apple - not including the apple being measured) This way, it encourages apples to chase dots that are not close to any other apples