pump_animation.js

/**
 * Pump Cycle Animation — Canvas 2D renderer + playback controller.
 * Adapted from DR MURPHY v2 pump-renderer.js for Gradio gr.HTML embedding.
 *
 * This file is loaded by scenario_gui.py via demo.load(js=...) at page load.
 * The drawPump function and animation controller are defined on globalThis
 * so they can be invoked from gr.HTML onerror handlers.
 *
 * Layout (left to right):
 *   [TANK] --ICV--> [CHAMBER with PISTON] --DCV--> [DISCHARGE]
 *   Below: Energy balance bar chart (WORK / ICV FLOW / DCV FLOW)
 */
function() {
    var _anims = {};

    // -----------------------------------------------------------------------
    // Color helpers
    // -----------------------------------------------------------------------
    function parseHex(hex) {
        var h = hex.replace('#', '');
        return [
            parseInt(h.substring(0, 2), 16),
            parseInt(h.substring(2, 4), 16),
            parseInt(h.substring(4, 6), 16),
        ];
    }

    function toHex(r, g, b) {
        function clamp(v) { return Math.max(0, Math.min(255, Math.round(v))); }
        return '#' +
            clamp(r).toString(16).padStart(2, '0') +
            clamp(g).toString(16).padStart(2, '0') +
            clamp(b).toString(16).padStart(2, '0');
    }

    function lerpColor(c1, c2, t) {
        var a = parseHex(c1), b = parseHex(c2);
        return toHex(
            a[0] + (b[0] - a[0]) * t,
            a[1] + (b[1] - a[1]) * t,
            a[2] + (b[2] - a[2]) * t
        );
    }

    function tempToColor(Tc_K) {
        var t = Math.max(0, Math.min(1, (Tc_K - 20) / 180));
        if (t < 0.44) return lerpColor('#2196F3', '#FF9800', t / 0.44);
        return lerpColor('#FF9800', '#F44336', (t - 0.44) / 0.56);
    }

    function enthalpyToColor(hc) {
        var t = Math.max(-1, Math.min(1, hc / 500));
        if (t <= 0) return lerpColor('#2196F3', '#b0b0b0', t + 1);
        return lerpColor('#b0b0b0', '#F44336', t);
    }

    // -----------------------------------------------------------------------
    // Drawing helpers
    // -----------------------------------------------------------------------
    function roundedRect(ctx, x, y, w, h, r) {
        r = Math.min(r, w / 2, h / 2);
        ctx.beginPath();
        ctx.moveTo(x + r, y);
        ctx.lineTo(x + w - r, y);
        ctx.arcTo(x + w, y, x + w, y + r, r);
        ctx.lineTo(x + w, y + h - r);
        ctx.arcTo(x + w, y + h, x + w - r, y + h, r);
        ctx.lineTo(x + r, y + h);
        ctx.arcTo(x, y + h, x, y + h - r, r);
        ctx.lineTo(x, y + r);
        ctx.arcTo(x, y, x + r, y, r);
        ctx.closePath();
    }

    function drawArrow(ctx, x1, y1, x2, y2, thickness, color) {
        var headLen = Math.max(8, thickness * 3);
        var dx = x2 - x1, dy = y2 - y1;
        var angle = Math.atan2(dy, dx);
        ctx.save();
        ctx.strokeStyle = color;
        ctx.fillStyle = color;
        ctx.lineWidth = thickness;
        ctx.lineCap = 'round';
        ctx.beginPath();
        ctx.moveTo(x1, y1);
        ctx.lineTo(x2 - headLen * Math.cos(angle), y2 - headLen * Math.sin(angle));
        ctx.stroke();
        ctx.beginPath();
        ctx.moveTo(x2, y2);
        ctx.lineTo(x2 - headLen * Math.cos(angle - Math.PI / 6), y2 - headLen * Math.sin(angle - Math.PI / 6));
        ctx.lineTo(x2 - headLen * Math.cos(angle + Math.PI / 6), y2 - headLen * Math.sin(angle + Math.PI / 6));
        ctx.closePath();
        ctx.fill();
        ctx.restore();
    }

    // -----------------------------------------------------------------------
    // Main draw function
    // -----------------------------------------------------------------------
    function drawPump(ctx, width, height, state, pumpGeom) {
        var angle_deg = state.angle_deg;
        var pc = state.pc;
        var Tc_K = state.Tc_K;
        var yp = state.yp;
        var yp_prev = state.yp_prev;
        var ICV_open_frac = state.ICV_open_frac;
        var DCV_open_frac = state.DCV_open_frac;
        var ICV_leak_kgpm = state.ICV_leak_kgpm;
        var DCV_leak_kgpm = state.DCV_leak_kgpm;
        var hc = state.hc;

        // 0. Clear
        ctx.clearRect(0, 0, width, height);
        ctx.fillStyle = '#f0f0ec';
        ctx.fillRect(0, 0, width, height);

        // 1. Layout constants
        var margin = 24;
        var pumpAreaH = height * 0.70;
        var chamberW = width * 0.42;
        var chamberH = pumpAreaH * 0.55;
        var chamberX = (width - chamberW) / 2;
        var chamberY = (pumpAreaH - chamberH) / 2 + 4;

        var boxW = 80, boxH = chamberH * 0.5;
        var tankX = margin, tankY = chamberY + (chamberH - boxH) / 2;
        var dischX = width - margin - boxW, dischY = tankY;

        var valveW = 14, valveH = 30, valveMaxTravel = 18;
        var pistonH = 14;

        // 2. Pipes
        var pipeY = chamberY + chamberH / 2 - 4, pipeH = 8;
        ctx.fillStyle = '#c0c0b8';
        ctx.fillRect(tankX + boxW, pipeY, chamberX - (tankX + boxW), pipeH);
        ctx.fillRect(chamberX + chamberW, pipeY, dischX - (chamberX + chamberW), pipeH);

        // 3. Tank
        ctx.fillStyle = '#2196F3';
        roundedRect(ctx, tankX, tankY, boxW, boxH, 4); ctx.fill();
        ctx.strokeStyle = '#1565C0'; ctx.lineWidth = 2;
        roundedRect(ctx, tankX, tankY, boxW, boxH, 4); ctx.stroke();
        ctx.fillStyle = '#ffffff';
        ctx.font = 'bold 13px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
        ctx.fillText('TANK', tankX + boxW / 2, tankY + boxH / 2);

        // 4. Discharge
        ctx.fillStyle = '#e8e8e4';
        roundedRect(ctx, dischX, dischY, boxW, boxH, 4); ctx.fill();
        ctx.strokeStyle = '#808078'; ctx.lineWidth = 2;
        roundedRect(ctx, dischX, dischY, boxW, boxH, 4); ctx.stroke();
        ctx.fillStyle = '#505050';
        ctx.font = 'bold 11px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
        ctx.fillText('DISCHARGE', dischX + boxW / 2, dischY + boxH / 2);

        // 5. Chamber
        var pistonFrac = Math.max(0, Math.min(1, yp));
        var pistonY = chamberY + pistonFrac * (chamberH - pistonH);
        var fluidH = pistonY - chamberY;
        if (fluidH > 0) {
            ctx.fillStyle = tempToColor(Tc_K);
            ctx.fillRect(chamberX, chamberY, chamberW, fluidH);
        }
        var belowPistonY = pistonY + pistonH;
        var belowH = chamberY + chamberH - belowPistonY;
        if (belowH > 0) {
            ctx.fillStyle = '#d8d8d2';
            ctx.fillRect(chamberX, belowPistonY, chamberW, belowH);
        }
        ctx.fillStyle = '#505050';
        ctx.fillRect(chamberX, pistonY, chamberW, pistonH);
        var rodW = 8, rodX = chamberX + chamberW / 2 - rodW / 2;
        ctx.fillRect(rodX, pistonY + pistonH, rodW, chamberY + chamberH - (pistonY + pistonH));
        ctx.fillStyle = '#3a3a3a';
        ctx.fillRect(chamberX, pistonY + 2, chamberW, 2);
        ctx.fillRect(chamberX, pistonY + pistonH - 4, chamberW, 2);
        ctx.strokeStyle = '#404040'; ctx.lineWidth = 3;
        ctx.strokeRect(chamberX, chamberY, chamberW, chamberH);

        // 6. ICV (left)
        var icvOF = Math.max(0, Math.min(1, ICV_open_frac));
        var icvX = chamberX - valveW + icvOF * (-valveMaxTravel);
        var icvY = chamberY + chamberH / 2 - valveH / 2;
        ctx.fillStyle = '#606060';
        ctx.fillRect(chamberX - 4, icvY - 2, 4, valveH + 4);
        ctx.fillStyle = '#1a4fd6';
        ctx.fillRect(icvX, icvY, valveW, valveH);
        ctx.strokeStyle = '#0d36a0'; ctx.lineWidth = 1;
        ctx.strokeRect(icvX, icvY, valveW, valveH);
        ctx.fillStyle = '#1a4fd6';
        ctx.font = 'bold 10px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'bottom';
        ctx.fillText('ICV', chamberX - valveW / 2 - valveMaxTravel / 2, icvY - 6);

        // 7. DCV (right)
        var dcvOF = Math.max(0, Math.min(1, DCV_open_frac));
        var dcvX = chamberX + chamberW + dcvOF * valveMaxTravel;
        var dcvY = chamberY + chamberH / 2 - valveH / 2;
        ctx.fillStyle = '#606060';
        ctx.fillRect(chamberX + chamberW, dcvY - 2, 4, valveH + 4);
        ctx.fillStyle = '#1a7a4a';
        ctx.fillRect(dcvX, dcvY, valveW, valveH);
        ctx.strokeStyle = '#0e5230'; ctx.lineWidth = 1;
        ctx.strokeRect(dcvX, dcvY, valveW, valveH);
        ctx.fillStyle = '#1a7a4a';
        ctx.font = 'bold 10px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'bottom';
        ctx.fillText('DCV', chamberX + chamberW + valveW / 2 + valveMaxTravel / 2, dcvY - 6);

        // 8. Flow arrows
        var maxFlow = 5.0;
        var flowColor = enthalpyToColor(hc || 0);
        if (ICV_leak_kgpm > 0.01) {
            var th = Math.max(1, Math.min(5, (ICV_leak_kgpm / maxFlow) * 5));
            drawArrow(ctx, tankX + boxW + 6, pipeY + pipeH / 2, chamberX - 6, pipeY + pipeH / 2, th, flowColor);
        }
        if (DCV_leak_kgpm > 0.01) {
            var th2 = Math.max(1, Math.min(5, (DCV_leak_kgpm / maxFlow) * 5));
            drawArrow(ctx, chamberX + chamberW + 6, pipeY + pipeH / 2, dischX - 6, pipeY + pipeH / 2, th2, flowColor);
        }

        // 9. Pressure text
        var textY = Math.max(chamberY + 18, pistonY - 16);
        ctx.font = 'bold 16px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
        ctx.fillStyle = 'rgba(0,0,0,0.45)';
        ctx.fillText(pc.toFixed(0) + ' bar', chamberX + chamberW / 2 + 1, textY + 1);
        ctx.fillStyle = '#ffffff';
        ctx.fillText(pc.toFixed(0) + ' bar', chamberX + chamberW / 2, textY);

        // 10. Angle indicator (top-right)
        ctx.fillStyle = '#1a1a1a';
        ctx.font = '14px "JetBrains Mono", monospace';
        ctx.textAlign = 'right'; ctx.textBaseline = 'top';
        ctx.fillText(angle_deg.toFixed(1) + '\u00B0', width - margin, margin / 2);

        // 11. Temperature badge (top-left)
        ctx.fillStyle = tempToColor(Tc_K);
        ctx.font = '13px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'top';
        ctx.fillText(Tc_K.toFixed(1) + ' K', margin, margin / 2);

        // 12. Geometry annotation
        if (pumpGeom) {
            ctx.fillStyle = '#888880';
            ctx.font = '10px "JetBrains Mono", monospace';
            ctx.textAlign = 'right'; ctx.textBaseline = 'bottom';
            ctx.fillText('\u2300' + pumpGeom.bore_mm.toFixed(1) + ' x ' + pumpGeom.stroke_mm.toFixed(0) + ' mm', width - margin, pumpAreaH - 4);
        }

        // 13. Stroke phase indicator (top-center)
        var isRetract = (angle_deg % 360) < 180;
        var phaseText = isRetract ? 'RETRACT (intake)' : 'EXTEND (compress/discharge)';
        var phaseColor = isRetract ? '#1a4fd6' : '#c01a1a';
        ctx.save();
        ctx.font = 'bold 13px "JetBrains Mono", monospace';
        ctx.textAlign = 'center'; ctx.textBaseline = 'top';
        ctx.fillStyle = phaseColor;
        ctx.fillText(phaseText, width / 2, margin / 2);
        ctx.restore();

        // 14. Valve force indicators
        var forceArrowLen = 22, springLen = 12;
        // ICV
        var fCX_icv = chamberX - valveW / 2 - valveMaxTravel / 2;
        var fY_icv = icvY - 20;
        drawArrow(ctx, fCX_icv - 2, fY_icv, fCX_icv - 2 + springLen, fY_icv, 2, '#b85c00');
        if (ICV_leak_kgpm > 0.005) {
            var pLen = Math.min(forceArrowLen, Math.max(3, (ICV_leak_kgpm / maxFlow) * forceArrowLen));
            drawArrow(ctx, fCX_icv + 2, fY_icv, fCX_icv + 2 - pLen, fY_icv, 2, '#1a4fd6');
        }
        // DCV
        var fCX_dcv = chamberX + chamberW + valveW / 2 + valveMaxTravel / 2;
        var fY_dcv = dcvY - 20;
        drawArrow(ctx, fCX_dcv + 2, fY_dcv, fCX_dcv + 2 - springLen, fY_dcv, 2, '#b85c00');
        if (DCV_leak_kgpm > 0.005) {
            var pLen2 = Math.min(forceArrowLen, Math.max(3, (DCV_leak_kgpm / maxFlow) * forceArrowLen));
            drawArrow(ctx, fCX_dcv - 2, fY_dcv, fCX_dcv - 2 + pLen2, fY_dcv, 2, '#1a7a4a');
        }

        // 15. Energy balance bar chart
        var barAreaTop = pumpAreaH + 8;
        var barAreaLeft = margin + 10;
        var barAreaWidth = width - 2 * margin - 20;
        var barHeight = 22, barGap = 10;
        var barLabelW = 90, barValueW = 80;
        var barMaxW = barAreaWidth - barLabelW - barValueW;

        // Separator
        ctx.save();
        ctx.strokeStyle = '#a0a09a'; ctx.lineWidth = 1;
        ctx.beginPath();
        ctx.moveTo(margin, pumpAreaH + 2);
        ctx.lineTo(width - margin, pumpAreaH + 2);
        ctx.stroke();
        ctx.restore();

        ctx.save();
        ctx.fillStyle = '#505050';
        ctx.font = 'bold 10px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'top';
        ctx.fillText('ENERGY BALANCE', barAreaLeft, barAreaTop);
        ctx.restore();

        var barsStartY = barAreaTop + 18;
        var deltaYp = yp - (yp_prev !== undefined ? yp_prev : yp);
        var workProxy = Math.abs(deltaYp * pc);
        var isCompression = deltaYp < 0;
        var workScale = 100, flowScale = maxFlow;

        // Bar 1: WORK
        var y1 = barsStartY;
        var barW1 = Math.min(barMaxW, Math.max(2, (workProxy / workScale) * barMaxW));
        var barColor1 = isCompression ? '#c01a1a' : '#1a7a4a';
        ctx.save();
        ctx.fillStyle = '#505050'; ctx.font = '11px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'middle';
        ctx.fillText('WORK', barAreaLeft, y1 + barHeight / 2);
        ctx.fillStyle = '#d8d8d4';
        ctx.fillRect(barAreaLeft + barLabelW, y1, barMaxW, barHeight);
        ctx.fillStyle = barColor1;
        ctx.fillRect(barAreaLeft + barLabelW, y1, barW1, barHeight);
        ctx.strokeStyle = '#a0a09a'; ctx.lineWidth = 1;
        ctx.strokeRect(barAreaLeft + barLabelW, y1, barMaxW, barHeight);
        ctx.fillStyle = '#1a1a1a'; ctx.font = '10px "JetBrains Mono", monospace';
        ctx.fillText(workProxy.toFixed(1) + (isCompression ? ' COMP' : ' EXP'), barAreaLeft + barLabelW + barMaxW + 6, y1 + barHeight / 2);
        ctx.restore();

        // Bar 2: ICV FLOW
        var y2 = barsStartY + barHeight + barGap;
        var barW2 = Math.min(barMaxW, Math.max(0, (ICV_leak_kgpm / flowScale) * barMaxW));
        ctx.save();
        ctx.fillStyle = '#505050'; ctx.font = '11px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'middle';
        ctx.fillText('ICV FLOW', barAreaLeft, y2 + barHeight / 2);
        ctx.fillStyle = '#d8d8d4';
        ctx.fillRect(barAreaLeft + barLabelW, y2, barMaxW, barHeight);
        ctx.fillStyle = '#1a4fd6';
        ctx.fillRect(barAreaLeft + barLabelW, y2, barW2, barHeight);
        ctx.strokeStyle = '#a0a09a'; ctx.lineWidth = 1;
        ctx.strokeRect(barAreaLeft + barLabelW, y2, barMaxW, barHeight);
        ctx.fillStyle = '#1a1a1a'; ctx.font = '10px "JetBrains Mono", monospace';
        ctx.fillText(ICV_leak_kgpm.toFixed(2) + ' kg/min', barAreaLeft + barLabelW + barMaxW + 6, y2 + barHeight / 2);
        ctx.restore();

        // Bar 3: DCV FLOW
        var y3 = barsStartY + 2 * (barHeight + barGap);
        var barW3 = Math.min(barMaxW, Math.max(0, (DCV_leak_kgpm / flowScale) * barMaxW));
        ctx.save();
        ctx.fillStyle = '#505050'; ctx.font = '11px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'middle';
        ctx.fillText('DCV FLOW', barAreaLeft, y3 + barHeight / 2);
        ctx.fillStyle = '#d8d8d4';
        ctx.fillRect(barAreaLeft + barLabelW, y3, barMaxW, barHeight);
        ctx.fillStyle = '#1a7a4a';
        ctx.fillRect(barAreaLeft + barLabelW, y3, barW3, barHeight);
        ctx.strokeStyle = '#a0a09a'; ctx.lineWidth = 1;
        ctx.strokeRect(barAreaLeft + barLabelW, y3, barMaxW, barHeight);
        ctx.fillStyle = '#1a1a1a'; ctx.font = '10px "JetBrains Mono", monospace';
        ctx.fillText(DCV_leak_kgpm.toFixed(2) + ' kg/min', barAreaLeft + barLabelW + barMaxW + 6, y3 + barHeight / 2);
        ctx.restore();

        // 16. Enthalpy legend
        ctx.save();
        ctx.fillStyle = '#888880';
        ctx.font = '10px "JetBrains Mono", monospace';
        ctx.textAlign = 'left'; ctx.textBaseline = 'top';
        ctx.fillText('hc: ' + (hc || 0).toFixed(1) + ' kJ/kg', margin + 90, margin / 2 + 2);
        ctx.restore();
    }

    // -----------------------------------------------------------------------
    // Animation controller
    // -----------------------------------------------------------------------

    function _drawFrame(st, frameIdx) {
        var f = st.frames;
        var i = Math.min(frameIdx, st.totalFrames - 1);
        var prevI = Math.max(0, i - 1);
        drawPump(st.ctx, st.canvas.width, st.canvas.height, {
            angle_deg:      f.angle_deg[i] || 0,
            pc:             f.pc[i] || 0,
            Tc_K:           f.Tc_K[i] || 20,
            yp:             f.yp[i] || 0,
            yp_prev:        f.yp[prevI] || 0,
            ICV_open_frac:  f.ICV_open_frac[i] || 0,
            DCV_open_frac:  f.DCV_open_frac[i] || 0,
            ICV_leak_kgpm:  f.ICV_leak_kgpm[i] || 0,
            DCV_leak_kgpm:  f.DCV_leak_kgpm[i] || 0,
            hc:             (f.hc && f.hc[i]) || 0,
        }, st.pumpGeom);

        // Update angle readout
        var angleEl = document.getElementById('pump-angle-' + st.uid);
        if (angleEl) angleEl.textContent = (f.angle_deg[i] || 0).toFixed(1) + '\u00B0';

        // Update scrubber position
        var scrubEl = document.getElementById('pump-scrub-' + st.uid);
        if (scrubEl) scrubEl.value = i;
    }

    function _startPlayback(st) {
        var framesPerTick = Math.max(1, Math.round(st.totalFrames / 360 * st.speed));
        function tick() {
            if (!st.isPlaying) return;
            // Check canvas still in DOM (Gradio may have replaced it)
            if (!document.getElementById('pump-anim-' + st.uid)) {
                st.isPlaying = false;
                return;
            }
            st.currentFrame = (st.currentFrame + framesPerTick) % st.totalFrames;
            _drawFrame(st, st.currentFrame);
            st.animId = requestAnimationFrame(tick);
        }
        st.animId = requestAnimationFrame(tick);
    }

    globalThis.initPumpAnimation = function(uid) {
        var container = document.getElementById('pump-anim-' + uid);
        if (!container) return;

        // Cancel any previous animation
        if (_anims._current && _anims._current.animId) {
            cancelAnimationFrame(_anims._current.animId);
            _anims._current.isPlaying = false;
        }

        var canvas = container.querySelector('canvas');
        var dataDiv = document.getElementById('pump-data-' + uid);
        if (!canvas || !dataDiv) return;

        var data = JSON.parse(dataDiv.textContent);
        var ctx = canvas.getContext('2d');

        var st = {
            uid: uid,
            ctx: ctx,
            canvas: canvas,
            frames: data.frames,
            pumpGeom: data.pumpGeom,
            totalFrames: data.frames.angle_deg.length,
            currentFrame: 0,
            isPlaying: false,
            speed: 1,
            animId: null,
        };

        _anims[uid] = st;
        _anims._current = st;

        // Draw first frame
        _drawFrame(st, 0);
    };

    globalThis.togglePumpAnim = function(uid) {
        var st = _anims[uid];
        if (!st) return;
        st.isPlaying = !st.isPlaying;

        var btn = document.getElementById('pump-play-' + uid);
        if (btn) btn.textContent = st.isPlaying ? '\u23F8' : '\u25B6';

        if (st.isPlaying) {
            _startPlayback(st);
        } else {
            if (st.animId) cancelAnimationFrame(st.animId);
            st.animId = null;
        }
    };

    globalThis.setPumpAnimSpeed = function(uid, speed) {
        var st = _anims[uid];
        if (st) st.speed = parseFloat(speed);
    };

    globalThis.scrubPumpAnim = function(uid, frame) {
        var st = _anims[uid];
        if (!st) return;
        st.currentFrame = parseInt(frame);
        // Pause on manual scrub
        if (st.isPlaying) {
            st.isPlaying = false;
            if (st.animId) cancelAnimationFrame(st.animId);
            st.animId = null;
            var btn = document.getElementById('pump-play-' + st.uid);
            if (btn) btn.textContent = '\u25B6';
        }
        _drawFrame(st, st.currentFrame);
    };
}