import { assert } from "../../../core/assert.js";
import { float_to_uint8 } from "../../../core/binary/float_to_uint8.js";
import { max2 } from "../../../core/math/max2.js";
import { CanvasView } from "../../../view/elements/CanvasView.js";
import EmptyView from "../../../view/elements/EmptyView.js";

/**
 * Debug view: renders each z-slice of a 3D field as a 2D canvas tile, with up to four
 * channels mapped to RGBA. Used by {@link prototype.js}; not a production runtime view.
 *
 * The visualiser is data-agnostic — it takes an array of one to four Float32Array
 * channels, each of length `res_x * res_y * res_z`, and a per-channel mapping to colour
 * via {@link setDrawMethod}. Default mapping clamps to `[0, 1]` per channel and writes
 * opaque pixels.
 */
export class SliceVisualiser extends EmptyView {

    /**
     * @type {Float32Array[]}
     */
    #channels = [];

    /**
     * @type {CanvasView[]}
     */
    #views = [];

    /**
     * @type {[number, number, number]}
     */
    #resolution = [0, 0, 0];

    /**
     * @type {ImageData|null}
     */
    #imageData = null;

    /**
     * Per-pixel transform. Receives the destination RGBA buffer, the byte offset of the
     * current pixel, a Float32Array of channel samples for that pixel, and the channel
     * count.
     * @type {(out: Uint8ClampedArray, out_offset: number, sample: Float32Array, channel_count: number) => void}
     */
    #drawMethod = (out, out_offset, sample, channel_count) => {
        for (let i = 0; i < channel_count; i++) {
            out[out_offset + i] = float_to_uint8(sample[i]);
        }
        if (channel_count < 4) {
            out[out_offset + 3] = 255;
        }
    };

    constructor() {
        super({ classList: ['slice-visualizer'] });

        this.on.linked.add(this.layout, this);
        this.size.onChanged.add(this.layout, this);
    }

    /**
     * Replace the per-pixel transform. The function receives the destination RGBA
     * `Uint8ClampedArray`, the byte offset of the current pixel within it, a Float32Array
     * of channel samples for that pixel, and the configured channel count.
     * @param {(out: Uint8ClampedArray, out_offset: number, sample: Float32Array, channel_count: number) => void} fn
     */
    setDrawMethod(fn) {
        this.#drawMethod = fn;
    }

    /**
     * Attach the channels and resolution to visualise. Recomputes the layout and
     * rebuilds the per-slice canvases (creating or destroying child views as needed).
     *
     * @param {Float32Array[]} channels   1-4 channels; each length = res_x * res_y * res_z.
     * @param {number[]} resolution       `[res_x, res_y, res_z]`
     */
    setChannels(channels, resolution) {
        assert.isArray(channels, "channels");
        assert.greaterThanOrEqual(channels.length, 1, "channels.length");
        assert.lessThanOrEqual(channels.length, 4, "channels.length");
        assert.equal(resolution.length, 3, "resolution must have 3 entries");

        const res_x = resolution[0];
        const res_y = resolution[1];
        const res_z = resolution[2];
        const cell_count = res_x * res_y * res_z;

        for (let i = 0; i < channels.length; i++) {
            assert.equal(channels[i].length, cell_count,
                `channel[${i}] sized to resolution`);
        }

        this.#channels = channels.slice();
        this.#resolution[0] = res_x;
        this.#resolution[1] = res_y;
        this.#resolution[2] = res_z;

        // Add canvases for new slices.
        while (this.#views.length < res_z) {
            const view = new CanvasView();
            view.size.set(res_x, res_y);
            view.css({ position: "absolute", left: 0, top: 0 });
            this.#views.push(view);
            this.addChild(view);
        }

        // Drop canvases for removed slices.
        while (this.#views.length > res_z) {
            const view = this.#views.pop();
            this.removeChild(view);
        }

        // Resize existing canvases if resolution changed.
        for (let i = 0; i < res_z; i++) {
            const view = this.#views[i];
            if (view.size.x !== res_x || view.size.y !== res_y) {
                view.size.set(res_x, res_y);
            }
        }

        // Recreate the shared ImageData buffer at the new slice size.
        if (res_z > 0) {
            this.#imageData = this.#views[0].context2d.getImageData(0, 0, res_x, res_y);
        } else {
            this.#imageData = null;
        }

        this.layout();
    }

    /**
     * Render the current channel data into each per-slice canvas. Cheap to call every
     * frame for small grids; expensive at high resolutions.
     */
    draw() {
        if (this.#imageData === null) {
            return;
        }

        const res_x = this.#resolution[0];
        const res_y = this.#resolution[1];
        const res_z = this.#resolution[2];
        const slice_size = res_x * res_y;
        const channels = this.#channels;
        const channel_count = channels.length;

        const out = this.#imageData.data;
        const sample = new Float32Array(channel_count);

        for (let z = 0; z < res_z; z++) {
            const view = this.#views[z];
            const z_off = z * slice_size;

            for (let p = 0; p < slice_size; p++) {
                const c = z_off + p;
                for (let k = 0; k < channel_count; k++) {
                    sample[k] = channels[k][c];
                }
                this.#drawMethod(out, p * 4, sample, channel_count);
            }

            view.context2d.putImageData(this.#imageData, 0, 0);
        }
    }

    /**
     * Reflow the per-slice canvases into a 2D grid that fits the view's current width.
     * Called automatically on link and on size change; safe to invoke manually.
     */
    layout() {
        const res_x = this.#resolution[0];
        const res_y = this.#resolution[1];
        const res_z = this.#resolution[2];

        if (res_x === 0 || res_z === 0) {
            return;
        }

        const width = this.size.x;
        const grid_width = max2(1, Math.floor(width / res_x));

        for (let i = 0; i < res_z; i++) {
            const view = this.#views[i];
            const grid_x = i % grid_width;
            const grid_y = Math.floor(i / grid_width);

            view.position.set(grid_x * res_x, grid_y * res_y);
        }
    }
}
