import { DataTexture2DArray, LinearFilter, NearestFilter, RedFormat, UnsignedByteType } from "three";
import { assert } from "../../../../../core/assert.js";
import { Base64 } from "../../../../../core/binary/base64/Base64.js";
import { array_copy } from "../../../../../core/collection/array/array_copy.js";
import { typed_array_copy } from "../../../../../core/collection/array/typed/typed_array_copy.js";
import Vector2 from "../../../../../core/geom/Vector2.js";
import { max2 } from "../../../../../core/math/max2.js";
import { min2 } from "../../../../../core/math/min2.js";
import { countTask } from "../../../../../core/process/task/util/countTask.js";
import { sampler2d_scale } from "../../../../graphics/texture/sampler/resize/sampler2d_scale.js";
import { Sampler2D } from "../../../../graphics/texture/sampler/Sampler2D.js";


/**
 * Represents how various materials are mixed across the terrain
 */
export class SplatMapping {
    constructor() {
        /**
         * How large is the map
         * @type {Vector2}
         */
        this.size = new Vector2(1, 1);

        /**
         * number of layers
         * @type {number}
         */
        this.depth = 1;

        const size = this.size;

        const width = size.x;
        const height = size.y;

        /**
         * Weights of materials identified by material texture. Larger weight causes material to appear more prominently
         * @type {DataTexture2DArray}
         */
        this.weightTexture = new DataTexture2DArray(new Uint8Array(width * height), width, height, 1);
        this.weightTexture.generateMipmaps = false;
        this.weightTexture.format = RedFormat;
        this.weightTexture.type = UnsignedByteType;
        this.weightTexture.magFilter = LinearFilter;
        this.weightTexture.minFilter = NearestFilter;
        this.weightTexture.needsUpdate = true;

    }

    toJSON() {
        const encoded = Base64.encode(this.weightData.buffer);

        return {
            size: this.size.toJSON(),
            depth: this.depth,
            data: encoded
        };
    }

    fromJSON({
                 size = { x: 1, y: 1 },
                 depth = 0,
                 data
             } = {}) {
        this.resize(size.x, size.y, depth);

        let uint_data;

        if (typeof data === "string") {

            const decoded_data = Base64.decode(data);
            uint_data = new Uint8Array(decoded_data);

        } else if (Array.isArray(data)) {

            // deprecated
            console.warn('Array JSON format is deprecated, please upgrade your data as soon as possible');
            uint_data = data;

        } else {

            throw new Error('Unsupported data format');

        }

        array_copy(uint_data, 0, this.weightData, 0, data.length);

        this.weightTexture.needsUpdate = true;
    }

    /**
     *
     * @param {Sampler2D} source
     * @param {number} index
     * @param {number} channel
     */
    writeLayerFromSampler(source, index, channel = 0) {
        const weight_data = this.weightData;

        const width = this.size.x;
        const height = this.size.y;

        const dataOffset = index * width * height;


        if (source.width === width && source.height === height) {

            // special fast path
            const pixel_count = width * height;

            const source_channel_count = source.itemSize;
            const source_data = source.data;

            for (let i = 0; i < pixel_count; i++) {
                weight_data[dataOffset + i] = source_data[i * source_channel_count + channel]
            }

        } else {

            for (let y = 0; y < height; y++) {

                const v = (y + 0.5) / height;

                for (let x = 0; x < width; x++) {

                    const u = (x + 0.5) / width;

                    const targetIndex = y * width + x;
                    const targetAddress = dataOffset + targetIndex;

                    let value = source.sampleChannelBilinearUV(u, v, channel);

                    if (value > 0 && value < 1) {
                        // avoid collapse to 0
                        value = 1;
                    }

                    weight_data[targetAddress] = value;
                }
            }

        }

        // mark for update
        this.weightTexture.needsUpdate = true;
    }

    /**
     *
     * @param {Sampler2D} destination
     * @param {number} index
     * @param {number} channel
     */
    readLayerToSampler(destination, index, channel = 0) {
        const weightData = this.weightData;

        const width = this.size.x;
        const height = this.size.y;

        assert.equal(destination.width, width, 'destination dimensions mismatch');
        assert.equal(destination.height, height, 'destination dimensions mismatch');

        const dataOffset = index * width * height;

        for (let y = 0; y < height; y++) {

            for (let x = 0; x < width; x++) {

                const targetIndex = y * width + x;
                const targetAddress = dataOffset + targetIndex;

                const value = weightData[targetAddress];

                destination.writeChannel(x, y, channel, value);

            }
        }

    }

    /**
     *
     * @param {Uint8Array} source
     * @param {number} destinationX
     * @param {number} destinationY
     * @param {number} sourceWidth
     * @param {number} sourceHeight
     */
    writeWeightData(source, destinationX, destinationY, sourceWidth, sourceHeight) {
        const depth = this.depth;

        const sourceLayerSize = sourceWidth * sourceHeight;

        assert.equal(source.length, sourceLayerSize * depth, 'source size is incorrect');

        const w = this.size.x;
        const h = this.size.y;

        const destinationLayerSize = w * h;

        const _w = max2(0, min2(w - destinationX, sourceWidth));
        const _h = max2(0, min2(h - destinationY, sourceHeight));

        const dData = this.weightData;

        for (let i = 0; i < depth; i++) {

            const dLayerAddress = destinationLayerSize * i;
            const sLayerAddress = sourceLayerSize * i;

            for (let y = 0; y < _h; y++) {
                const dY = y + destinationY;

                const dRowAddress = dY * w + dLayerAddress;

                const sRowAddress = y * sourceWidth + sLayerAddress;

                for (let x = 0; x < _w; x++) {
                    const dX = x + destinationX;

                    const dAddress = dRowAddress + dX;

                    const sAddress = sRowAddress + x;

                    dData[dAddress] = source[sAddress];
                }
            }

        }
    }

    /**
     *
     * @param {Uint8Array} destination
     * @param {number} sourceX
     * @param {number} sourceY
     * @param {number} width
     * @param {number} height
     */
    readWeightData(destination, sourceX, sourceY, width, height) {
        const depth = this.depth;

        const w = this.size.x;
        const h = this.size.y;

        const sourceLayerSize = w * h;
        const destinationLayerSize = width * height;

        assert.equal(destination.length, destinationLayerSize * depth, 'destination size is incorrect');


        const _w = max2(0, min2(w - sourceX, width));
        const _h = max2(0, min2(h - sourceY, height));

        const sData = this.weightData;

        for (let i = 0; i < depth; i++) {

            const dLayerAddress = destinationLayerSize * i;
            const sLayerAddress = sourceLayerSize * i;

            for (let y = 0; y < _h; y++) {
                const dRowAddress = y * width + dLayerAddress;

                const sRowAddress = (y + sourceY) * w + sLayerAddress;

                for (let x = 0; x < _w; x++) {
                    const dAddress = dRowAddress + x;

                    const sAddress = sRowAddress + x + sourceX;

                    destination[dAddress] = sData[sAddress];
                }
            }

        }
    }

    /**
     *
     * @param {number} u
     * @param {number} v
     * @param {number} materialIndex
     * @returns {number} between 0 and 255
     */
    sampleWeight(u, v, materialIndex) {

        return this.getLayerWeightSampler(materialIndex).sampleChannelBilinearUV(u, v, 0);
    }

    /**
     *
     * @param {number} index
     * @returns {Sampler2D}
     */
    getLayerWeightSampler(index) {

        const width = this.size.x;
        const height = this.size.y;

        const layerSize = width * height;

        const startAddress = layerSize * index;
        const endAddress = startAddress + layerSize;

        //build a sampler for the layer
        return new Sampler2D(this.weightData.subarray(startAddress, endAddress), 1, width, height);
    }

    /**
     * @returns {Uint8ClampedArray|Uint8Array}
     */
    get weightData() {
        return this.weightTexture.image.data;
    }

    /**
     * @deprecated
     * @returns {Uint8ClampedArray|Uint8Array}
     */
    get materialData() {
        throw new Error('deprecated, use getLayerWeightSampler');
    }

    /**
     * @deprecated
     * @return {Sampler2D}
     */
    get materialSampler() {
        throw new Error('deprecated, use getLayerWeightSampler');
    }

    /**
     * Computes a making map, where each texel of each layer is given a rank based on it's weight relative to other layers
     * @param {Uint8Array} result
     * @returns {Task}
     */
    computeWeightRankingMap(result) {
        const w = this.size.x;
        const h = this.size.y;
        const depth = this.depth;

        const weightData = this.weightData;

        const sample_weights = [];
        const sample_materials = [];

        function compareSampleWeights(m0, m1) {
            const w0 = sample_weights[m0];
            const w1 = sample_weights[m1];

            return w1 - w0;
        }

        const layerSize = w * h;


        return countTask(0, w * h, index => {

            //read sample for all layers
            for (let d = 0; d < depth; d++) {

                const layerOffset = d * layerSize;

                const address = index + layerOffset;

                sample_materials[d] = d;
                sample_weights[d] = weightData[address];
            }

            //sort the layers
            sample_materials.sort(compareSampleWeights);

            //write ranking
            for (let rank = 0; rank < depth; rank++) {
                const material = sample_materials[rank];

                const layerOffset = material * layerSize;

                const address = index + layerOffset;

                result[address] = rank;
            }

        });

    }

    dispose() {
        this.weightTexture.dispose();
    }

    /**
     * Cut a specified layer, depth will be reduced by 1
     * @param {number} index
     * @returns {boolean}
     */
    removeWeightLayer(index) {

        if (index >= this.depth) {
            //layer doesn't exist
            return false;
        }

        this.depth--;

        const layerByteSize = this.size.x * this.size.y;

        const startAddress = index * layerByteSize;
        const endAddress = (index + 1) * layerByteSize;

        const weightTexture = this.weightTexture;

        // release memory
        weightTexture.dispose();

        const image = weightTexture.image;

        const oldData = image.data;

        oldData.copyWithin(startAddress, endAddress);

        const newData = new Uint8Array(this.depth * layerByteSize);


        newData.set(oldData.subarray(0, newData.length));

        image.data = newData;

        image.depth = this.depth;

        weightTexture.needsUpdate = true;

        return true;
    }

    /**
     *
     * @param {number} layer_index
     * @param {number} value
     */
    fillLayerWeights(layer_index, value) {
        assert.isNonNegativeInteger(layer_index, 'layer_index');
        assert.isNumber(value, 'value');

        const w = this.size.x;
        const h = this.size.y;

        const layerSize = w * h;

        const startAddress = layerSize * layer_index;
        const endAddress = startAddress + layerSize;

        const weightData = this.weightData;

        for (let i = startAddress; i < endAddress; i++) {
            weightData[i] = value;
        }

        this.weightTexture.needsUpdate = true;
    }

    addWeightLayer() {
        // release currently held GPU memory
        this.weightTexture.dispose();

        const layerByteSize = this.size.x * this.size.y;

        this.depth++;

        const depth = this.depth;

        const data = new Uint8Array(layerByteSize * depth);

        data.set(this.weightData);

        this.weightTexture.image.data = data;
        this.weightTexture.image.depth = depth;
        this.weightTexture.needsUpdate = true;
    }

    /**
     *
     * @param {number} width
     * @param {number} height
     * @param {number} depth
     * @param {boolean} [keep_data]
     */
    resize(width, height, depth, keep_data = true) {
        const oldWidth = this.size.x;
        const oldHeight = this.size.y;
        const oldDepth = this.depth;

        this.depth = depth;

        if (width !== oldWidth || height !== oldHeight || depth !== oldDepth) {
            this.size.set(width, height);
        }

        const weightTexture = this.weightTexture;
        const weightImage = weightTexture.image;

        // glTexStorage3D requires all dims > 0; uploading a 0-depth array texture
        // produces GL_INVALID_VALUE. Track size/depth at the JS level (above) but
        // leave the GPU texture at its prior valid state.
        if (width <= 0 || height <= 0 || depth <= 0) {
            return;
        }

        if (weightImage.width !== width || weightImage.height !== height || weightImage.depth !== depth) {
            // drop old data
            weightTexture.dispose();

            // this extra check is in case dimensions were modified from outside

            const oldWeightData = weightImage.data;

            const oldLayerSize = oldWidth * oldHeight;
            const newLayerSize = width * height;

            weightImage.width = width;
            weightImage.height = height;
            weightImage.depth = depth;
            weightImage.data = new Uint8Array(newLayerSize * depth);

            if (keep_data) {

                if (oldWidth === width && oldHeight === height) {
                    // dimensions haven't changed, we can copy data directly
                    typed_array_copy(oldWeightData, weightImage.data);
                } else {

                    // different dimensions, copy layer-by-layer
                    for (let d = 0; d < min2(depth, oldDepth); d++) {

                        const source = new Sampler2D(oldWeightData.subarray(d * oldLayerSize, (d + 1) * oldLayerSize), 1, oldWidth, oldHeight);
                        const target = new Sampler2D(weightImage.data.subarray(d * newLayerSize, (d + 1) * newLayerSize), 1, width, height);

                        sampler2d_scale(source, target);
                    }
                }

            }

            weightTexture.needsUpdate = true;

        }
    }
}
