static/js/simulation.js

class BrownianSimulation {
    constructor(canvas) {
        this.canvas = canvas;
        this.ctx = canvas.getContext('2d');
        this.width = canvas.width;
        this.height = canvas.height;
        
        // Load mosquito images
        this.mosquitoImage = new Image();
        this.mosquitoImage.src = '/static/images/mosquito_1.png';
        this.mosquitoSize = this.width * 0.03; // 3x actual mosquito size
        this.mosquitoType = 'mosquito_1'; // Default mosquito type
        
        // Set up mosquito selection
        this.setupMosquitoSelection();
        
        // Load trap image
        this.trapImage = new Image();
        this.trapImage.src = '/static/images/tim.png';
        this.defaultTrapImageSrc = '/static/images/tim.png';
        this.trapSize = this.width * 0.1; // Actual trap size
        
        // Set up trap image upload
        this.setupTrapImageUpload();
        
        // Set up trap selection
        this.setupTrapSelection();
        
        // Potential properties
        this.potentialCenter = {
            x: this.width / 2,
            y: this.height / 2
        };
        
        // Add click event listener
        this.canvas.addEventListener('click', (e) => {
            const rect = this.canvas.getBoundingClientRect();
            this.potentialCenter.x = e.clientX - rect.left;
            this.potentialCenter.y = e.clientY - rect.top;
        });
        
        // Remove these as they'll come from Python
        this.x = this.width / 2;
        this.y = this.height / 2;
        
        // Update API endpoint to be more specific
        const protocol = window.location.protocol;
        const host = window.location.host;
        this.apiUrl = `${protocol}//${host}/simulate_multiple_steps`;
        
        // Initialize error message element
        this.errorMessage = document.getElementById('error-message');
        
        // Initialize particles array
        this.particles = [];
        
        // Store complete trajectory history for download
        this.completeTrajectoryData = [];
        this.trajectoryTimeOffset = 0;
        this.cueTypeHistory = []; // Track cue type changes over time
        this.lastCueType = 'visualco2'; // Default cue type
        
        // Add number of particles control
        this.numParticlesInput = document.getElementById('numParticles');
        this.numParticlesValue = document.getElementById('numParticlesValue');
        
        this.numParticlesInput.addEventListener('input', (e) => {
            this.numParticles = parseInt(e.target.value);
            this.numParticlesValue.textContent = this.numParticles;
            // Reset trajectory data when number of mosquitoes changes
            this.resetTrajectoryData();
        });
        
        this.numParticles = 3; // Start with 3 particles
        
        // Create and add warning message element
        this.warningMessage = document.createElement('div');
        this.warningMessage.style.color = 'orange';
        this.warningMessage.style.display = 'none';
        this.warningMessage.style.position = 'absolute';
        this.warningMessage.style.top = '10px';
        this.warningMessage.style.left = '10px';
        document.body.appendChild(this.warningMessage);
        
        this.amountOfDataToRequest = 1000;
        this.requestFrequency = 800;
        this.targetFrameTime = 50; // 50ms = 0.05 seconds
        this.lastFrameTime = 0;
        this.currentFrameIndex = 0;
        this.pendingParticles = null;
        this.lastDataRequestTime = 0;
        
        // Add model type control
        this.setupModelTypeSelection();
        
        // Set up download trajectory functionality
        this.setupDownloadTrajectory();
    }
    
    // Add this new method to handle mosquito selection
    setupMosquitoSelection() {
        const mosquito1 = document.getElementById('mosquito1');
        
        if (mosquito1) {
            mosquito1.addEventListener('click', () => {
                this.selectMosquito('mosquito_1');
                mosquito1.classList.add('selected');
            });
        }
    }
    
    selectMosquito(type) {
        this.mosquitoType = type;
        this.mosquitoImage.src = `/static/images/${type}.png`;
    }
    
    // Add this new method to handle trap image upload
    setupTrapImageUpload() {
        const trapImageUpload = document.getElementById('trapImageUpload');
        const resetTrapImage = document.getElementById('resetTrapImage');
        
        if (trapImageUpload && resetTrapImage) {
            trapImageUpload.addEventListener('change', (e) => {
                const file = e.target.files[0];
                if (file) {
                    const reader = new FileReader();
                    reader.onload = (event) => {
                        this.trapImage.src = event.target.result;
                        // Clear selection when custom image is uploaded
                        this.clearTrapSelection();
                    };
                    reader.readAsDataURL(file);
                }
            });
            
            resetTrapImage.addEventListener('click', () => {
                this.trapImage.src = this.defaultTrapImageSrc;
                trapImageUpload.value = ''; // Clear the file input
                // Select the default trap option
                this.selectTrap('tim');
            });
        }
    }
    
    // Add this new method to handle trap selection
    setupTrapSelection() {
        const trapTim = document.getElementById('trap-tim');
        const trapCircle = document.getElementById('trap-circle');
        
        if (trapTim && trapCircle) {
            trapTim.addEventListener('click', () => {
                this.selectTrap('tim');
            });
            
            trapCircle.addEventListener('click', () => {
                this.selectTrap('circle');
            });
        }
    }
    
    selectTrap(type) {
        // Clear all selections
        this.clearTrapSelection();
        
        // Set the selected trap
        switch(type) {
            case 'tim':
                this.trapImage.src = '/static/images/tim.png';
                document.getElementById('trap-tim').classList.add('selected');
                break;
            case 'circle':
                this.trapImage.src = '/static/images/circle.png';
                document.getElementById('trap-circle').classList.add('selected');
                break;
        }
        
        // Clear any custom uploaded image
        const trapImageUpload = document.getElementById('trapImageUpload');
        if (trapImageUpload) {
            trapImageUpload.value = '';
        }
    }
    
    clearTrapSelection() {
        const trapOptions = document.querySelectorAll('.trap-option');
        trapOptions.forEach(option => {
            option.classList.remove('selected');
        });
    }
    
    // Add this new method to handle trajectory download
    setupDownloadTrajectory() {
        const downloadBtn = document.getElementById('download-trajectory');
        if (downloadBtn) {
            downloadBtn.addEventListener('click', () => {
                this.downloadTrajectoryCSV();
            });
        }
    }
    
    downloadTrajectoryCSV() {
        if (!this.completeTrajectoryData || this.completeTrajectoryData.length === 0) {
            alert('No trajectory data available. Please run a simulation first.');
            return;
        }
        
        // Prepare CSV data
        let csvContent = 'Timepoint,MosquitoID,X,Y,VX,VY,CueType\n';
        
        // Find the maximum number of timepoints across all particles
        const maxTimepoints = Math.max(...this.completeTrajectoryData.map(p => p.x.length));
        
        // Generate timepoints based on the simulation parameters
        const timeStep = this.targetFrameTime / 1000; // Convert to seconds
        
        for (let t = 0; t < maxTimepoints; t++) {
            const timepoint = t * timeStep;
            
            // Find the cue type for this specific timepoint
            let cueTypeForTimepoint = 'visualco2'; // Default
            if (this.cueTypeHistory.length > 0) {
                // Find the cue type that was active at this timepoint
                const historyEntry = this.cueTypeHistory.find(entry => 
                    Math.abs(entry.timepoint - timepoint) < timeStep / 2
                );
                if (historyEntry) {
                    cueTypeForTimepoint = historyEntry.cueType;
                }
            }
            
            for (let mosquitoId = 0; mosquitoId < this.completeTrajectoryData.length; mosquitoId++) {
                const particle = this.completeTrajectoryData[mosquitoId];
                
                // Check if this timepoint exists for this particle
                if (t < particle.x.length) {
                    const x = particle.x[t];
                    const y = particle.y[t];
                    const vx = particle.vx[t];
                    const vy = particle.vy[t];
                    
                    csvContent += `${timepoint.toFixed(3)},${mosquitoId + 1},${x.toFixed(2)},${y.toFixed(2)},${vx.toFixed(2)},${vy.toFixed(2)},${cueTypeForTimepoint}\n`;
                }
            }
        }
        
        // Create and download the file
        const blob = new Blob([csvContent], { type: 'text/csv;charset=utf-8;' });
        const link = document.createElement('a');
        const url = URL.createObjectURL(blob);
        link.setAttribute('href', url);
        link.setAttribute('download', `mosquito_trajectory_complete_${new Date().toISOString().slice(0, 19).replace(/:/g, '-')}.csv`);
        link.style.visibility = 'hidden';
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);
    }
    
    // Add this new method to handle model type selection
    setupModelTypeSelection() {
        const modelOptions = document.querySelectorAll('.model-option');
        
        if (modelOptions.length > 0) {
            modelOptions.forEach(option => {
                option.addEventListener('click', () => {
                    // Remove selected class from all options
                    modelOptions.forEach(opt => opt.classList.remove('selected'));
                    
                    // Add selected class to clicked option
                    option.classList.add('selected');
                    
                    // Update model type
                    this.modelType = option.dataset.value;
                    
                    // Track cue type change for trajectory data
                    if (this.lastCueType !== this.modelType) {
                        this.lastCueType = this.modelType;
                    }
                });
            });
            
            // Set initial model type from the selected option
            const selectedOption = document.querySelector('.model-option.selected');
            if (selectedOption) {
                this.modelType = selectedOption.dataset.value;
            }
        }
    }
    
    // Update particle position
    async update() {
        try {
            const url = `${this.apiUrl}?centerX=${this.potentialCenter.x}&centerY=${this.potentialCenter.y}&numParticles=${this.numParticles}&model=${this.modelType}&amountOfDataToRequest=${this.amountOfDataToRequest}&targetFrameTime=${this.targetFrameTime}`;
            const response = await fetch(url, {
                method: 'GET'
            });
            if (!response.ok) {
                throw new Error(`HTTP error! status: ${response.status}`);
            }
            const data = await response.json();
            return data; // Return data instead of setting this.particles directly
            this.errorMessage.style.display = 'none';
        } catch (error) {
            console.error('Error fetching simulation data:', error);
            this.errorMessage.textContent = `Cannot connect to simulation backend: ${error.message}`;
            this.errorMessage.style.display = 'block';
            return null;
        }
    }
    
    // Draw the current state
    draw() {
        // Clear canvas
        this.ctx.clearRect(0, 0, this.width, this.height);
        
        // Draw trap (instead of red dot)
        if (this.trapImage.complete) {
            // Adjust trap size based on the current trap image
            let trapWidth = this.trapSize;
            let trapHeight = this.trapSize;
            
            // If it's the Tim trap, make height 1.6 times the width
            if (this.trapImage.src.includes('tim.png')) {
                trapHeight = this.trapSize * 1.6;
            }
            
            this.ctx.drawImage(
                this.trapImage,
                this.potentialCenter.x - trapWidth/2,
                this.potentialCenter.y - trapHeight/2,
                trapWidth,
                trapHeight
            );
        }
        
        // Draw all mosquitos
        if (this.mosquitoImage.complete && this.particles && this.particles.length > 0) {
            for (const particle of this.particles) {
                this.ctx.save();
                // Draw the current frame from the array of positions
                const currentIndex = Math.min(this.currentFrameIndex, particle.x.length - 1);
                this.ctx.translate(particle.x[currentIndex], particle.y[currentIndex]);
                
                // Use different rotation angle based on mosquito type
                const rotationAngle = this.mosquitoType === 'mosquito_1' ? Math.PI/2 : Math.PI;
                this.ctx.rotate(Math.atan2(particle.vy[currentIndex], particle.vx[currentIndex]) + rotationAngle);
                
                this.ctx.drawImage(this.mosquitoImage, -this.mosquitoSize/2, -this.mosquitoSize/2, 
                    this.mosquitoSize, this.mosquitoSize);
                this.ctx.restore();
            }
        }
    }
    
    // Modify animate to handle precise timing and frame advancement
    async animate(currentTime) {
        // Calculate time since last frame
        const deltaTime = currentTime - this.lastFrameTime;

        // Check for slow frames and display warning
        if (deltaTime >= this.targetFrameTime * 1.5) {
            this.warningMessage.textContent = `Performance warning: Frame took ${Math.round(deltaTime)}ms (target: ${this.targetFrameTime}ms)`;
            this.warningMessage.style.display = 'block';
        } else {
            this.warningMessage.style.display = 'none';
        }

        if (deltaTime >= this.targetFrameTime) {
            // Draw current frame
            this.draw();
            
            // Advance to next frame
            if (this.particles && this.particles.length > 0) {
                const maxFrames = this.particles[0].x.length;
                
                // Move to next frame or reset if at the end
                this.currentFrameIndex++;
                
                // If we've shown all frames from the current batch and have pending data
                if (this.currentFrameIndex >= maxFrames - 1 && this.pendingParticles) {
                    // Add the current batch to complete trajectory before switching
                    this.addToCompleteTrajectory(this.particles);
                    
                    // Switch to the pending data
                    this.particles = this.pendingParticles;
                    this.pendingParticles = null;
                    this.currentFrameIndex = 0;
                }
            }
            
            this.lastFrameTime = currentTime;
        }
        
        requestAnimationFrame((timestamp) => this.animate(timestamp));
    }

    async updateCallback() {
        // Only request new data if we don't already have pending data
        console.log(this.numParticles);
        if (!this.pendingParticles) {
            const newData = await this.update();
            if (newData) {
                // Check if the number of mosquitoes has changed
                if (this.particles.length > 0 && newData.length !== this.particles.length) {
                    // Number of mosquitoes changed, reset trajectory data
                    this.resetTrajectoryData();
                }
                
                // If we have no particles yet, set them directly
                if (!this.particles || this.particles.length === 0) {
                    this.particles = newData;
                    this.currentFrameIndex = 0;
                    // Initialize complete trajectory data
                    this.initializeCompleteTrajectory(newData);
                } else {
                    // Otherwise, store as pending
                    this.pendingParticles = newData;
                }
            }
        }
    }
    
    // Initialize complete trajectory data structure
    initializeCompleteTrajectory(newData) {
        this.completeTrajectoryData = [];
        this.trajectoryTimeOffset = 0;
        this.cueTypeHistory = [];
        this.lastCueType = this.modelType || 'visualco2';
        
        for (let i = 0; i < newData.length; i++) {
            this.completeTrajectoryData.push({
                x: [...newData[i].x],
                y: [...newData[i].y],
                vx: [...newData[i].vx],
                vy: [...newData[i].vy]
            });
        }
        
        // Record the initial cue type for all timepoints in this batch
        const timepointsInBatch = newData[0].x.length;
        for (let t = 0; t < timepointsInBatch; t++) {
            this.cueTypeHistory.push({
                timepoint: t * (this.targetFrameTime / 1000),
                cueType: this.lastCueType
            });
        }
    }
    
    // Add new trajectory data to the complete history
    addToCompleteTrajectory(newData) {
        if (this.completeTrajectoryData.length === 0) {
            this.initializeCompleteTrajectory(newData);
            return;
        }
        
        // Check if the number of mosquitoes has changed significantly
        // If the change is more than just adding mosquitoes, reset the trajectory
        if (newData.length < this.completeTrajectoryData.length) {
            // Number of mosquitoes decreased, reset trajectory
            this.initializeCompleteTrajectory(newData);
            return;
        }
        
        // Check if cue type has changed
        const currentCueType = this.modelType || 'visualco2';
        if (currentCueType !== this.lastCueType) {
            this.lastCueType = currentCueType;
        }
        
        // Ensure we have the same number of particles
        while (this.completeTrajectoryData.length < newData.length) {
            this.completeTrajectoryData.push({
                x: [],
                y: [],
                vx: [],
                vy: []
            });
        }
        
        // Add new data points to each particle's trajectory
        for (let i = 0; i < newData.length; i++) {
            if (i < this.completeTrajectoryData.length) {
                // Add new data points
                this.completeTrajectoryData[i].x.push(...newData[i].x);
                this.completeTrajectoryData[i].y.push(...newData[i].y);
                this.completeTrajectoryData[i].vx.push(...newData[i].vx);
                this.completeTrajectoryData[i].vy.push(...newData[i].vy);
            }
        }
        
        // Add cue type history for this batch
        const timepointsInBatch = newData[0].x.length;
        const startTime = this.trajectoryTimeOffset;
        for (let t = 0; t < timepointsInBatch; t++) {
            this.cueTypeHistory.push({
                timepoint: startTime + (t * (this.targetFrameTime / 1000)),
                cueType: this.lastCueType
            });
        }
        
        // Update time offset for the next batch
        this.trajectoryTimeOffset += newData[0].x.length * (this.targetFrameTime / 1000);
    }

    // Reset trajectory data when number of mosquitoes changes
    resetTrajectoryData() {
        this.completeTrajectoryData = [];
        this.trajectoryTimeOffset = 0;
        this.cueTypeHistory = [];
        this.lastCueType = this.modelType || 'visualco2';
        this.currentFrameIndex = 0;
        this.pendingParticles = null;
        this.particles = []; // Clear current particles as well
    }
    
    // todo:
    // - if user moves, call requestDataCallback() and updateCallback()
}

// Modify initialization
window.onload = () => {
    const canvas = document.getElementById('simulation');
    const simulation = new BrownianSimulation(canvas);
    
    // Initial data request
    simulation.updateCallback();
    
    // Start animation
    requestAnimationFrame((timestamp) => simulation.animate(timestamp));
    
    // Set up periodic data requests
    setInterval(() => simulation.updateCallback(), simulation.requestFrequency);
    
    // Proactively clean up the server-side session when the user leaves
    const sendCleanup = () => {
        const url = `${window.location.origin}/cleanup`;
        // Prefer sendBeacon for reliability during unload
        if (navigator.sendBeacon) {
            const blob = new Blob([], { type: 'application/octet-stream' });
            navigator.sendBeacon(url, blob);
        } else {
            // Fallback to a synchronous XHR as a last resort
            try {
                const xhr = new XMLHttpRequest();
                xhr.open('POST', url, false);
                xhr.setRequestHeader('Content-Type', 'text/plain');
                xhr.send('');
            } catch (e) {
                // ignore
            }
        }
    };
    
    // Use multiple event listeners to ensure cleanup happens
    window.addEventListener('pagehide', sendCleanup);
    window.addEventListener('beforeunload', sendCleanup);
    window.addEventListener('unload', sendCleanup);
    document.addEventListener('visibilitychange', () => {
        if (document.visibilityState === 'hidden') {
            sendCleanup();
        }
    });
};