// ======= Analytics & Identity Bootstrap =======
const ANALYTICS_ENDPOINT = '/api/track';
const COOKIE_NAME = 'vid';

// Generate a stable visitor id (persists across sessions)
const visitorId = (() => {
  const key = 'dp_sgd_visitor_id';
  let id = localStorage.getItem(key);
  if (!id) { 
    id = crypto.randomUUID?.() || (String(Date.now()) + Math.random().toString(16).slice(2)); 
    localStorage.setItem(key, id); 
  }
  return id;
})();

// Generate a stable session id (per browser tab/session)
const sessionId = (() => {
  const key = 'dp_sgd_session_id';
  let id = sessionStorage.getItem(key);
  if (!id) { 
    id = crypto.randomUUID?.() || (String(Date.now()) + Math.random().toString(16).slice(2)); 
    sessionStorage.setItem(key, id); 
  }
  return id;
})();

// Expose globally for compatibility with other scripts
window.__visitor_id = visitorId;
window.__session_id = sessionId;

// Minimal user context (non-PII by default). Call identify({ id, role, org, plan }) if you have a login.
let userContext = { vid: visitorId, id: null, role: null, org: null, plan: null };

async function initIdentity() {
  try {
    const r = await fetch('/api/whoami', { credentials: 'same-origin' });
    const info = await r.json();
    if (info && info.vid) userContext.vid = info.vid;
  } catch {}
}
initIdentity();

// Track page view after DOM is ready to ensure track() is defined
if (document.readyState === 'loading') {
  document.addEventListener('DOMContentLoaded', () => {
    track('page_view', { path: location.pathname, title: document.title });
  });
} else {
  track('page_view', { path: location.pathname, title: document.title });
}

function identify(user) {
  userContext = { ...userContext, ...{
    id: user?.id ?? null,
    role: user?.role ?? null,
    org: user?.org ?? null,
    plan: user?.plan ?? null,
  }};
  track('identify', { user: { id: userContext.id, role: userContext.role, org: userContext.org, plan: userContext.plan } });
}

// Fire-and-forget tracker
function track(eventType, payload = {}) {
  const body = {
    t: Date.now(),
    session_id: sessionId,        // Use snake_case to match API
    sessionId: sessionId,         // Keep camelCase for backward compatibility
    event: eventType,             // New field name
    eventType: eventType,         // Keep for backward compatibility
    path: location.pathname,
    payload,
    user: { id: userContext.id, role: userContext.role, org: userContext.org, plan: userContext.plan },
    visitor_id: userContext.vid,  // Use snake_case to match API
    vid: userContext.vid          // Keep for backward compatibility
  };
  const data = new Blob([JSON.stringify(body)], { type: 'application/json' });
  if (!(navigator.sendBeacon && navigator.sendBeacon(ANALYTICS_ENDPOINT, data))) {
    fetch(ANALYTICS_ENDPOINT, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body), credentials: 'same-origin' }).catch(()=>{});
  }
}

// Global click listener (optional; captures generic UI clicks)
document.addEventListener('click', (e) => {
  const target = e.target?.closest?.('[data-track], button, a, .nav-link, .tab');
  if (!target) return;
  const name = target.getAttribute('data-track') || target.id || target.textContent?.trim()?.slice(0, 60);
  if (!name) return;
  track('ui_click', { name });
}, { capture: true });
// ======= End Analytics Bootstrap =======

class DPSGDExplorer {
    constructor() {
        this.trainingChart = null;
        this.privacyChart = null;
        this.gradientChart = null;
        this.isTraining = false;
        this.currentView = 'epochs'; // 'epochs' or 'iterations'
        this.epochsData = [];
        this.iterationsData = [];
        this.abortController = null;  // For canceling training requests
        this.eventSource = null;      // For SSE streaming
        this.initializeUI();
    }

    initializeUI() {
        // Initialize parameter controls
        this.initializeSliders();
        this.initializePresets();
        this.initializeTabs();
        this.initializeCharts();
        
        // Add event listeners
        document.getElementById('train-button')?.addEventListener('click', () => this.toggleTraining());
        
        
        document.getElementById('train-button')?.addEventListener('click', () => { try { track('train_toggle', (this.getParameters?.()||{})); } catch (e) {} });
// Add view toggle listeners
        document.getElementById('view-epochs')?.addEventListener('click', () => this.switchView('epochs'));
        document.getElementById('view-iterations')?.addEventListener('click', () => this.switchView('iterations'));

        document.getElementById('dataset-select')?.addEventListener('change', (e) => {
            track('dataset_change', { dataset: e.target.value });
        });

        document.getElementById('model-select')?.addEventListener('change', (e) => {
            track('model_change', { model_architecture: e.target.value });
        });
        document.getElementById('preset-high-privacy')?.addEventListener('click', () => {
            track('preset_apply', { preset: 'high-privacy' });
        });
        document.getElementById('preset-balanced')?.addEventListener('click', () => {
            track('preset_apply', { preset: 'balanced' });
        });
        document.getElementById('preset-high-utility')?.addEventListener('click', () => {
            track('preset_apply', { preset: 'high-utility' });
        });
    }

    initializeSliders() {
        // Parameter sliders
        const sliders = {
            'clipping-norm': document.getElementById('clipping-norm'),
            'noise-multiplier': document.getElementById('noise-multiplier'),
            'batch-size': document.getElementById('batch-size'),
            'learning-rate': document.getElementById('learning-rate'),
            'epochs': document.getElementById('epochs')
        };

        // Add event listeners to sliders
        for (const [id, slider] of Object.entries(sliders)) {
            if (slider) {
                slider.addEventListener('input', (e) => {
                    const value = parseFloat(e.target.value);
                    document.getElementById(`${id}-value`).textContent = value.toFixed(1);
                    
                    // Update privacy budget
                    this.updatePrivacyBudget();
                    
                    
                    try { track('param_change', { param: id, value }); } catch (e) {};
// Update gradient visualization when clipping norm changes
                    if (id === 'clipping-norm') {
                        this.updateGradientVisualization(value);
                    }
                });
            }
        }

        // Add event listener for the visual clipping norm slider
        const visualSlider = document.getElementById('clipping-norm-visual');
        if (visualSlider) {
            visualSlider.addEventListener('input', (e) => {
                const value = parseFloat(e.target.value);
                document.getElementById('clipping-norm-visual-value').textContent = value.toFixed(1);
                this.updateGradientVisualization(value);
            });
        }
    }

    initializePresets() {
        // Presets based on research (Optax/TF Privacy benchmarks)
        // With proper noise scaling: noise_stddev = C * σ / batch_size
        const presets = {
            'high-privacy': {
                // Strong privacy (ε≈1-3), ~95% accuracy achievable
                // Based on: noise=1.3, clip=1.5, LR=0.25, 15 epochs → ~95%
                clippingNorm: 1.5,
                noiseMultiplier: 1.3,
                batchSize: 256,
                learningRate: 0.25,
                epochs: 30
            },
            'balanced': {
                // Moderate privacy (ε≈3-5), ~96% accuracy
                // Based on: noise=1.1, clip=1.0, LR=0.15, 60 epochs → ~96.6%
                clippingNorm: 1.0,
                noiseMultiplier: 1.1,
                batchSize: 256,
                learningRate: 0.15,
                epochs: 30
            },
            'high-utility': {
                // Lower privacy (ε≈8+), ~97% accuracy
                // Based on: noise=0.7, clip=1.5, LR=0.25, 45 epochs → ~97%
                clippingNorm: 1.5,
                noiseMultiplier: 0.7,
                batchSize: 256,
                learningRate: 0.25,
                epochs: 30
            }
        };

        // Add event listeners to preset buttons
        for (const [preset, values] of Object.entries(presets)) {
            document.getElementById(`preset-${preset}`)?.addEventListener('click', () => {
                this.applyPreset(values);
            });
        }
    }

    initializeTabs() {
        const tabs = document.querySelectorAll('.tab');
        tabs.forEach(tab => {
            
            tab.addEventListener('click', () => { try { track('tab_click', { tab: tab.dataset?.tab || tab.id || 'unknown' }); } catch (e) {} });
tab.addEventListener('click', () => {
                const tabsContainer = tab.closest('.tabs');
                tabsContainer.querySelectorAll('.tab').forEach(t => t.classList.remove('active'));
                tab.classList.add('active');

                const tabName = tab.getAttribute('data-tab');
                const panel = tab.closest('.panel');
                panel.querySelectorAll('.tab-content').forEach(content => {
                    content.classList.remove('active');
                });
                panel.querySelector(`#${tabName}-tab`)?.classList.add('active');
            });
        });
    }

    initializeCharts() {
        const trainingCtx = document.getElementById('training-chart')?.getContext('2d');
        const privacyCtx = document.getElementById('privacy-chart')?.getContext('2d');
        const gradientCtx = document.getElementById('gradient-chart')?.getContext('2d');

        if (trainingCtx) {
            this.trainingChart = new Chart(trainingCtx, {
                type: 'line',
                data: {
                    labels: [],
                    datasets: [
                        {
                            label: 'Accuracy',
                            borderColor: '#4caf50',
                            backgroundColor: 'rgba(76, 175, 80, 0.1)',
                            data: [],
                            yAxisID: 'y',
                            borderWidth: 3,
                            pointRadius: 4,
                            pointHoverRadius: 6,
                            tension: 0.1
                        },
                        {
                            label: 'Loss',
                            borderColor: '#f44336',
                            backgroundColor: 'rgba(244, 67, 54, 0.1)',
                            data: [],
                            yAxisID: 'y1',
                            borderWidth: 3,
                            pointRadius: 4,
                            pointHoverRadius: 6,
                            tension: 0.1,
                            borderDash: [5, 5]  // Dashed line to differentiate from accuracy
                        }
                    ]
                },
                options: {
                    responsive: true,
                    maintainAspectRatio: false,
                    interaction: {
                        mode: 'index',
                        intersect: false,
                    },
                    plugins: {
                        legend: {
                            display: true,
                            position: 'top',
                            labels: {
                                usePointStyle: true,
                                padding: 20,
                                font: {
                                    size: 12,
                                    weight: 'bold'
                                }
                            }
                        },
                        tooltip: {
                            mode: 'index',
                            intersect: false,
                            backgroundColor: 'rgba(0, 0, 0, 0.8)',
                            titleColor: '#fff',
                            bodyColor: '#fff',
                            borderColor: '#ddd',
                            borderWidth: 1
                        }
                    },
                    scales: {
                        y: {
                            type: 'linear',
                            display: true,
                            position: 'left',
                            title: {
                                display: true,
                                text: 'Accuracy (%)',
                                color: '#4caf50',
                                font: {
                                    size: 14,
                                    weight: 'bold'
                                }
                            },
                            min: 0,
                            max: 100,
                            ticks: {
                                color: '#4caf50',
                                font: {
                                    weight: 'bold'
                                },
                                callback: function(value) {
                                    return value + '%';
                                }
                            },
                            grid: {
                                color: 'rgba(76, 175, 80, 0.2)'
                            }
                        },
                        y1: {
                            type: 'linear',
                            display: true,
                            position: 'right',
                            title: {
                                display: true,
                                text: 'Loss',
                                color: '#f44336',
                                font: {
                                    size: 14,
                                    weight: 'bold'
                                }
                            },
                            min: 0,
                            max: 3,  // More reasonable max for loss
                            ticks: {
                                color: '#f44336',
                                font: {
                                    weight: 'bold'
                                },
                                callback: function(value) {
                                    return value.toFixed(1);
                                }
                            },
                            grid: {
                                drawOnChartArea: false,  // Don't overlay grid lines
                                color: 'rgba(244, 67, 54, 0.2)'
                            },
                        },
                        x: {
                            title: {
                                display: true,
                                text: 'Training Progress',
                                font: {
                                    size: 12,
                                    weight: 'bold'
                                }
                            },
                            ticks: {
                                font: {
                                    size: 11
                                }
                            }
                        }
                    }
                }
            });
        }

        if (privacyCtx) {
            this.privacyChart = new Chart(privacyCtx, {
                type: 'line',
                data: {
                    labels: [],
                    datasets: [{
                        label: 'Privacy Budget (ε)',
                        borderColor: '#3f51b5',
                        data: []
                    }]
                },
                options: {
                    responsive: true,
                    maintainAspectRatio: false,
                    scales: {
                        y: {
                            beginAtZero: true,
                            title: {
                                display: true,
                                text: 'Privacy Budget (ε)'
                            }
                        }
                    }
                }
            });
        }

        if (gradientCtx) {
            this.gradientChart = new Chart(gradientCtx, {
                type: 'scatter',
                data: {
                    datasets: [
                        {
                            label: 'Before Clipping',
                            borderColor: '#2196f3',
                            backgroundColor: 'rgba(33, 150, 243, 0.1)',
                            data: [],
                            showLine: true
                        },
                        {
                            label: 'After Clipping',
                            borderColor: '#f44336',
                            backgroundColor: 'rgba(244, 67, 54, 0.1)',
                            data: [],
                            showLine: true
                        }
                    ]
                },
                options: {
                    responsive: true,
                    maintainAspectRatio: false,
                    scales: {
                        x: {
                            type: 'linear',
                            position: 'bottom',
                            title: {
                                display: true,
                                text: 'Gradient Norm'
                            },
                            min: 0
                        },
                        y: {
                            type: 'linear',
                            position: 'left',
                            title: {
                                display: true,
                                text: 'Density'
                            },
                            min: 0
                        }
                    },
                    plugins: {
                        annotation: {
                            annotations: {
                                line1: {
                                    type: 'line',
                                    xMin: 1,
                                    xMax: 1,
                                    borderColor: '#f44336',
                                    borderWidth: 2,
                                    borderDash: [5, 5],
                                    label: {
                                        content: 'Clipping Threshold',
                                        display: true,
                                        position: 'top'
                                    }
                                }
                            }
                        }
                    }
                }
            });
        }
    }

    async updatePrivacyBudget() {
        const params = this.getParameters();
        try {
            const response = await fetch('/api/privacy-budget', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(params)
            });
            const data = await response.json();
            
            // Update UI
            const budgetValue = document.getElementById('budget-value');
            const budgetFill = document.getElementById('budget-fill');
            
            if (budgetValue && budgetFill) {
                budgetValue.textContent = data.epsilon.toFixed(2);
                budgetFill.style.width = `${Math.min(data.epsilon / 10 * 100, 100)}%`;
                
                
            try { track('privacy_budget_update', { epsilon: data.epsilon }); } catch (e) {};
// Update class for coloring
                budgetFill.classList.remove('low', 'medium', 'high');
                if (data.epsilon <= 1) {
                    budgetFill.classList.add('low');
                } else if (data.epsilon <= 5) {
                    budgetFill.classList.add('medium');
                } else {
                    budgetFill.classList.add('high');
                }
            }
        } catch (error) {
            console.error('Error calculating privacy budget:', error);
        }
    }

    async toggleTraining() {
        if (this.isTraining) {
            this.stopTraining();
        } else {
            await this.startTraining();
        }
    }

    async startTraining() {
        const trainButton = document.getElementById('train-button');
        const trainingStatus = document.getElementById('training-status');
        const trainingStatusText = document.getElementById('training-status-text');
        const currentEpochEl = document.getElementById('current-epoch');
        const totalEpochsEl = document.getElementById('total-epochs');
        
        if (!trainButton || this.isTraining) return;

        this.isTraining = true;
        this.epochsData = [];  // Reset epoch data for streaming
        
        trainButton.textContent = 'Stop Training';
        trainButton.classList.add('running');
        trainingStatus.style.display = 'flex';
        
        // Show initialization status
        if (trainingStatusText) {
            trainingStatusText.textContent = 'Initializing model...';
            trainingStatusText.style.color = '#ff9800';  // Orange for initializing
        }
        if (currentEpochEl) currentEpochEl.textContent = '0';
        if (totalEpochsEl) totalEpochsEl.textContent = this.getParameters().epochs;

        // Reset charts
        this.resetCharts();

        console.log('Starting streaming training with parameters:', this.getParameters());

        // === Analytics: training started ===
        try {
            track('train_start', {
                ...this.getParameters(),
                view: this.currentView
            });
        } catch (e) {}

        // Use fetch with POST to initiate SSE stream (EventSource only supports GET)
        try {
            const response = await fetch('/api/train-stream', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify(this.getParameters())
            });

            if (!response.ok) {
                throw new Error('Failed to start training');
            }

            const reader = response.body.getReader();
            const decoder = new TextDecoder();
            let buffer = '';

            while (true) {
                const { done, value } = await reader.read();
                
                console.log('[Stream] Read chunk - done:', done, 'value size:', value?.length, 'isTraining:', this.isTraining);
                
                if (done || !this.isTraining) {
                    console.log('[Stream] Stream ended or training stopped');
                    break;
                }

                const chunk = decoder.decode(value, { stream: true });
                console.log('[Stream] Decoded chunk:', chunk.substring(0, 200));
                buffer += chunk;
                
                // Process complete SSE messages
                const lines = buffer.split('\n');
                buffer = lines.pop() || '';  // Keep incomplete line in buffer

                console.log('[Stream] Processing', lines.length, 'lines, buffer remaining:', buffer.length);

                for (const line of lines) {
                    if (line.startsWith('data: ')) {
                        try {
                            const data = JSON.parse(line.slice(6));
                            console.log('[Stream] Parsed SSE data type:', data.type);
                            this.handleStreamingData(data);
                        } catch (parseError) {
                            console.warn('[Stream] Failed to parse SSE data:', parseError, 'line:', line);
                        }
                    }
                }
            }

        } catch (error) {
            if (!this.isTraining) {
                console.log('Training was stopped');
                return;
            }
            
            console.error('Training error:', error);
            
            // === Analytics: training failed ===
            try {
                track('train_error', {
                    message: error.message || 'unknown',
                    params: this.getParameters()
                });
            } catch (e) {}
            
            // Show error message to user
            const errorMessage = document.createElement('div');
            errorMessage.className = 'error-message';
            errorMessage.textContent = error.message || 'An error occurred during training';
            const labMain = document.querySelector('.lab-main');
            if (labMain) {
                labMain.insertBefore(errorMessage, labMain.firstChild);
                setTimeout(() => errorMessage.remove(), 5000);
            }
        } finally {
            try {
                track('train_end', { ended_at: Date.now() });
            } catch (e) {}
            this.stopTraining();
        }
    }

    handleStreamingData(data) {
        const trainingStatusText = document.getElementById('training-status-text');
        const currentEpochEl = document.getElementById('current-epoch');
        const totalEpochsEl = document.getElementById('total-epochs');
        const chartInfo = document.getElementById('chart-info');

        console.log('[SSE] Received:', data.type, data);

        switch (data.type) {
            case 'status':
                // Update status message
                console.log('[SSE] Status update:', data.message);
                if (trainingStatusText) {
                    trainingStatusText.textContent = data.message;
                    trainingStatusText.style.color = data.message.includes('Initializing') ? '#ff9800' : '#4caf50';
                }
                if (currentEpochEl) currentEpochEl.textContent = data.epoch;
                if (totalEpochsEl) totalEpochsEl.textContent = data.total_epochs;
                break;

            case 'progress':
                // Update progress - add new epoch data to chart
                console.log('[SSE] Progress update - Epoch:', data.epoch, 'Accuracy:', data.epoch_data?.accuracy);
                if (trainingStatusText) {
                    trainingStatusText.textContent = `Training epoch ${data.epoch}...`;
                    trainingStatusText.style.color = '#4caf50';
                }
                if (currentEpochEl) currentEpochEl.textContent = data.epoch;
                if (totalEpochsEl) totalEpochsEl.textContent = data.total_epochs;
                
                // Add epoch data to our collection
                this.epochsData.push(data.epoch_data);
                
                // Update chart with new data point
                console.log('[SSE] Updating chart with epoch data, chart exists:', !!this.trainingChart);
                this.updateChartRealtime(data.epoch_data);
                
                if (chartInfo) {
                    chartInfo.textContent = `Showing ${this.epochsData.length} data points (epochs)`;
                }
                break;

            case 'complete':
                // Training complete - update all final results
                console.log('Training complete:', data);
                
                // Store complete data
                this.epochsData = data.epochs_data || this.epochsData;
                this.iterationsData = data.iterations_data || [];
                
                // Update final results
                this.updateResults(data);
                
                // === Analytics: training succeeded ===
                try {
                    track('train_success', {
                        trainer_type: data.trainer_type,
                        dataset: data.dataset,
                        model_architecture: data.model_architecture,
                        final_metrics: data.final_metrics,
                        privacy_budget: data.privacy_budget,
                        epochs: this.epochsData.length
                    });
                } catch (e) {}
                break;

            case 'error':
                console.error('Training error from server:', data.message);
                const errorMessage = document.createElement('div');
                errorMessage.className = 'error-message';
                errorMessage.textContent = data.message || 'An error occurred during training';
                const labMain = document.querySelector('.lab-main');
                if (labMain) {
                    labMain.insertBefore(errorMessage, labMain.firstChild);
                    setTimeout(() => errorMessage.remove(), 5000);
                }
                break;
        }
    }

    updateChartRealtime(epochData) {
        console.log('[Chart] updateChartRealtime called, chart exists:', !!this.trainingChart, 'epochData:', epochData);
        
        if (!this.trainingChart) {
            console.error('[Chart] Training chart not initialized!');
            return;
        }

        // Add new data point to chart
        const label = `Epoch ${epochData.epoch}`;
        
        this.trainingChart.data.labels.push(label);
        this.trainingChart.data.datasets[0].data.push(epochData.accuracy);
        this.trainingChart.data.datasets[1].data.push(epochData.loss);

        console.log('[Chart] Updated data - labels:', this.trainingChart.data.labels.length, 
                    'accuracies:', this.trainingChart.data.datasets[0].data,
                    'losses:', this.trainingChart.data.datasets[1].data);

        // Auto-adjust loss scale
        const losses = this.trainingChart.data.datasets[1].data;
        const maxLoss = Math.max(...losses);
        const minLoss = Math.min(...losses);
        this.trainingChart.options.scales.y1.max = Math.max(maxLoss * 1.1, 3);
        this.trainingChart.options.scales.y1.min = Math.max(0, minLoss * 0.9);

        // Update chart with animation
        this.trainingChart.update('none');  // 'none' for faster updates during streaming
        console.log('[Chart] Chart updated');
    }

    stopTraining() {
        // Mark as not training - this will cause the stream reader to stop
        this.isTraining = false;
        
        // Abort any pending training request
        if (this.abortController) {
            this.abortController.abort();
            this.abortController = null;
        }
        
        // Close any active event source
        if (this.eventSource) {
            this.eventSource.close();
            this.eventSource = null;
        }
        
        const trainButton = document.getElementById('train-button');
        if (trainButton) {
            trainButton.textContent = 'Run Training';
            trainButton.classList.remove('running');
        }
        const trainingStatus = document.getElementById('training-status');
        if (trainingStatus) {
            trainingStatus.style.display = 'none';
        }
    }

    resetCharts() {
        if (this.trainingChart) {
            this.trainingChart.data.labels = [];
            this.trainingChart.data.datasets[0].data = [];
            this.trainingChart.data.datasets[1].data = [];
            this.trainingChart.update();
        }

        if (this.privacyChart) {
            this.privacyChart.data.labels = [];
            this.privacyChart.data.datasets[0].data = [];
            this.privacyChart.update();
        }

        if (this.gradientChart) {
            this.gradientChart.data.datasets[0].data = [];
            this.gradientChart.data.datasets[1].data = [];
            this.gradientChart.update();
        }
    }

    switchView(view) {
        this.currentView = view;
        
        
        try { track('view_switch', { view }); } catch (e) {};
// Update button states
        document.querySelectorAll('.view-toggle').forEach(btn => {
            btn.classList.remove('active');
        });
        document.getElementById(`view-${view}`).classList.add('active');
        
        // Update chart with current data
        if (view === 'epochs' && this.epochsData.length > 0) {
            this.updateChartsWithData(this.epochsData, 'epochs');
        } else if (view === 'iterations' && this.iterationsData.length > 0) {
            this.updateChartsWithData(this.iterationsData, 'iterations');
        }
    }

    updateCharts(data) {
        if (!this.trainingChart || !data) return;

        console.log('Updating charts with data:', data); // Debug log

        // Store data for view switching
        if (data.epochs_data) {
            this.epochsData = data.epochs_data;
        }
        if (data.iterations_data) {
            this.iterationsData = data.iterations_data;
        }

        // Use current view to determine which data to display
        if (this.currentView === 'epochs' && this.epochsData.length > 0) {
            this.updateChartsWithData(this.epochsData, 'epochs');
        } else if (this.currentView === 'iterations' && this.iterationsData.length > 0) {
            this.updateChartsWithData(this.iterationsData, 'iterations');
        } else if (this.epochsData.length > 0) {
            // Fallback to epochs if iterations not available
            this.updateChartsWithData(this.epochsData, 'epochs');
        }
    }

    updateChartsWithData(chartData, dataType) {
        if (!this.trainingChart || !chartData) return;

        // Update training metrics chart
        const labels = chartData.map(d => 
            dataType === 'epochs' ? `Epoch ${d.epoch}` : `Iter ${d.iteration}`
        );
        const accuracies = chartData.map(d => d.accuracy);
        const losses = chartData.map(d => d.loss);

        console.log(`${dataType} - Accuracies:`, accuracies);
        console.log(`${dataType} - Losses:`, losses);

        this.trainingChart.data.labels = labels;
        this.trainingChart.data.datasets[0].data = accuracies;
        this.trainingChart.data.datasets[1].data = losses;

        // Auto-adjust loss scale based on actual data
        const maxLoss = Math.max(...losses);
        const minLoss = Math.min(...losses);
        this.trainingChart.options.scales.y1.max = Math.max(maxLoss * 1.1, 3);
        this.trainingChart.options.scales.y1.min = Math.max(0, minLoss * 0.9);

        // Update chart info
        const chartInfo = document.getElementById('chart-info');
        if (chartInfo) {
            chartInfo.textContent = `Showing ${chartData.length} data points (${dataType})`;
        }

        this.trainingChart.update();

        // Update current epoch display
        const currentEpoch = document.getElementById('current-epoch');
        const totalEpochs = document.getElementById('total-epochs');
        if (currentEpoch && totalEpochs && dataType === 'epochs') {
            currentEpoch.textContent = chartData.length;
            totalEpochs.textContent = this.getParameters().epochs;
        }

        // Update privacy budget chart (only for epochs view)
        if (this.privacyChart && dataType === 'epochs') {
            const privacyBudgets = chartData.map((_, i) => 
                this.calculateEpochPrivacy(i + 1)
            );
            this.privacyChart.data.labels = labels;
            this.privacyChart.data.datasets[0].data = privacyBudgets;
            this.privacyChart.update();
        }

        // Update gradient visualization
        if (this.gradientChart) {
            const clippingNorm = this.getParameters().clipping_norm;
            
            // Generate gradient data if not provided in chartData
            let gradientData;
            if (chartData[chartData.length - 1]?.gradient_info) {
                gradientData = chartData[chartData.length - 1].gradient_info;
            } else {
                // Generate synthetic gradient data
                const beforeClipping = [];
                const afterClipping = [];
                
                // Generate log-normal distributed gradients
                const mu = Math.log(clippingNorm) - 0.5;
                const sigma = 0.8;
                
                for (let i = 0; i < 100; i++) {
                    const u1 = Math.random();
                    const u2 = Math.random();
                    const z = Math.sqrt(-2.0 * Math.log(u1)) * Math.cos(2.0 * Math.PI * u2);
                    const norm = Math.exp(mu + sigma * z);
                    
                    const density = Math.exp(-(Math.pow(Math.log(norm) - mu, 2) / (2 * sigma * sigma))) / 
                                  (norm * sigma * Math.sqrt(2 * Math.PI));
                    const y = 0.2 + 0.8 * (density / 0.8) + 0.1 * (Math.random() - 0.5);
                    
                    beforeClipping.push({ x: norm, y: y });
                    afterClipping.push({ x: Math.min(norm, clippingNorm), y: y });
                }
                
                gradientData = {
                    before_clipping: beforeClipping.sort((a, b) => a.x - b.x),
                    after_clipping: afterClipping.sort((a, b) => a.x - b.x)
                };
            }
            
            // Update gradient chart
            this.gradientChart.data.datasets[0].data = gradientData.before_clipping;
            this.gradientChart.data.datasets[1].data = gradientData.after_clipping;
            
            // Update clipping threshold line
            this.gradientChart.options.plugins.annotation.annotations.line1 = {
                type: 'line',
                xMin: clippingNorm,
                xMax: clippingNorm,
                borderColor: '#f44336',
                borderWidth: 2,
                borderDash: [5, 5],
                label: {
                    content: `Clipping Threshold (C=${clippingNorm.toFixed(1)})`,
                    display: true,
                    position: 'top'
                }
            };
            
            // Update x-axis scale based on clipping norm
            this.gradientChart.options.scales.x.max = Math.max(clippingNorm * 2.5, 5);
            
            this.gradientChart.update('active');
        }
    }

    updateResults(data) {
        // Hide no-results message and show results content
        document.getElementById('no-results').style.display = 'none';
        document.getElementById('results-content').style.display = 'block';

        // Update metrics
        document.getElementById('accuracy-value').textContent = 
            data.final_metrics.accuracy.toFixed(1) + '%';
        document.getElementById('loss-value').textContent = 
            data.final_metrics.loss.toFixed(3);
        document.getElementById('training-time-value').textContent = 
            data.final_metrics.training_time.toFixed(1) + 's';

        // Update privacy budget display (make it dynamic)
        const privacyBudgetElement = document.getElementById('privacy-budget-value');
        if (privacyBudgetElement) {
            privacyBudgetElement.textContent = `ε=${data.privacy_budget.toFixed(1)}`;
        }

        // Update privacy-utility trade-off explanation dynamically
        const tradeoffElement = document.getElementById('tradeoff-explanation');
        if (tradeoffElement) {
            const accuracy = data.final_metrics.accuracy.toFixed(1);
            const epsilon = data.privacy_budget.toFixed(1);
            
            // Generate realistic trade-off assessment
            let tradeoffAssessment;
            if (data.final_metrics.accuracy >= 85) {
                tradeoffAssessment = "This is an excellent trade-off for most applications.";
            } else if (data.final_metrics.accuracy >= 75) {
                tradeoffAssessment = "This is a good trade-off for most applications.";
            } else if (data.final_metrics.accuracy >= 65) {
                tradeoffAssessment = "This trade-off may be acceptable for privacy-critical applications.";
            } else if (data.final_metrics.accuracy >= 50) {
                tradeoffAssessment = "Low utility - consider reducing noise or increasing clipping norm.";
            } else {
                tradeoffAssessment = "Very poor utility - privacy parameters need significant adjustment.";
            }
            
            tradeoffElement.textContent = 
                `This model achieved ${accuracy}% accuracy with a privacy budget of ε=${epsilon}. ${tradeoffAssessment}`;
        }

        // Update recommendations
        const recommendationList = document.querySelector('.recommendation-list');
        recommendationList.innerHTML = '';
        data.recommendations.forEach(rec => {
            const item = document.createElement('li');
            item.className = 'recommendation-item';
            item.innerHTML = `
                <span class="recommendation-icon">${rec.icon}</span>
                <span>${rec.text}</span>
            `;
            recommendationList.appendChild(item);
        });
    }

    getParameters() {
        return {
            clipping_norm: parseFloat(document.getElementById('clipping-norm').value),
            noise_multiplier: parseFloat(document.getElementById('noise-multiplier').value),
            batch_size: parseInt(document.getElementById('batch-size').value),
            learning_rate: parseFloat(document.getElementById('learning-rate').value),
            epochs: parseInt(document.getElementById('epochs').value),
            dataset: document.getElementById('dataset-select').value,
            model_architecture: document.getElementById('model-select').value
        };
    }

    applyPreset(values) {
        document.getElementById('clipping-norm').value = values.clippingNorm;
        document.getElementById('noise-multiplier').value = values.noiseMultiplier;
        document.getElementById('batch-size').value = values.batchSize;
        document.getElementById('learning-rate').value = values.learningRate;
        document.getElementById('epochs').value = values.epochs;

        // Update displayed values
        document.getElementById('clipping-norm-value').textContent = values.clippingNorm;
        document.getElementById('noise-multiplier-value').textContent = values.noiseMultiplier;
        document.getElementById('batch-size-value').textContent = values.batchSize;
        document.getElementById('learning-rate-value').textContent = values.learningRate;
        document.getElementById('epochs-value').textContent = values.epochs;

        this.updatePrivacyBudget();
    }

    calculateEpochPrivacy(epoch) {
        const params = this.getParameters();
        
        // Get dataset size based on selection
        let datasetSize;
        switch(params.dataset) {
            case 'cifar10':
                datasetSize = 50000; // CIFAR-10 training set size
                break;
            case 'fashion-mnist':
                datasetSize = 60000; // Fashion-MNIST training set size
                break;
            case 'mnist':
            default:
                datasetSize = 60000; // MNIST training set size
                break;
        }
        
        const samplingRate = params.batch_size / datasetSize;
        const steps = epoch * (1 / samplingRate);
        const delta = 1e-5;
        const c = Math.sqrt(2 * Math.log(1.25 / delta));
        return Math.min((c * samplingRate * Math.sqrt(steps)) / params.noise_multiplier, 10);
    }

    updateGradientVisualization(clippingNorm) {
        if (!this.gradientChart) return;

        // Generate random gradient norms following a log-normal distribution
        const numPoints = 100;
        const beforeClipping = [];
        const afterClipping = [];

        // Parameters for log-normal distribution
        const mu = Math.log(clippingNorm) - 0.5;
        const sigma = 0.8;

        // Generate gradient norms
        for (let i = 0; i < numPoints; i++) {
            // Generate log-normal distributed gradient norms
            const u1 = Math.random();
            const u2 = Math.random();
            const z = Math.sqrt(-2.0 * Math.log(u1)) * Math.cos(2.0 * Math.PI * u2);
            const norm = Math.exp(mu + sigma * z);

            // Calculate density using kernel density estimation
            const density = Math.exp(-(Math.pow(Math.log(norm) - mu, 2) / (2 * sigma * sigma))) / (norm * sigma * Math.sqrt(2 * Math.PI));
            
            // Normalize density and add some randomness
            const y = 0.2 + 0.8 * (density / 0.8) + 0.1 * (Math.random() - 0.5);
            
            beforeClipping.push({ x: norm, y: y });
            afterClipping.push({ x: Math.min(norm, clippingNorm), y: y });
        }

        // Sort points by x-value for smoother lines
        beforeClipping.sort((a, b) => a.x - b.x);
        afterClipping.sort((a, b) => a.x - b.x);

        // Update chart data
        this.gradientChart.data.datasets[0].data = beforeClipping;
        this.gradientChart.data.datasets[1].data = afterClipping;
        
        // Update clipping threshold line
        this.gradientChart.options.plugins.annotation.annotations.line1 = {
            type: 'line',
            xMin: clippingNorm,
            xMax: clippingNorm,
            borderColor: '#f44336',
            borderWidth: 2,
            borderDash: [5, 5],
            label: {
                content: `Clipping Threshold (C=${clippingNorm.toFixed(1)})`,
                display: true,
                position: 'top'
            }
        };

        // Update x-axis scale based on clipping norm
        this.gradientChart.options.scales.x.max = Math.max(clippingNorm * 2.5, 5);
        
        // Update the chart with animation
        this.gradientChart.update('active');
    }

    updateGradientVisualizationWithData(beforeClipping, afterClipping, clippingNorm) {
        if (!this.gradientChart) return;

        // Update chart data with real training data
        this.gradientChart.data.datasets[0].data = beforeClipping;
        this.gradientChart.data.datasets[1].data = afterClipping;
        
        // Update clipping threshold line
        this.gradientChart.options.plugins.annotation.annotations.line1 = {
            type: 'line',
            xMin: clippingNorm,
            xMax: clippingNorm,
            borderColor: '#f44336',
            borderWidth: 2,
            borderDash: [5, 5],
            label: {
                content: `Clipping Threshold (C=${clippingNorm.toFixed(1)})`,
                display: true,
                position: 'top'
            }
        };

        // Update x-axis scale based on clipping norm
        this.gradientChart.options.scales.x.max = Math.max(clippingNorm * 2.5, 5);
        
        // Update the chart with animation
        this.gradientChart.update('active');
    }
}

// Initialize the application when the DOM is loaded
document.addEventListener('DOMContentLoaded', () => {
    window.dpsgdExplorer = new DPSGDExplorer();
});

function setOptimalParameters() {
    // Research-validated optimal parameters for DP-SGD on MNIST
    // Based on Optax/TF Privacy: achieves ~96-97% accuracy with reasonable privacy
    document.getElementById('clipping-norm').value = '1.0';    // Standard clipping norm
    document.getElementById('noise-multiplier').value = '1.1'; // Moderate noise (ε≈3-5)
    document.getElementById('batch-size').value = '256';       // Large batches for stability
    document.getElementById('learning-rate').value = '0.15';   // Higher LR works well for DP-SGD
    document.getElementById('epochs').value = '30';            // Sufficient for convergence
    
    // Update displays
    updateClippingNormDisplay();
    updateNoiseMultiplierDisplay();
    updateBatchSizeDisplay();
    updateLearningRateDisplay();
    updateEpochsDisplay();
}