import { NDArray } from './ndarray_core.js';
import { NDWasm, fromWasm } from './ndwasm.js';
/** 
 * NDWasmBlas: BLAS (Basic Linear Algebra Subprograms)
 * Handles O(n^2) and O(n^3) matrix-matrix and matrix-vector operations.
 * @namespace NDWasmBlas
 */
export const NDWasmBlas = {

    /**
     * Calculates the trace of a 2D square matrix (sum of diagonal elements).
     * Complexity: O(n)
     * @param {NDArray} a
     * @returns {number} The sum of the diagonal elements.
     * @throws {Error} If the array is not 2D or not a square matrix.
     */
    trace(a) {
        // 1. Validation: Must be 2D
        if (a.ndim !== 2) {
            throw new Error(`Trace is only defined for 2D matrices, but this array is ${a.ndim}D.`);
        }

        // 2. Validation: Must be square (N x N)
        const rows = a.shape[0];
        const cols = a.shape[1];
        if (rows !== cols) {
            throw new Error(`Trace is only defined for square matrices, but this matrix is ${rows}x${cols}.`);
        }

        let sum = 0;
        const data = a.data;
        const offset = a.offset;
        const s0 = a.strides[0]; // Row stride
        const s1 = a.strides[1]; // Column stride

        // 3. Calculation: Sum A[i, i]
        // Memory index for element [i, i] is: offset + i * strides[0] + i * strides[1]
        // We can optimize the loop by pre-calculating the step size: s0 + s1
        const step = s0 + s1;
        for (let i = 0; i < rows; i++) {
            sum += data[offset + i * step];
        }

        return sum;
    },
    

    /**
     * General Matrix Multiplication (GEMM): C = A * B.
     * Complexity: O(m * n * k)
     * @memberof NDWasmBlas
     * @param {NDArray} a - Left matrix of shape [m, n].
     * @param {NDArray} b - Right matrix of shape [n, k].
     * @returns {NDArray} Result matrix of shape [m, k].
     */
    matmul(a, b) {
        if (a.shape[1] !== b.shape[0]) {
            throw new Error(`Matrix inner dimensions must match: ${a.shape[1]} != ${b.shape[0]}`);
        }
        if(b.ndim!==2 && b.ndim!==1){
            throw new Error(`Right operand must be a 2D matrix (or 1D vector), but got ${b.ndim}D.`);
        }
        const m = a.shape[0];
        const n = a.shape[1];
        const k = b.ndim === 2 ? b.shape[1] : 1;
        const outShape = [m, k];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        return NDWasm._compute([a, b], outShape, a.dtype, (aPtr, bPtr, outPtr) => {
            return NDWasm.runtime.exports[`MatMul${suffix}`](aPtr, bPtr, outPtr, m, n, k);
        });
    },

    /**
     * matPow computes A^k (Matrix Power).
     * Matrix Functions (O(n^3))
     * @memberof NDWasmBlas
     * @param {NDArray} a - Matrix of shape [n, n].
     * @returns {NDArray} Result matrix of shape [n, n].
     */
    matPow(a, k) {
        if (a.shape[0] !== a.shape[1]) {
            throw new Error(`Matrix must be square: ${a.shape[0]} != ${a.shape[1]}`);
        }
        const n = a.shape[0];
        const outShape = [n, n];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        return NDWasm._compute([a], outShape, a.dtype, (aPtr,  outPtr) => {
            return NDWasm.runtime.exports[`MatrixPower${suffix}`](aPtr, outPtr, n, k);
        });
    },


    /**
     * Batched Matrix Multiplication: C[i] = A[i] * B[i].
     * Common in deep learning inference.
     * Complexity: O(batch * m * n * k)
     * @memberof NDWasmBlas
     * @param {NDArray} a - Batch of matrices of shape [batch, m, n].
     * @param {NDArray} b - Batch of matrices of shape [batch, n, k].
     * @returns {NDArray} Result batch of shape [batch, m, k].
     */
    matmulBatch(a, b) {
        if (a.ndim !== 3 || b.ndim !== 3 || a.shape[0] !== b.shape[0]) {
            throw new Error("Input must be 3D batches with same batch size.");
        }
        if (a.shape[2] !== b.shape[1]) {
            throw new Error("Batch matrix inner dimensions must match.");
        }
        const batch = a.shape[0];
        const m = a.shape[1];
        const n = a.shape[2];
        const k = b.shape[2];
        const outShape = [batch, m, k];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        return NDWasm._compute([a, b], outShape, a.dtype, (aPtr, bPtr, outPtr) => {
            return NDWasm.runtime.exports[`MatMulBatch${suffix}`](aPtr, bPtr, outPtr, batch, m, n, k);
        });
    },

    /**
     * Symmetric Rank-K Update: C = alpha * A * A^T + beta * C.
     * Used for efficiently computing covariance matrices or Gram matrices.
     * Complexity: O(n^2 * k)
     * @memberof NDWasmBlas
     * @param {NDArray} a - Input matrix of shape [n, k].
     * @returns {NDArray} Symmetric result matrix of shape [n, n].
     */
    syrk(a) {
        const n = a.shape[0];
        const k = a.shape[1];
        const outShape = [n, n];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        return NDWasm._compute([a], outShape, a.dtype, (aPtr, outPtr) => {
            return NDWasm.runtime.exports[`Syrk${suffix}`](aPtr, outPtr, n, k);
        });
    },

    /**
     * Triangular System Solver: Solves A * X = B for X, where A is a triangular matrix.
     * Complexity: O(m^2 * n)
     * @memberof NDWasmBlas
     * @param {NDArray} a - Triangular matrix of shape [m, m].
     * @param {NDArray} b - Right-hand side matrix/vector of shape [m, n].
     * @returns {NDArray} Solution matrix X of shape [m, n].
     */
    trsm(a, b, lower = false) {
        if (a.shape[0] !== a.shape[1] || a.shape[0] !== b.shape[0]) {
            throw new Error("Dimension mismatch for triangular solver.");
        }
        const m = a.shape[0];
        const n = b.ndim === 1 ? 1 : b.shape[1];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        // We use _compute but must ensure b is copied to the output buffer first 
        // because Go's TRSM often operates on the B matrix storage.
        const wa = a.toWasm(NDWasm.runtime);
        const wb = b.toWasm(NDWasm.runtime); // wb already contains b's data
        try {
            NDWasm.runtime.exports[`Trsm${suffix}`](wa.ptr, wb.ptr, m, n, lower ? 1 : 0);
            return fromWasm(wb, b.shape, b.dtype);
        } finally {
            wa.dispose();
            wb.dispose();
        }
    },

    /**
     * Matrix-Vector Multiplication: y = A * x.
     * Complexity: O(m * n)
     * @memberof NDWasmBlas
     * @param {NDArray} a - Matrix of shape [m, n].
     * @param {NDArray} x - Vector of shape [n].
     * @returns {NDArray} Result vector of shape [m].
     */
    matVecMul(a, x) {
        if (a.shape[1] !== x.size) {
            throw new Error("Matrix-Vector dimension mismatch.");
        }
        const m = a.shape[0];
        const n = a.shape[1];
        const suffix = NDWasm.runtime._getSuffix(a.dtype);

        return NDWasm._compute([a, x], [m], a.dtype, (aPtr, xPtr, outPtr) => {
            return NDWasm.runtime.exports[`MatVecMul${suffix}`](aPtr, xPtr, outPtr, m, n);
        });
    },

    /**
     * Vector Outer Product (Rank-1 Update): A = x * y^T.
     * Complexity: O(m * n)
     * @memberof NDWasmBlas
     * @param {NDArray} x - Vector of shape [m].
     * @param {NDArray} y - Vector of shape [n].
     * @returns {NDArray} Result matrix of shape [m, n].
     */
    ger(x, y) {
        const m = x.size;
        const n = y.size;
        const outShape = [m, n];
        const suffix = NDWasm.runtime._getSuffix(x.dtype);

        return NDWasm._compute([x, y], outShape, x.dtype, (xPtr, yPtr, outPtr) => {
            return NDWasm.runtime.exports[`Ger${suffix}`](xPtr, yPtr, outPtr, m, n);
        });
    }
};
