/**
 * GradCAM — Gradient Saliency Map for MobileNet transfer learning heads.
 *
 * Ported verbatim from reference/mould-detect-hybrid-binary-and-multi-classification-app/js/grad-cam.js
 * Wrapped as ES5 IIFE for Babel 6 compatibility.
 *
 * Algorithm: Gradient Saliency (Simonyan et al., 2013)
 *   1. Build differentiable path: image → MobileNet.infer(img, true) → classifier → class score
 *   2. tf.grad(scoreFn)(imageTensor) → [224,224,3] gradients
 *   3. L2 norm across colour channels → [224,224] saliency map
 *   4. 5×5 box smooth via tf.depthwiseConv2d
 *   5. Normalise to [0,1]
 *
 * Public API:
 *   GradCAM.compute(imageTensor, classifierModel, mobilenetModel, classIndex)
 *     → Promise<{ heatmap: Float32Array, width: number, height: number }>
 *   GradCAM.overlay(heatmap, width, height, imageData, alpha) → ImageData
 *   GradCAM.colourmap(value) → [r, g, b]
 *
 * Depends on: TF.js (tf global)
 * No DOM. No module-specific state. All functions are pure given inputs.
 */
var GradCAM = (function () {

    var IMG_SIZE = 224;

    function compute(imageTensor, classifierModel, mobilenetModel, classIndex) {
        var imgSqueezed = tf.tidy(function () { return imageTensor.squeeze(); });

        var scoreFn = function (img) {
            return tf.tidy(function () {
                var batched = img.expandDims(0);
                var preprocessed = tf.sub(tf.mul(batched, 2.0), 1.0);
                var features = mobilenetModel.infer(preprocessed, true);
                var logits = classifierModel.predict(features);
                return logits.slice([0, classIndex], [1, 1]).squeeze();
            });
        };

        var gradFn = tf.grad(scoreFn);
        var gradients = gradFn(imgSqueezed);

        var saliency = tf.tidy(function () {
            var sq = tf.square(gradients);
            var sumSq = tf.sum(sq, 2);
            return tf.sqrt(sumSq);
        });

        var smoothed = _boxSmooth(saliency, 5);

        var normalised = tf.tidy(function () {
            var mn = tf.min(smoothed);
            var mx = tf.max(smoothed);
            var range = tf.sub(mx, mn);
            return tf.where(
                tf.greater(range, 1e-8),
                tf.div(tf.sub(smoothed, mn), range),
                tf.zerosLike(smoothed)
            );
        });

        // Helper: dispose all intermediate tensors regardless of success/failure
        function _disposeAll() {
            try { imgSqueezed.dispose(); } catch (e) {}
            try { gradients.dispose();   } catch (e) {}
            try { saliency.dispose();    } catch (e) {}
            try { smoothed.dispose();    } catch (e) {}
            try { normalised.dispose();  } catch (e) {}
        }

        return normalised.data().then(function (heatmapData) {
            _disposeAll();
            return { heatmap: heatmapData, width: IMG_SIZE, height: IMG_SIZE };
        })['catch'](function (err) {
            console.error('[GradCAM] compute() failed, releasing tensors:', err.message);
            _disposeAll();
            throw err;
        });
    }

    function _boxSmooth(t, k) {
        // CR-01 FIX: wrap in tf.tidy so kernel, t4, and conv are all disposed
        // automatically. squeeze([0,3]) extracts the 2D result; tf.tidy disposes
        // every other tensor created inside this scope.
        return tf.tidy(function () {
            var t4     = t.reshape([1, t.shape[0], t.shape[1], 1]);
            var kernel = tf.fill([k, k, 1, 1], 1 / (k * k));
            var conv   = tf.depthwiseConv2d(t4, kernel, 1, 'same');
            return conv.squeeze([0, 3]);
        });
    }

    function overlay(heatmap, width, height, imageData, alpha) {
        alpha = (alpha === undefined) ? 0.5 : alpha;
        var out = new Uint8ClampedArray(width * height * 4);
        var src = imageData.data;
        for (var i = 0; i < width * height; i++) {
            var heat = heatmap[i];
            var rgb = colourmap(heat);
            var hr = rgb[0]; var hg = rgb[1]; var hb = rgb[2];
            var sr = src[i * 4];
            var sg = src[i * 4 + 1];
            var sb = src[i * 4 + 2];
            out[i * 4]     = Math.round(sr * (1 - alpha) + hr * alpha);
            out[i * 4 + 1] = Math.round(sg * (1 - alpha) + hg * alpha);
            out[i * 4 + 2] = Math.round(sb * (1 - alpha) + hb * alpha);
            out[i * 4 + 3] = 255;
        }
        return new ImageData(out, width, height);
    }

    function colourmap(t) {
        var r = Math.min(1, Math.max(0, 1.5 - Math.abs(4 * t - 3)));
        var g = Math.min(1, Math.max(0, 1.5 - Math.abs(4 * t - 2)));
        var b = Math.min(1, Math.max(0, 1.5 - Math.abs(4 * t - 1)));
        return [Math.round(r * 255), Math.round(g * 255), Math.round(b * 255)];
    }

    return { compute: compute, overlay: overlay, colourmap: colourmap };
})();

window.GradCAM = GradCAM;