NdArray.java

package org.sadisamir.ndarray;

import java.util.Arrays;
import java.util.Objects;

/**
 * A multidimensional array of floats supporting 1D and 2D layouts.
 * Provides NumPy-style operations for scientific computing in Java.
 */
public class NdArray {

    private static final int ONE_DIMENSIONAL = 1;
    private static final int TWO_DIMENSIONAL = 2;
    private static final int ELLIPSIS_THRESHOLD = 7;
    private static final int EDGE_ITEMS_TO_SHOW = 3;

    private final float[] flatData;
    private final int numDimensions;
    private final int[] shape;
    private final int totalElements;

    // ==================== Constructors ====================

    /**
     * Creates a 1D array from flat data.
     *
     * @param data the source data (copied defensively)
     * @throws NullPointerException if data is null
     */
    public NdArray(float[] data) {
        Objects.requireNonNull(data, "Input array must not be null");
        this.flatData = Arrays.copyOf(data, data.length);
        this.numDimensions = ONE_DIMENSIONAL;
        this.shape = new int[] { data.length };
        this.totalElements = data.length;
    }

    /**
     * Creates a 2D array from a matrix.
     *
     * @param matrix the source matrix (copied defensively)
     * @throws NullPointerException if matrix is null
     * @throws IllegalArgumentException if any row is null or rows have inconsistent lengths
     */
    public NdArray(float[][] matrix) {
        Objects.requireNonNull(matrix, "Input matrix must not be null");
        int[] dimensions = validate2DStructure(matrix);
        int rowCount = dimensions[0];
        int colCount = dimensions[1];

        this.flatData = flatten2DArray(matrix, rowCount, colCount);
        this.numDimensions = TWO_DIMENSIONAL;
        this.shape = new int[] { rowCount, colCount };
        this.totalElements = rowCount * colCount;
    }

    private NdArray(float[] flatData, int numDimensions, int[] shape) {
        this.flatData = Arrays.copyOf(flatData, flatData.length);
        this.numDimensions = numDimensions;
        this.shape = Arrays.copyOf(shape, shape.length);
        this.totalElements = flatData.length;
    }

    // ==================== 2D Array Helpers ====================

    private static int[] validate2DStructure(float[][] matrix) {
        int rowCount = matrix.length;
        if (rowCount == 0) {
            return new int[] { 0, 0 };
        }

        float[] firstRow = matrix[0];
        if (firstRow == null) {
            throw new IllegalArgumentException("rows must not be null");
        }
        int expectedColCount = firstRow.length;

        for (int rowIndex = 1; rowIndex < rowCount; rowIndex++) {
            float[] currentRow = matrix[rowIndex];
            if (currentRow == null) {
                throw new IllegalArgumentException("rows must not be null");
            }
            if (currentRow.length != expectedColCount) {
                throw new IllegalArgumentException("all rows must have the same length");
            }
        }

        return new int[] { rowCount, expectedColCount };
    }

    private static float[] flatten2DArray(float[][] matrix, int rowCount, int colCount) {
        float[] flattened = new float[rowCount * colCount];
        int flatIndex = 0;

        for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
            float[] currentRow = matrix[rowIndex];
            for (int colIndex = 0; colIndex < colCount; colIndex++) {
                flattened[flatIndex++] = currentRow[colIndex];
            }
        }

        return flattened;
    }

    // ==================== Factory Methods ====================

    /**
     * Creates a 1D array from the given data.
     */
    public static NdArray array(float[] data) {
        return new NdArray(data);
    }

    /**
     * Creates a 2D array from the given matrix.
     */
    public static NdArray array(float[][] data) {
        return new NdArray(data);
    }

    /**
     * Creates a 1D array of zeros with the specified size.
     *
     * @param size number of elements
     * @throws IllegalArgumentException if size is negative
     */
    public static NdArray zeros(int size) {
        if (size < 0) {
            throw new IllegalArgumentException(
                    "Size must be non-negative, but got: " + size);
        }
        return new NdArray(new float[size]);
    }

    /**
     * Creates a 1D array with values from 0 (inclusive) to stop (exclusive).
     */
    public static NdArray arange(float stop) {
        return arange(0.0f, stop, 1.0f);
    }

    /**
     * Creates a 1D array with values from start (inclusive) to stop (exclusive).
     */
    public static NdArray arange(float start, float stop) {
        return arange(start, stop, 1.0f);
    }

    /**
     * Creates a 1D array with evenly spaced values within a given interval.
     *
     * @param start beginning of interval (inclusive)
     * @param stop end of interval (exclusive)
     * @param step spacing between values
     * @throws IllegalArgumentException if any argument is non-finite or step is zero
     */
    public static NdArray arange(float start, float stop, float step) {
        validateArangeArguments(start, stop, step);

        int count = computeArangeLength(start, stop, step);
        float[] values = new float[count];
        float currentValue = start;

        for (int index = 0; index < count; index++) {
            values[index] = currentValue;
            currentValue += step;
        }

        return new NdArray(values);
    }

    private static void validateArangeArguments(float start, float stop, float step) {
        if (!Float.isFinite(start) || !Float.isFinite(stop) || !Float.isFinite(step)) {
            throw new IllegalArgumentException(
                    "start, stop and step must be finite values");
        }
        if (step == 0.0f) {
            throw new IllegalArgumentException("step must not be zero");
        }
    }

    private static int computeArangeLength(float start, float stop, float step) {
        boolean noElementsNeeded = (step > 0.0f && start >= stop)
                                || (step < 0.0f && start <= stop);
        if (noElementsNeeded) {
            return 0;
        }
        return (int) Math.ceil((stop - start) / step);
    }

    // ==================== Accessors ====================

    /**
     * Returns the number of dimensions (1 or 2).
     */
    public int getNdim() {
        return numDimensions;
    }

    /**
     * Returns a copy of the shape array.
     */
    public int[] getShape() {
        return Arrays.copyOf(shape, shape.length);
    }

    /**
     * Returns the total number of elements.
     */
    public int getSize() {
        return totalElements;
    }

    /**
     * Returns a copy of the underlying data as a flat array.
     */
    public float[] toArray() {
        return Arrays.copyOf(flatData, flatData.length);
    }

    /**
     * Converts a 2D array to a matrix representation.
     *
     * @throws IllegalStateException if this is not a 2D array
     */
    public float[][] toMatrix() {
        if (numDimensions != TWO_DIMENSIONAL) {
            throw new IllegalStateException(
                    "toMatrix is only available for 2D ndarrays");
        }

        int rowCount = shape[0];
        int colCount = shape[1];
        float[][] matrix = new float[rowCount][colCount];
        int flatIndex = 0;

        for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
            for (int colIndex = 0; colIndex < colCount; colIndex++) {
                matrix[rowIndex][colIndex] = flatData[flatIndex++];
            }
        }

        return matrix;
    }

    // ==================== Operations ====================

    /**
     * Returns a new array with element-wise addition.
     *
     * @throws NullPointerException if other is null
     * @throws IllegalArgumentException if shapes don't match
     */
    public NdArray add(NdArray other) {
        Objects.requireNonNull(other, "Operand must not be null");
        requireMatchingShape(other);

        float[] result = new float[totalElements];
        for (int index = 0; index < totalElements; index++) {
            result[index] = flatData[index] + other.flatData[index];
        }

        return new NdArray(result, numDimensions, shape);
    }

    /**
     * Returns a new array with a scalar added to every element.
     */
    public NdArray add(float scalar) {
        float[] result = new float[totalElements];
        for (int index = 0; index < totalElements; index++) {
            result[index] = flatData[index] + scalar;
        }

        return new NdArray(result, numDimensions, shape);
    }

    /**
     * Adds another array to this one in place.
     *
     * @throws NullPointerException if other is null
     * @throws IllegalArgumentException if shapes don't match
     */
    public void addInPlace(NdArray other) {
        Objects.requireNonNull(other, "Operand must not be null");
        requireMatchingShape(other);

        for (int index = 0; index < totalElements; index++) {
            flatData[index] += other.flatData[index];
        }
    }

    /**
     * Adds a scalar to this array in place.
     */
    public void addInPlace(float scalar) {
        for (int index = 0; index < totalElements; index++) {
            flatData[index] += scalar;
        }
    }

    /**
     * Returns the sum of all elements in the array.
     */
    public float sum() {
        double total = 0.0d;
        for (float value : flatData) {
            total += value;
        }
        return (float) total;
    }

    /**
     * Returns a reshaped view of this array.
     *
     * @throws NullPointerException if newShape is null
     * @throws IllegalArgumentException if shape is invalid or incompatible
     */
    public NdArray reshape(int... newShape) {
        Objects.requireNonNull(newShape, "Shape must not be null");

        if (newShape.length == 0) {
            throw new IllegalArgumentException(
                    "shape must have at least one dimension");
        }

        int[] targetShape = Arrays.copyOf(newShape, newShape.length);
        long targetSize = computeShapeProduct(targetShape);

        if (targetSize != totalElements) {
            throw new IllegalArgumentException(
                    "shape mismatch: cannot reshape array of size " + totalElements
                            + " into shape " + Arrays.toString(targetShape));
        }

        return createReshapedArray(targetShape);
    }

    private long computeShapeProduct(int[] dimensions) {
        long product = 1L;
        for (int dimension : dimensions) {
            if (dimension < 0) {
                throw new IllegalArgumentException(
                        "shape dimensions must be >= 0");
            }
            product *= dimension;
        }
        return product;
    }

    private NdArray createReshapedArray(int[] targetShape) {
        int targetDimensions = targetShape.length;

        if (targetDimensions == 1) {
            return new NdArray(flatData, ONE_DIMENSIONAL, targetShape);
        }
        if (targetDimensions == 2) {
            return new NdArray(flatData, TWO_DIMENSIONAL, targetShape);
        }

        throw new IllegalArgumentException(
                "reshape currently supports only 1D and 2D shapes");
    }

    // ==================== Validation ====================

    private void requireMatchingShape(NdArray other) {
        if (this.numDimensions != other.numDimensions
                || !Arrays.equals(this.shape, other.shape)) {
            throw new IllegalArgumentException(
                    "shape mismatch: arrays must have same shape");
        }
    }

    // ==================== String Formatting ====================

    @Override
    public String toString() {
        if (numDimensions == TWO_DIMENSIONAL) {
            return format2DArray();
        }
        return format1DArray();
    }

    private String format1DArray() {
        StringBuilder output = new StringBuilder("[");

        if (flatData.length <= ELLIPSIS_THRESHOLD) {
            appendElementRange(output, 0, flatData.length);
        } else {
            appendElementRange(output, 0, EDGE_ITEMS_TO_SHOW);
            output.append(" ... ");
            appendElementRange(output, flatData.length - EDGE_ITEMS_TO_SHOW, flatData.length);
        }

        output.append(']');
        return output.toString();
    }

    private String format2DArray() {
        int rowCount = shape[0];
        int colCount = shape[1];
        StringBuilder output = new StringBuilder("[");
        int flatIndex = 0;

        for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
            if (rowIndex == 0) {
                output.append('[');
            } else {
                output.append("\n [");
            }

            for (int colIndex = 0; colIndex < colCount; colIndex++) {
                if (colIndex > 0) {
                    output.append(' ');
                }
                output.append(formatFloat(flatData[flatIndex++]));
            }
            output.append(']');
        }

        output.append(']');
        return output.toString();
    }

    private void appendElementRange(StringBuilder output, int startIndex, int endIndex) {
        for (int index = startIndex; index < endIndex; index++) {
            if (index > startIndex) {
                output.append(' ');
            }
            output.append(formatFloat(flatData[index]));
        }
    }

    private static String formatFloat(float value) {
        if (value == 0.0f) {
            return "0";
        }

        long asWholeNumber = (long) value;
        if (value == asWholeNumber) {
            return Long.toString(asWholeNumber);
        }

        return Float.toString(value);
    }
}