// This is the core JavaScript for the Spectrogram Editor.
// It handles file loading (.spec), visualization, tool interaction, and saving.

// --- Global Variables ---
let currentSpecData = null; // Stores the currently displayed/edited spectrogram data
let originalSpecData = null; // Stores the pristine, initially loaded spectrogram data
let dctSize = 512; // Default DCT size, read from header

let undoStack = [];
let redoStack = [];
const MAX_HISTORY_SIZE = 50;

let activeTool = null; // 'line', 'ellipse', 'noise', etc.
let isDrawing = false;
let startX, startY; // For tracking mouse down position

let shapes = []; // Array to store all drawn shapes (lines, ellipses, etc.)

// Web Audio Context
const audioContext = new (window.AudioContext || window.webkitAudioContext)();

// Audio Constants (should match C++ side)
const SAMPLE_RATE = 32000; 
const MAX_FREQ = SAMPLE_RATE / 2; // Nyquist frequency
const MIN_FREQ = 20; // Lower bound for log scale visualization

const SDF_FALLOFF_FACTOR = 10.0; // Adjust this value to control the softness of SDF edges.

// --- Utility Functions for Audio Processing ---
// JavaScript equivalent of C++ idct_512
function javascript_idct_512(input) {
    const output = new Float32Array(dctSize);
    const PI = Math.PI;
    const N = dctSize;

    for (let n = 0; n < N; ++n) {
        let sum = input[0] / 2.0;
        for (let k = 1; k < N; ++k) {
            sum += input[k] * Math.cos((PI / N) * k * (n + 0.5));
        }
        output[n] = sum * (2.0 / N);
    }
    return output;
}

// Hanning window for smooth audio transitions (JavaScript equivalent)
function hanningWindow(size) {
    const window = new Float32Array(size);
    const PI = Math.PI;
    for (let i = 0; i < size; i++) {
        window[i] = 0.5 * (1 - Math.cos((2 * PI * i) / (size - 1)));
    }
    return window;
}

const hanningWindowArray = hanningWindow(dctSize); // Pre-calculate window

// --- Signed Distance Functions (SDFs) ---
// Generic 2D vector operations
function vec2(x, y) { return { x: x, y: y }; }
function length(v) { return Math.sqrt(v.x * v.x + v.y * v.y); }
function dot(v1, v2) { return v1.x * v2.x + v1.y * v2.y; }
function sub(v1, v2) { return vec2(v1.x - v2.x, v1.y - v2.y); }
function mul(v, s) { return vec2(v.x * s, v.y * s); }
function div(v, s) { return vec2(v.x / s, v.y / s); }
function normalize(v) { return div(v, length(v)); }
function clamp(x, minVal, maxVal) { return Math.max(minVal, Math.min(x, maxVal)); }
function abs(v) { return vec2(Math.abs(v.x), Math.abs(v.y)); }
function max(v1, v2) { return vec2(Math.max(v1.x, v2.x), Math.max(v1.y, v2.y)); }
function sign(x) { return (x > 0) ? 1 : ((x < 0) ? -1 : 0); }

// sdSegment(p, a, b) - signed distance to a line segment
// p: point, a: segment start, b: segment end
function sdSegment(p, a, b) {
    const pa = sub(p, a);
    const ba = sub(b, a);
    const h = clamp(dot(pa, ba) / dot(ba, ba), 0.0, 1.0);
    return length(sub(pa, mul(ba, h)));
}

// sdEllipse(p, r) - signed distance to an ellipse (p relative to center, r is half-extents)
// p: point relative to ellipse center, r: half-extents (rx, ry)
function sdEllipse(p, r) {
    const k0 = vec2(1, length(div(p, r)));
    const k1 = vec2(length(div(p, r)), 1);
    const f = ((dot(div(mul(p, p), k0), vec2(1, 1)) < dot(div(mul(p, p), k1), vec2(1, 1))) ? k0 : k1);
    return length(sub(p, mul(r, normalize(mul(f, p))))) * sign(length(p) - r.x); // Simplified, original has length(p)-r.x which is only for circular
}

// sdBox(p, r) - signed distance to a rectangle (p relative to center, r is half-extents)
// p: point relative to box center, r: half-extents (hx, hy)
function sdBox(p, r) {
    const q = sub(abs(p), r);
    return length(max(q, vec2(0, 0))) + Math.min(0.0, Math.max(q.x, q.y));
}

// --- Utility to map canvas coords to spectrogram bins/frames (LOG SCALE) ---
// Maps a linear frequency bin index to its corresponding frequency in Hz
function binIndexToFreq(binIndex) {
    return (binIndex / dctSize) * MAX_FREQ;
}

// Maps a frequency in Hz to its corresponding linear bin index
function freqToBinIndex(freq) {
    return Math.floor((freq / MAX_FREQ) * dctSize);
}

// Maps a frequency (Hz) to its corresponding log-scaled bin index
function freqToBinIndexLog(freq) {
    if (freq < MIN_FREQ) freq = MIN_FREQ; // Clamp minimum frequency
    const logMin = Math.log(MIN_FREQ);
    const logMax = Math.log(MAX_FREQ);
    const logFreq = Math.log(freq);
    const normalizedLog = (logFreq - logMin) / (logMax - logMin);
    return Math.floor(normalizedLog * dctSize);
}

// Maps a log-scaled bin index to its corresponding frequency in Hz
function binIndexToFreqLog(binIndex) {
    const normalizedLog = binIndex / dctSize;
    const logMin = Math.log(MIN_FREQ);
    const logMax = Math.log(MAX_FREQ);
    const logFreq = normalizedLog * (logMax - logMin) + logMin;
    return Math.exp(logFreq);
}

// Converts a frequency (Hz) to a Y-coordinate on the canvas (log scale)
function freqToCanvasYLog(freq, canvasHeight) {
    if (freq < MIN_FREQ) freq = MIN_FREQ; // Clamp minimum frequency
    const logMin = Math.log(MIN_FREQ);
    const logMax = Math.log(MAX_FREQ);
    const logFreq = Math.log(freq);
    const normalizedLog = (logFreq - logMin) / (logMax - logMin);
    return canvasHeight * (1 - normalizedLog); // Y-axis is inverted
}

// Converts a Y-coordinate on the canvas to a frequency (Hz) (log scale)
function canvasYToFreqLog(canvasY, canvasHeight) {
    const normalizedLog = 1 - (canvasY / canvasHeight);
    const logMin = Math.log(MIN_FREQ);
    const logMax = Math.log(MAX_FREQ);
    const logFreq = normalizedLog * (logMax - logMin) + logMin;
    return Math.exp(logFreq);
}

// Converts canvas Y-coordinate to log-scaled bin index
function canvasYToBinIndexLog(canvasY, specData) {
    const freq = canvasYToFreqLog(canvasY, canvas.height);
    return freqToBinIndex(freq); // Use linear bin index from calculated log freq
}

// Converts log-scaled bin index to canvas Y-coordinate
function binIndexToCanvasYLog(binIndex, specData) {
    const freq = binIndexToFreq(binIndex);
    return freqToCanvasYLog(freq, canvas.height);
}

// Helper to get frequency delta from canvas delta (for ellipse radius in freq)
function canvasDeltaYToFreqDeltaLog(canvasDeltaY, canvasHeight) {
    // This is an approximation as delta in log scale is not linear
    // For small deltas around a center, it can be approximated
    const centerCanvasY = canvasHeight / 2;
    const freqAtCenter = canvasYToFreqLog(centerCanvasY, canvasHeight);
    const freqAtCenterPlusDelta = canvasYToFreqLog(centerCanvasY - canvasDeltaY, canvasHeight);
    return Math.abs(freqAtCenterPlusDelta - freqAtCenter);
}

// Initial setup for canvas size (can be updated on window resize)
window.addEventListener('resize', () => {
    if (originalSpecData) {
        canvas.width = window.innerWidth * 0.7;
        canvas.height = 400; // Fixed height
        redrawCanvas();
    }
});

// Initial call to set button states
updateUndoRedoButtons();

// --- Utility for sizeof(float) in JS context ---
// This is a workaround since typeof(float) is not directly available.
// Float32Array.BYTES_PER_ELEMENT is used in handleFileSelect.
function sizeof(type) {
    if (type === 'float') {
        return Float32Array.BYTES_PER_ELEMENT;
    }
    return 0; 
}


// --- File Handling ---
const specFileInput = document.getElementById('specFileInput');
specFileInput.addEventListener('change', handleFileSelect);

async function handleFileSelect(event) {
    const file = event.target.files[0];
    if (!file) {
        return;
    }

    try {
        const buffer = await file.arrayBuffer();
        const dataView = new DataView(buffer);

        // Parse SPEC header
        const header = {
            magic: String.fromCharCode(...new Uint8Array(buffer.slice(0, 4))),
            version: dataView.getInt32(4, true),
            dct_size: dataView.getInt32(8, true),
            num_frames: dataView.getInt32(12, true)
        };

        if (header.magic !== "SPEC" || header.version !== 1) {
            console.error("Invalid SPEC file format.");
            alert("Invalid SPEC file format. Please load a valid .spec file.");
            return;
        }

        dctSize = header.dct_size;
        const dataStart = 16;
        const numBytes = header.num_frames * header.dct_size * Float32Array.BYTES_PER_ELEMENT;
        const spectralDataFloat = new Float32Array(buffer, dataStart, header.num_frames * header.dct_size);

        originalSpecData = { header: header, data: new Float32Array(spectralDataFloat) }; // Store pristine copy
        currentSpecData = { header: header, data: new Float32Array(spectralDataFloat) }; // Editable copy

        shapes = []; // Clear shapes on new file load
        undoStack = []; // Clear undo history
        redoStack = []; // Clear redo history

        console.log("Loaded SPEC file:", header);
        redrawCanvas(); // Redraw with new data

    } catch (error) {
        console.error("Error loading SPEC file:", error);
        alert("Failed to load SPEC file. Check console for details.");
    }
}

// --- Spectrogram Visualization ---
const canvas = document.getElementById('spectrogramCanvas');
const ctx = canvas.getContext('2d');

// Add canvas event listeners
canvas.addEventListener('mousedown', handleMouseDown);
canvas.addEventListener('mousemove', handleMouseMove);
canvas.addEventListener('mouseup', handleMouseUp);
canvas.addEventListener('mouseout', handleMouseUp); // Treat mouse out as mouse up

// Function to get a color based on intensity (0 to 1)
function getColorForIntensity(intensity) {
    // Example: Blue to white/yellow gradient
    const h = (1 - intensity) * 240; // Hue from blue (240) to red (0), inverse for intensity
    const s = 100; // Saturation
    const l = intensity * 50 + 50; // Lightness from 50 to 100
    return `hsl(${h}, ${s}%, ${l}%)`;
}

function drawSpectrogram(specData) {
    const width = canvas.width; 
    const height = canvas.height;

    ctx.clearRect(0, 0, width, height);
    ctx.fillStyle = '#ffffff';
    ctx.fillRect(0, 0, width, height);

    if (!specData || !specData.data || specData.header.num_frames === 0 || specData.data.length === 0) {
        console.warn("No spectrogram data or invalid header/data to draw.");
        return;
    }

    const numFrames = specData.header.num_frames;
    const frameWidth = width / numFrames; // Width of each time frame

    // Draw each frame's spectral data with log frequency scale
    for (let frameIndex = 0; frameIndex < numFrames; frameIndex++) {
        const frameDataStart = frameIndex * dctSize;
        const xPos = frameIndex * frameWidth;

        // To draw with log scale, we iterate over canvas y-coordinates
        // and map them back to frequency bins
        for (let y = 0; y < height; y++) {
            const binIndex = canvasYToBinIndexLog(y, specData);
            if (binIndex < 0 || binIndex >= dctSize) continue; // Out of bounds

            const value = specData.data[frameDataStart + binIndex];
            const intensity = Math.min(1, Math.abs(value) / 1.0); // Assuming values are normalized to [-1, 1]
            
            ctx.fillStyle = getColorForIntensity(intensity);
            ctx.fillRect(xPos, height - y - 1, frameWidth, 1); // Draw a 1-pixel height line for each y
        }
    }

    // Draw active shapes on top (previews for current drawing tool)
    shapes.forEach(shape => {
        drawShape(shape);
    });
}

function drawShape(shape) {
    // This draws the final, persistent shape. Preview is drawn in handleMouseMove.
    ctx.strokeStyle = shape.color || 'red';
    ctx.lineWidth = shape.width || 2;
    
    switch (shape.type) {
        case 'line':
            ctx.beginPath();
            ctx.moveTo(shape.x1, shape.y1);
            ctx.lineTo(shape.x2, shape.y2);
            ctx.stroke();
            break;
        case 'ellipse':
            ctx.beginPath();
            ctx.ellipse(shape.cx, shape.cy, shape.rx, shape.ry, 0, 0, 2 * Math.PI);
            ctx.stroke();
            break;
        case 'noise_rect': // Noise is visualized as a rectangle
            ctx.fillStyle = 'rgba(0, 0, 255, 0.2)';
            ctx.fillRect(shape.x, shape.y, shape.width, shape.height);
            ctx.strokeStyle = 'blue';
            ctx.strokeRect(shape.x, shape.y, shape.width, shape.height);
            break;
    }
}

// --- Mouse Event Handlers ---
function getMousePos(event) {
    const rect = canvas.getBoundingClientRect();
    return {
        x: event.clientX - rect.left,
        y: event.clientY - rect.top
    };
}

function handleMouseDown(event) {
    if (!activeTool || !currentSpecData) return;
    isDrawing = true;
    const pos = getMousePos(event);
    startX = pos.x;
    startY = pos.y;
}

function handleMouseMove(event) {
    if (!isDrawing || !activeTool) return;
    const pos = getMousePos(event);
    
    redrawCanvas(); // Clear and redraw persistent state

    ctx.strokeStyle = 'rgba(0, 0, 0, 0.5)'; // Preview color
    ctx.lineWidth = 1;
    ctx.setLineDash([5, 5]); // Dashed line for preview

    switch (activeTool) {
        case 'line':
            ctx.beginPath();
            ctx.moveTo(startX, startY);
            ctx.lineTo(pos.x, pos.y);
            ctx.stroke();
            break;
        case 'ellipse':
            // Draw preview ellipse based on start and current pos (bounding box)
            const rx = Math.abs(pos.x - startX) / 2;
            const ry = Math.abs(pos.y - startY) / 2;
            const cx = startX + (pos.x - startX) / 2;
            const cy = startY + (pos.y - startY) / 2;
            if (rx > 0 && ry > 0) {
                ctx.beginPath();
                ctx.ellipse(cx, cy, rx, ry, 0, 0, 2 * Math.PI);
                ctx.stroke();
            }
            break;
        case 'noise':
            // Draw preview rectangle for noise area
            const rectX = Math.min(startX, pos.x);
            const rectY = Math.min(startY, pos.y);
            const rectW = Math.abs(pos.x - startX);
            const rectH = Math.abs(pos.y - startY);
            ctx.strokeRect(rectX, rectY, rectW, rectH);
            break;
    }
    ctx.setLineDash([]); // Reset line dash
}

function handleMouseUp(event) {
    if (!isDrawing || !activeTool || !currentSpecData) return;
    isDrawing = false;
    const endPos = getMousePos(event);

    let newShape = null;

    switch (activeTool) {
        case 'line': {
            const startCoords = canvasToSpectrogramCoords(startX, startY, currentSpecData);
            const endCoords = canvasToSpectrogramCoords(endPos.x, endPos.y, currentSpecData);

            newShape = {
                type: 'line',
                // Canvas coordinates for drawing visual representation (unchanged)
                x1: startX, y1: startY,
                x2: endPos.x, y2: endPos.y,
                // World coordinates for SDF calculations (frame and log-scaled frequency)
                frame1_world: startCoords.frame,
                freq1_world: binIndexToFreqLog(startCoords.bin),
                frame2_world: endCoords.frame,
                freq2_world: binIndexToFreqLog(endCoords.bin),
                amplitude: 0.5, // Default amplitude
                width: 2,       // Visual width in canvas pixels, not directly used by SDF, but kept for drawing
                color: 'red',
                falloff: SDF_FALLOFF_FACTOR, // SDF falloff factor
            };
            break;
        }
        case 'ellipse': {
            const rx = Math.abs(endPos.x - startX) / 2;
            const ry = Math.abs(endPos.y - startY) / 2;
            const cx = startX + (endPos.x - startX) / 2;
            const cy = startY + (endPos.y - startY) / 2;

            const centerCoords = canvasToSpectrogramCoords(cx, cy, currentSpecData);
            const halfWidthFrames = Math.floor((rx / canvas.width) * currentSpecData.header.num_frames);

            const startFreq = canvasYToFreqLog(startY, canvas.height);
            const endFreq = canvasYToFreqLog(endPos.y, canvas.height);
            const centerFreq = (startFreq + endFreq) / 2;
            const halfHeightFreq = Math.abs(startFreq - endFreq) / 2;


            newShape = {
                type: 'ellipse',
                // Canvas coordinates for drawing visual representation (unchanged)
                cx: cx, cy: cy,
                rx: rx, ry: ry,
                // World coordinates for SDF calculations
                center_frame_world: centerCoords.frame,
                center_freq_world: centerFreq,
                radius_frames_world: halfWidthFrames,
                radius_freq_world: halfHeightFreq,
                amplitude: 0.5,
                color: 'green',
                falloff: SDF_FALLOFF_FACTOR,
            };
            break;
        }
        case 'noise': {
            const rectX = Math.min(startX, endPos.x);
            const rectY = Math.min(startY, endPos.y);
            const rectW = Math.abs(endPos.x - startX);
            const rectH = Math.abs(endPos.y - startY);

            const startCoords = canvasToSpectrogramCoords(rectX, rectY, currentSpecData);
            const endCoords = canvasToSpectrogramCoords(rectX + rectW, rectY + rectH, currentSpecData);

            const centerFrame = Math.floor((startCoords.frame + endCoords.frame) / 2);
            const centerFreq = (binIndexToFreqLog(startCoords.bin) + binIndexToFreqLog(endCoords.bin)) / 2;
            const halfExtentFrames = Math.floor(Math.abs(endCoords.frame - startCoords.frame) / 2);
            const halfExtentFreq = Math.abs(binIndexToFreqLog(endCoords.bin) - binIndexToFreqLog(startCoords.bin)) / 2;

            newShape = {
                type: 'noise_rect',
                // Canvas coordinates for drawing visual representation (unchanged)
                x: rectX, y: rectY,
                width: rectW, height: rectH,
                // World coordinates for SDF calculations
                center_frame_world: centerFrame,
                center_freq_world: centerFreq,
                half_extent_frames_world: halfExtentFrames,
                half_extent_freq_world: halfExtentFreq,
                amplitude: 0.3, // Default noise amplitude
                density: 0.5,   // Default noise density
                color: 'blue',
                falloff: 0.0, // No falloff for pure noise inside rect
            };
            break;
        }
    }

    if (newShape) {
        // Capture the state *before* applying the new shape for undo
        const previousDataSnapshot = new Float32Array(currentSpecData.data); // Copy of actual data
        const previousShapesSnapshot = shapes.map(s => ({ ...s })); // Deep copy shapes array

        applyShapeToSpectrogram(newShape, currentSpecData); // Modify currentSpecData directly
        shapes.push(newShape);
        addAction({
            type: 'add_shape',
            shape: newShape,
            undo: () => {
                // To undo, restore previous shapes and previous data
                shapes = previousShapesSnapshot;
                currentSpecData.data = previousDataSnapshot;
            },
            redo: () => {
                // To redo, add the shape back and apply it to current data
                shapes.push(newShape);
                applyShapeToSpectrogram(newShape, currentSpecData);
            }
        });
    }
    redrawCanvas(); // Final redraw after action
    updateUndoRedoButtons();
}

// --- Spectrogram Data Manipulation ---
function applyShapeToSpectrogram(shape, targetSpecData) {
    if (!targetSpecData || !targetSpecData.data || targetSpecData.header.num_frames === 0) return;

    const numFrames = targetSpecData.header.num_frames;

    // Determine a bounding box for optimization (iterate only relevant cells)
    let minFrame = 0, maxFrame = numFrames - 1;
    let minBin = 0, maxBin = dctSize - 1;

    // Calculate tighter bounding boxes for each shape type
    switch (shape.type) {
        case 'line':
            minFrame = Math.min(shape.frame1_world, shape.frame2_world) - Math.ceil(shape.width / 2);
            maxFrame = Math.max(shape.frame1_world, shape.frame2_world) + Math.ceil(shape.width / 2);
            // For frequency, approximate by visual width or a fixed range if needed
            minBin = freqToBinIndex(Math.min(shape.freq1_world, shape.freq2_world)) - Math.ceil(shape.width / 2);
            maxBin = freqToBinIndex(Math.max(shape.freq1_world, shape.freq2_world)) + Math.ceil(shape.width / 2);
            break;
        case 'ellipse':
            minFrame = shape.center_frame_world - shape.radius_frames_world - 1;
            maxFrame = shape.center_frame_world + shape.radius_frames_world + 1;
            minBin = freqToBinIndex(shape.center_freq_world - shape.radius_freq_world) - 1; // Approx bin range from world freq
            maxBin = freqToBinIndex(shape.center_freq_world + shape.radius_freq_world) + 1; // Approx bin range from world freq
            break;
        case 'noise_rect':
            minFrame = shape.center_frame_world - shape.half_extent_frames_world - 1;
            maxFrame = shape.center_frame_world + shape.half_extent_frames_world + 1;
            minBin = freqToBinIndex(shape.center_freq_world - shape.half_extent_freq_world) - 1; // Approx bin range from world freq
            maxBin = freqToBinIndex(shape.center_freq_world + shape.half_extent_freq_world) + 1; // Approx bin range from world freq
            break;
    }

    minFrame = Math.max(0, minFrame);
    maxFrame = Math.min(targetSpecData.header.num_frames - 1, maxFrame); // Use targetSpecData.header.num_frames
    minBin = Math.max(0, minBin);
    maxBin = Math.min(dctSize - 1, maxBin);

    for (let f = minFrame; f <= maxFrame; f++) {
        for (let b = minBin; b <= maxBin; b++) {
            const p_world = vec2(f, binIndexToFreqLog(b));
            let distance = Infinity;

            switch (shape.type) {
                case 'line':
                    const a_line = vec2(shape.frame1_world, shape.freq1_world);
                    const b_line = vec2(shape.frame2_world, shape.freq2_world);
                    distance = sdSegment(p_world, a_line, b_line);
                    break;
                case 'ellipse':
                    const center_ellipse = vec2(shape.center_frame_world, shape.center_freq_world);
                    const r_ellipse = vec2(shape.radius_frames_world, shape.radius_freq_world);
                    distance = sdEllipse(sub(p_world, center_ellipse), r_ellipse);
                    break;
                case 'noise_rect':
                    const center_box = vec2(shape.center_frame_world, shape.center_freq_world);
                    const r_box = vec2(shape.half_extent_frames_world, shape.half_extent_freq_world);
                    distance = sdBox(sub(p_world, center_box), r_box);
                    if (distance <= 0 && Math.random() < shape.density) { // Only add noise inside the box with density
                        targetSpecData.data[f * dctSize + b] += (Math.random() * 2 - 1) * shape.amplitude;
                    }
                    break;
            }

            if (shape.type !== 'noise_rect') { // Noise is handled differently for amplitude
                const attenuation = Math.exp(-distance * distance * shape.falloff);
                targetSpecData.data[f * dctSize + b] += shape.amplitude * attenuation;
            }

            // Clamp final value
            targetSpecData.data[f * dctSize + b] = Math.max(-1, Math.min(1, targetSpecData.data[f * dctSize + b]));
        }
    }
}

// --- Tool Interactions (Button Clicks) ---
const lineToolButton = document.getElementById('lineTool');
const ellipseToolButton = document.getElementById('ellipseTool');
const noiseToolButton = document.getElementById('noiseTool');
const undoButton = document.getElementById('undoButton');
const redoButton = document.getElementById('redoButton'); 
const listenOriginalButton = document.getElementById('listenOriginalButton');
const listenGeneratedButton = document.getElementById('listenGeneratedButton');

lineToolButton.addEventListener('click', () => { activeTool = 'line'; console.log('Line tool selected'); });
ellipseToolButton.addEventListener('click', () => { activeTool = 'ellipse'; console.log('Ellipse tool selected'); });
noiseToolButton.addEventListener('click', () => { activeTool = 'noise'; console.log('Noise tool selected'); });

// --- Undo/Redo Logic ---
function addAction(action) {
    undoStack.push(action);
    if (undoStack.length > MAX_HISTORY_SIZE) {
        undoStack.shift(); 
    }
    redoStack = []; 
    updateUndoRedoButtons();
}

function handleUndo() {
    if (undoStack.length === 0) {
        console.log('Undo stack is empty.');
        return;
    }
    const actionToUndo = undoStack.pop();
    actionToUndo.undo(); 
    redoStack.push(actionToUndo);
    redrawCanvas(); 
    updateUndoRedoButtons();
}

function handleRedo() {
    if (redoStack.length === 0) {
        console.log('Redo stack is empty.');
        return;
    }

    const actionToRedo = redoStack.pop();
    actionToRedo.redo(); 
    undoStack.push(actionToRedo);
    redrawCanvas(); 
    updateUndoRedoButtons();
}

function redrawCanvas() {
    console.log('Redrawing canvas...');
    if (!originalSpecData) {
        ctx.clearRect(0, 0, canvas.width, canvas.height);
        ctx.fillStyle = '#ffffff';
        ctx.fillRect(0, 0, canvas.width, canvas.height);
        return;
    }

    // Start with a fresh copy of the original data
    currentSpecData.data = new Float32Array(originalSpecData.data); 

    // Replay all shapes from the `shapes` array to `currentSpecData`
    shapes.forEach(shape => {
        applyShapeToSpectrogram(shape, currentSpecData); 
    });

    drawSpectrogram(currentSpecData); 
}

function updateUndoRedoButtons() {
    undoButton.disabled = undoStack.length === 0;
    redoButton.disabled = redoStack.length === 0; 
}