# ndarray_prob.js

    import { NDArray, DTYPE_MAP } from './ndarray_core.js';

    /**
     * NDArray Probability & Random Module
     * Uses Cryptographically Strong Pseudo-Random Number Generator (CSPRNG).
     * All functions return a new NDArray instance.
     * @namespace NDProb
     */
    export const NDProb = {
        /**
         * Internal helper to get a Float64 between [0, 1) using crypto.
         * Generates high-quality uniform random floats.
         * @private
         */
        _cryptoUniform01(size) {
            const uint32 = new Uint32Array(size);
            // Browser & Node.js compatible crypto access
            const cryptoObj = globalThis.crypto;

            // getRandomValues has a limit (usually 65536 bytes), so we chunk if necessary
            const quota = 65536 / 4;
            for (let i = 0; i < size; i += quota) {
                const chunk = uint32.subarray(i, Math.min(i + quota, size));
                cryptoObj.getRandomValues(chunk);
            }

            const floats = new Float64Array(size);
            for (let i = 0; i < size; i++) {
                // Divide by 2^32 to get [0, 1)
                floats[i] = uint32[i] / 4294967296.0;
            }
            return floats;
        },

        /**
         * Uniform distribution over [low, high).
         * @memberof NDProb
         * @function
         * @param {Array} shape - Dimensions of the output array.
         * @param {number} [low=0] - Lower boundary.
         * @param {number} [high=1] - Upper boundary.
         * @param {string} [dtype='float64'] - Data type.
         * @returns {NDArray}
         */
        random(shape, low = 0, high = 1, dtype = 'float64') {
            const size = shape.reduce((a, b) => a * b, 1);
            const data = this._cryptoUniform01(size);
            const range = high - low;

            if (range !== 1 || low !== 0) {
                for (let i = 0; i < size; i++) data[i] = data[i] * range + low;
            }

            // Cast to specified dtype if necessary
            const finalData = dtype === 'float64' ? data : this._cast(data, dtype);
            return new NDArray(finalData, { shape, dtype });
        },

        /**
         * Normal (Gaussian) distribution using Box-Muller transform.
         * @param {Array} shape - Dimensions of the output array.
         * @param {number} [mean=0] - Mean of the distribution.
         * @param {number} [std=1] - Standard deviation.
         * @param {string} [dtype='float64'] - Data type.
         * @memberof NDProb
         * @returns {NDArray}
         */
        normal(shape, mean = 0, std = 1, dtype = 'float64') {
            const size = shape.reduce((a, b) => a * b, 1);
            // Box-Muller requires two uniform variables to produce two normal variables
            const u1 = this._cryptoUniform01(Math.ceil(size / 2) * 2);
            const u2 = this._cryptoUniform01(u1.length);
            const data = new Float64Array(size);

            for (let i = 0; i < size; i += 2) {
                const r = Math.sqrt(-2.0 * Math.log(Math.max(u1[i], 1e-10)));
                const theta = 2.0 * Math.PI * u2[i];

                data[i] = r * Math.cos(theta) * std + mean;
                if (i + 1 < size) {
                    data[i + 1] = r * Math.sin(theta) * std + mean;
                }
            }

            const finalData = dtype === 'float64' ? data : this._cast(data, dtype);
            return new NDArray(finalData, { shape, dtype });
        },

        /**
         * Bernoulli distribution (0 or 1 with probability p).
         * @param {Array} shape
         * @param {number} [p=0.5] - Probability of success (1).
         * @param {string} [dtype='int32']
         * @memberof NDProb
         * @returns {NDArray}
         */
        bernoulli(shape, p = 0.5, dtype = 'int32') {
            const size = shape.reduce((a, b) => a * b, 1);
            const u = this._cryptoUniform01(size);
            const data = new DTYPE_MAP[dtype](size);

            for (let i = 0; i < size; i++) {
                data[i] = u[i] < p ? 1 : 0;
            }
            return new NDArray(data, { shape, dtype });
        },

        /**
         * Exponential distribution: f(x; λ) = λe^(-λx).
         * Inverse transform sampling: x = -ln(1-u) / λ.
         * @param {Array} shape
         * @param {number} [lambda=1.0] - Rate parameter.
         * @memberof NDProb
         * @returns {NDArray}
         */
        exponential(shape, lambda = 1.0, dtype = 'float64') {
            const size = shape.reduce((a, b) => a * b, 1);
            const u = this._cryptoUniform01(size);
            const data = new Float64Array(size);

            for (let i = 0; i < size; i++) {
                // Using Math.max to avoid log(0)
                data[i] = -Math.log(Math.max(1.0 - u[i], 1e-10)) / lambda;
            }

            const finalData = dtype === 'float64' ? data : this._cast(data, dtype);
            return new NDArray(finalData, { shape, dtype });
        },

        /**
         * Poisson distribution using Knuth's algorithm.
         * Note: For very large lambda, this becomes slow; but for most use cases it's fine.
         * @param {Array} shape
         * @param {number} [lambda=1.0] - Mean of the distribution.
         * @memberof NDProb
         * @returns {NDArray}
         */
        poisson(shape, lambda = 1.0, dtype = 'int32') {
            const size = shape.reduce((a, b) => a * b, 1);
            const data = new DTYPE_MAP[dtype](size);
            const L = Math.exp(-lambda);

            for (let i = 0; i < size; i++) {
                let k = 0;
                let p = 1.0;
                do {
                    k++;
                    // We fetch one by one here as Poisson count per element is variable
                    p *= this._cryptoUniform01(1)[0];
                } while (p > L);
                data[i] = k - 1;
            }
            return new NDArray(data, { shape, dtype });
        },

        /**
         * @private
         * @internal
         */
        _cast(data, dtype) {
            const out = new DTYPE_MAP[dtype](data.length);
            out.set(data);
            return out;
        }
    };