import { BufferAttribute, BufferGeometry, StreamDrawUsage } from 'three';
import { assert } from "../../../../../../core/assert.js";
import { BitSet } from "../../../../../../core/binary/BitSet.js";
import {
    compute_typed_array_constructor_from_data_type
} from "../../../../../../core/binary/type/DataType2TypedArrayConstructorMapping.js";
import { typed_array_copy } from "../../../../../../core/collection/array/typed/typed_array_copy.js";
import { aabb3_from_v3_array } from "../../../../../../core/geom/3d/aabb/aabb3_from_v3_array.js";
import { max3 } from "../../../../../../core/math/max3.js";
import ObservedValue from "../../../../../../core/model/ObservedValue.js";

const SHRINK_THRESHOLD = 64;
const GROW_MIN_STEP = 32;

export class ParticlePool {
    /**
     *
     * @type {BitSet}
     */
    occupancy = new BitSet();

    /**
     *
     * @type {BufferAttribute[]}
     */
    attributes = [];

    /**
     * @type {ObservedValue.<BufferGeometry>}
     */
    geometry = new ObservedValue(null);

    /**
     *
     * @type {number}
     */
    growFactor = 1.3;

    /**
     *
     * @type {number}
     */
    shrinkFactor = 0.5;

    /**
     *
     * @type {number}
     */
    capacity = 16;

    /**
     *
     * @param {AttributeGroupSpec} spec
     * @constructor
     */
    constructor(spec) {
        assert.equal(spec.isAttributeGroupSpec, true, 'spec.VertexDataSpec !== true');

        /**
         *
         * @type {AttributeGroupSpec}
         */
        this.spec = spec;


        this.build();
    }

    size() {
        return this.occupancy.size();
    }

    /**
     *
     * @returns {number}
     */
    create() {
        //find an empty place in the pool
        const index = this.occupancy.nextClearBit(0);

        //reserve place
        this.occupancy.set(index, true);

        //ensure attribute arrays are large enough
        const min_capacity = index + 1;

        if (this.capacity < min_capacity) {
            this.growCapacity(min_capacity);
        }

        return index;
    }

    /**
     * Returns true if particle has discontinuous unused slots (holes)
     * @returns {boolean}
     */
    hasHoles() {
        const occupancy = this.occupancy;

        return occupancy.nextClearBit(0) < occupancy.size();
    }

    /**
     *
     * @param {number} index
     */
    remove(index) {
        this.occupancy.set(index, false);
    }

    /**
     * Remove holes from the occupancy set
     */
    compact() {
        const occupancy = this.occupancy;

        const specAttributes = this.spec.attributes;
        const numAttributes = specAttributes.length;

        let holeIndex, lastIndex, attributeIndex;

        holeIndex = occupancy.nextClearBit(0);
        lastIndex = occupancy.previousSetBit(occupancy.size());

        const attributes = this.attributes;

        while (holeIndex < lastIndex) {

            //move occupancy flag from the last to the hole
            occupancy.set(holeIndex, true);
            occupancy.set(lastIndex, false);

            //move attribute values
            for (attributeIndex = 0; attributeIndex < numAttributes; attributeIndex++) {
                const attribute = attributes[attributeIndex];

                const itemSize = attribute.itemSize;

                /**
                 *
                 * @type {Float32Array}
                 */
                const attribute_array = attribute.array;

                //do swaps in the old array
                const targetIndex = holeIndex * itemSize;
                const sourceIndex = lastIndex * itemSize;

                attribute_array.copyWithin(targetIndex, sourceIndex, sourceIndex + itemSize);
            }

            //move hole and last cursors
            holeIndex = occupancy.nextClearBit(holeIndex + 1);
            lastIndex = occupancy.previousSetBit(lastIndex);
        }

        this.geometry.getValue().setDrawRange(0, occupancy.size());
    }

    updateAttributes() {
        const attributes = this.attributes;
        const numAttributes = attributes.length;

        const activeParticleLimit = this.occupancy.size();

        if (activeParticleLimit <= 0) {
            // no particles, no need to update attributes
            return;
        }

        let i = 0;

        for (; i < numAttributes; i++) {
            const attribute = attributes[i];

            attribute.updateRange.count = activeParticleLimit * attribute.itemSize;
            attribute.updateRange.offset = 0;

            //mark attribute for update
            attribute.needsUpdate = true;
        }
    }

    update() {
        //some of the particles may have died and formed holes in the data array, compaction removes them
        this.compact();

        //attempt attribute size shrinkage
        const size = this.occupancy.size();
        this.shrinkCapacity(size);

        this.updateAttributes();
    }

    /**
     * Swap two particles identified by index
     * NOTE: Only attributes are swapped, occupancy is unaffected.
     * @param {number} indexA
     * @param {number} indexB
     */
    swap(indexA, indexB) {
        //assert.ok(this.occupancy.get(indexA), `indexA=${indexA} is unoccupied`);
        //assert.ok(this.occupancy.get(indexB), `indexB=${indexB} is unoccupied`);

        const attributes = this.attributes;
        const numAttributes = attributes.length;

        let i, j;

        for (i = 0; i < numAttributes; i++) {
            const attribute = attributes[i];

            const itemSize = attribute.itemSize;

            const offsetA = indexA * itemSize;
            const offsetB = indexB * itemSize;

            const array = attribute.array;

            for (j = 0; j < itemSize; j++) {
                const addressA = offsetA + j;
                const addressB = offsetB + j;

                const temp = array[addressA];

                array[addressA] = array[addressB];

                array[addressB] = temp;
            }
        }
    }

    build() {
        const geometry = new BufferGeometry();
        geometry.dynamic = true;

        const spec_attributes = this.spec.attributes;
        const l = spec_attributes.length;

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

            /**
             * @type {AttributeSpec}
             */
            const attributeSpec = spec_attributes[i];

            const itemSize = attributeSpec.itemSize;
            const TypedArrayConstructor = compute_typed_array_constructor_from_data_type(attributeSpec.type);

            const newArrayLength = itemSize * this.capacity;
            const newArray = new TypedArrayConstructor(newArrayLength);

            const old_attribute = this.attributes[i];

            if (old_attribute !== undefined && old_attribute.count !== 0) {
                //copy old data
                typed_array_copy(old_attribute.array, newArray);
            }

            const newAttribute = new BufferAttribute(newArray, itemSize, attributeSpec.normalized);

            newAttribute.setUsage(StreamDrawUsage);
            newAttribute.name = attributeSpec.name;

            this.attributes[i] = newAttribute;

            geometry.setAttribute(attributeSpec.name, newAttribute);
        }

        const oldGeometry = this.geometry.getValue();

        this.geometry.set(geometry);

        if (oldGeometry !== null) {
            //dispose of old geometry
            oldGeometry.dispose();
        }

        geometry.setDrawRange(0, this.occupancy.size());
    }

    /**
     *
     * @param {number} capacity
     */
    growCapacity(capacity) {
        const newCapacity = Math.floor(
            max3(
                capacity,
                this.capacity * this.growFactor,
                this.capacity + GROW_MIN_STEP
            )
        );

        this.setCapacity(newCapacity);
    }

    /**
     *
     * @param {number} newCapacity
     */
    shrinkCapacity(newCapacity) {
        if (
            newCapacity < this.capacity * this.shrinkFactor
            && newCapacity < this.capacity - SHRINK_THRESHOLD
        ) {

            this.setCapacity(newCapacity);

        }
    }

    /**
     *
     * @param {number} newCapacity
     */
    setCapacity(newCapacity) {
        if (this.capacity === newCapacity) {
            //do nothing
            return;
        }

        const occupancy_size = this.occupancy.size();

        if (newCapacity < occupancy_size) {
            throw new Error(`Can not set capacity (=${newCapacity}) below occupancy set size (=${occupancy_size})`);
        }

        this.capacity = newCapacity;

        this.build();
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @returns {number}
     */
    readAttributeScalar(particleIndex, attributeIndex) {
        const attribute = this.attributes[attributeIndex];

        const array = attribute.array;

        return array[particleIndex];
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number} value
     */
    writeAttributeScalar(particleIndex, attributeIndex, value) {
        const attribute = this.attributes[attributeIndex];

        const array = attribute.array;

        array[particleIndex] = value;
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number[]} result
     */
    readAttributeVector2(particleIndex, attributeIndex, result) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 2;

        const array = attribute.array;

        result[0] = array[address];
        result[1] = array[address + 1];
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number} x
     * @param {number} y
     */
    writeAttributeVector2(particleIndex, attributeIndex, x, y) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 2;

        const array = attribute.array;

        array[address] = x;
        array[address + 1] = y;
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number[]} result
     */
    readAttributeVector3(particleIndex, attributeIndex, result) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 3;

        const array = attribute.array;

        result[0] = array[address];
        result[1] = array[address + 1];
        result[2] = array[address + 2];
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number} x
     * @param {number} y
     * @param {number} z
     */
    writeAttributeVector3(particleIndex, attributeIndex, x, y, z) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 3;

        const array = attribute.array;

        array[address] = x;
        array[address + 1] = y;
        array[address + 2] = z;
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number[]} result
     */
    readAttributeVector4(particleIndex, attributeIndex, result) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 4;

        const array = attribute.array;

        result[0] = array[address];
        result[1] = array[address + 1];
        result[2] = array[address + 2];
        result[3] = array[address + 3];
    }

    /**
     *
     * @param {int} particleIndex
     * @param {int} attributeIndex
     * @param {number} x
     * @param {number} y
     * @param {number} z
     * @param {number} w
     */
    writeAttributeVector4(particleIndex, attributeIndex, x, y, z, w) {
        const attribute = this.attributes[attributeIndex];

        const address = particleIndex * 4;

        const array = attribute.array;

        array[address] = x;
        array[address + 1] = y;
        array[address + 2] = z;
        array[address + 3] = w;
    }

    /**
     * Write the same attribute value for every particle.
     * Assumes particles are packed (no holes).
     * @param {number} attribute_index
     * @param {number[]|Float32Array|Uint8Array|Uint16Array|Int8Array|Int32Array} value
     */
    writeAttributeForAllArray(attribute_index, value) {
        const attribute = this.attributes[attribute_index];

        const itemSize = attribute.itemSize;

        if (value.length < itemSize) {
            throw new Error(`Argument underflow, attribute.itemSize is '${itemSize}', but supplier argument's length is '${value.length}'`);
        }

        const n = this.size();

        const destination = attribute.array;

        let k;
        for (let i = 0; i < n; i++) {
            const offset = i * itemSize;

            for (k = 0; k < itemSize; k++) {
                destination[offset + k] = value[k];
            }
        }

        // mark for update
        attribute.needsUpdate = true;
        attribute.updateRange.count = n;
        attribute.updateRange.offset = 0;
    }

    /**
     *
     * @param {number} attribute_index
     * @param {number[]} result
     */
    computeAttributeVector3AxisAlignedBoundingBox(attribute_index, result) {
        const attribute = this.attributes[attribute_index];
        const item_size = attribute.itemSize;

        if (item_size !== 3) {
            throw new Error(`Attribute item_size is expected to be 3, instead was '${item_size}'`);
        }

        const particle_count = this.size();

        const data = attribute.array;

        aabb3_from_v3_array(result, data, particle_count * 3);

    }

    dispose() {
        const bufferGeometry = this.geometry.getValue();

        if (bufferGeometry !== null) {
            // release used GPU memory
            bufferGeometry.dispose();
        }
    }
}
