/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 *  * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 *  * License for the specific language governing permissions and limitations
 *  * under the License.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */

package org.nd4j.linalg.dataset.api;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.nd4j.common.base.Preconditions;
import org.nd4j.linalg.exception.ND4JIllegalStateException;
import org.nd4j.common.primitives.Pair;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.impl.broadcast.BroadcastMulOp;
import org.nd4j.linalg.api.shape.Shape;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.nd4j.linalg.indexing.NDArrayIndex;

import java.util.Arrays;

@Slf4j
public class DataSetUtil {
    public static INDArray tailor2d(@NonNull DataSet dataSet, boolean areFeatures) {
        return tailor2d(areFeatures ? dataSet.getFeatures() : dataSet.getLabels(),
                        areFeatures ? dataSet.getFeaturesMaskArray() : dataSet.getLabelsMaskArray());
    }

    public static INDArray tailor2d(@NonNull INDArray data, INDArray mask) {
        switch (data.rank()) {
            case 1:
            case 2:
                return data;
            case 3:
                return tailor3d2d(data, mask);
            case 4:
                return tailor4d2d(data);
            default:
                throw new RuntimeException("Unsupported data rank");
        }
    }

    /**
     * @deprecated
     */
    public static INDArray tailor3d2d(DataSet dataset, boolean areFeatures) {
        INDArray data = areFeatures ? dataset.getFeatures() : dataset.getLabels();
        INDArray mask = areFeatures ? dataset.getFeaturesMaskArray() : dataset.getLabelsMaskArray();
        return tailor3d2d(data, mask);
    }

    public static INDArray tailor3d2d(@NonNull INDArray data, INDArray mask) {
        //Check mask shapes:
        if (mask != null) {
            if (data.size(0) != mask.size(0) || data.size(2) != mask.size(1)) {
                throw new IllegalArgumentException(
                                "Invalid mask array/data combination: got data with shape [minibatch, vectorSize, timeSeriesLength] = "
                                                + Arrays.toString(data.shape())
                                                + "; got mask with shape [minibatch,timeSeriesLength] = "
                                                + Arrays.toString(mask.shape())
                                                + "; minibatch and timeSeriesLength dimensions must match");
            }
        }


        if (data.ordering() != 'f' || data.isView() || !Shape.strideDescendingCAscendingF(data)) {
            data = data.dup('f');
        }
        //F order: strides are like [1, miniBatch, minibatch*size] - i.e., each time step array is contiguous in memory
        //This can be reshaped to 2d with a no-copy op
        //Same approach as RnnToFeedForwardPreProcessor in DL4J
        //I.e., we're effectively stacking time steps for all examples

        val shape = data.shape();
        INDArray as2d;
        if (shape[0] == 1) {
            as2d = data.tensorAlongDimension(0, 1, 2).permutei(1, 0); //Edge case: miniBatchSize==1
        } else if (shape[2] == 1) {
            as2d = data.tensorAlongDimension(0, 1, 0); //Edge case: timeSeriesLength=1
        } else {
            INDArray permuted = data.permute(0, 2, 1); //Permute, so we get correct order after reshaping
            as2d = permuted.reshape('f', shape[0] * shape[2], shape[1]);
        }

        if (mask == null) {
            return as2d;
        }

        //With stride 1 along the examples (dimension 0), we are concatenating time series - same as the
        if (mask.ordering() != 'f' || mask.isView() || !Shape.strideDescendingCAscendingF(mask)) {
            mask = mask.dup('f');
        }

        INDArray mask1d = mask.reshape('f', new long[] {mask.length(), 1});

        //Assume masks are 0s and 1s: then sum == number of elements
        int numElements = mask.sumNumber().intValue();
        if (numElements == mask.length()) {
            return as2d; //All are 1s
        }
        if (numElements == 0) {
            return null;
        }

        int[] rowsToPull = new int[numElements];
        float[] floatMask1d = mask1d.data().asFloat();
        int currCount = 0;
        for (int i = 0; i < floatMask1d.length; i++) {
            if (floatMask1d[i] != 0.0f) {
                rowsToPull[currCount++] = i;
            }
        }

        INDArray subset = Nd4j.pullRows(as2d, 1, rowsToPull); //Tensor along dimension 1 == rows
        return subset;
    }

    public static INDArray tailor4d2d(DataSet dataset, boolean areFeatures) {
        return tailor4d2d(areFeatures ? dataset.getFeatures() : dataset.getLabels());
    }

    public static INDArray tailor4d2d(@NonNull INDArray data) {
        long instances = data.size(0);
        long channels = data.size(1);
        long height = data.size(2);
        long width = data.size(3);

        INDArray in2d = Nd4j.create(channels, height * width * instances);

        long tads = data.tensorsAlongDimension(3, 2, 0);
        for (int i = 0; i < tads; i++) {
            INDArray thisTAD = data.tensorAlongDimension(i, 3, 2, 0);
            in2d.putRow(i, Nd4j.toFlattened(thisTAD));
        }
        return in2d.transposei();
    }

    public static void setMaskedValuesToZero(INDArray data, INDArray mask) {
        if (mask == null || data.rank() != 3)
            return;

        Nd4j.getExecutioner().exec(new BroadcastMulOp(data, mask, data, 0, 2));
    }

    /**
     * Merge all of the features arrays into one minibatch.
     *
     * @param featuresToMerge     features to merge. Note that first index is the input array (example) index, the second
     *                            index is the input array.
     *                            Thus to merge 10 examples with 3 input arrays each, featuresToMerge will be indexed
     *                            like featuresToMerge[0..9][0..2]
     * @param featureMasksToMerge May be null. If non-null: feature masks to merge
     * @return Merged features, and feature masks. Note that feature masks may be added automatically, if required - even
     * if no feature masks were present originally
     */
    public static Pair<INDArray[], INDArray[]> mergeFeatures(@NonNull INDArray[][] featuresToMerge, INDArray[][] featureMasksToMerge) {
        int nInArrs = featuresToMerge[0].length;
        INDArray[] outF = new INDArray[nInArrs];
        INDArray[] outM = null;

        for (int i = 0; i < nInArrs; i++) {
            Pair<INDArray, INDArray> p = mergeFeatures(featuresToMerge, featureMasksToMerge, i);
            outF[i] = p.getFirst();
            if (p.getSecond() != null) {
                if (outM == null) {
                    outM = new INDArray[nInArrs];
                }
                outM[i] = p.getSecond();
            }
        }
        return new Pair<>(outF, outM);
    }

    /**
     * Merge the specified features and mask arrays (i.e., concatenate the examples)
     *
     * @param featuresToMerge     Features to merge
     * @param featureMasksToMerge Mask arrays to merge. May be null
     * @return Merged features and mask. Mask may be null
     */
    public static Pair<INDArray, INDArray> mergeFeatures(@NonNull INDArray[] featuresToMerge,
                    INDArray[] featureMasksToMerge) {
        Preconditions.checkNotNull(featuresToMerge[0], "Encountered null feature array when merging");
        int rankFeatures = featuresToMerge[0].rank();

        switch (rankFeatures) {
            case 2:
                return DataSetUtil.merge2d(featuresToMerge, featureMasksToMerge);
            case 3:
                return DataSetUtil.mergeTimeSeries(featuresToMerge, featureMasksToMerge);
            case 4:
                return DataSetUtil.merge4d(featuresToMerge, featureMasksToMerge);
            default:
                throw new IllegalStateException("Cannot merge examples: features rank must be in range 2 to 4"
                                + " inclusive. First example features shape: "
                                + Arrays.toString(featuresToMerge[0].shape()));
        }
    }

    /**
     * Extract out the specified column, and merge the specified features and mask arrays (i.e., concatenate the examples)
     *
     * @param featuresToMerge     Features to merge. Will use featuresToMerge[all][inOutIdx]
     * @param featureMasksToMerge Mask arrays to merge. May be null
     * @return Merged features and mask. Mask may be null
     */
    public static Pair<INDArray, INDArray> mergeFeatures(INDArray[][] featuresToMerge, INDArray[][] featureMasksToMerge,
                    int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(featuresToMerge, featureMasksToMerge, inOutIdx);
        return mergeFeatures(p.getFirst(), p.getSecond());
    }

    /**
     * Merge the specified labels and label mask arrays (i.e., concatenate the examples)
     *
     * @param labelsToMerge     Features to merge
     * @param labelMasksToMerge Mask arrays to merge. May be null
     * @return Merged features and mask. Mask may be null
     */
    public static Pair<INDArray, INDArray> mergeLabels(INDArray[] labelsToMerge, INDArray[] labelMasksToMerge) {
        Preconditions.checkNotNull(labelsToMerge[0], "Cannot merge data: Encountered null labels array");
        int rankFeatures = labelsToMerge[0].rank();

        switch (rankFeatures) {
            case 2:
                return DataSetUtil.merge2d(labelsToMerge, labelMasksToMerge);
            case 3:
                return DataSetUtil.mergeTimeSeries(labelsToMerge, labelMasksToMerge);
            case 4:
                return DataSetUtil.merge4d(labelsToMerge, labelMasksToMerge);
            default:
                throw new ND4JIllegalStateException("Cannot merge examples: labels rank must be in range 2 to 4"
                                + " inclusive. First example features shape: "
                                + Arrays.toString(labelsToMerge[0].shape()));
        }
    }

    /**
     * Extract out the specified column, and merge the specified label and label mask arrays
     * (i.e., concatenate the examples)
     *
     * @param labelsToMerge     Features to merge. Will use featuresToMerge[all][inOutIdx]
     * @param labelMasksToMerge Mask arrays to merge. May be null
     * @return Merged features and mask. Mask may be null
     */
    public static Pair<INDArray, INDArray> mergeLabels(@NonNull INDArray[][] labelsToMerge,
                    INDArray[][] labelMasksToMerge, int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(labelsToMerge, labelMasksToMerge, inOutIdx);
        return mergeLabels(p.getFirst(), p.getSecond());
    }

    private static Pair<INDArray[], INDArray[]> selectColumnFromMDSData(@NonNull INDArray[][] arrays,
                    INDArray[][] masks, int inOutIdx) {
        INDArray[] a = new INDArray[arrays.length];
        INDArray[] m = new INDArray[a.length];
        for (int i = 0; i < a.length; i++) {
            a[i] = arrays[i][inOutIdx];
            if (masks != null && masks[i] != null) {
                m[i] = masks[i][inOutIdx];
            }
        }
        return new Pair<>(a, m);
    }

    /**
     * Merge the specified 2d arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @param inOutIdx Index to extract out before merging
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> merge2d(@NonNull INDArray[][] arrays, INDArray[][] masks, int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(arrays, masks, inOutIdx);
        return merge2d(p.getFirst(), p.getSecond());
    }

    /**
     * Merge the specified 2d arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> merge2d(INDArray[] arrays, INDArray[] masks) {
        long cols = arrays[0].columns();

        INDArray[] temp = new INDArray[arrays.length];
        boolean hasMasks = false;
        for (int i = 0; i < arrays.length; i++) {
            Preconditions.checkNotNull(arrays[i], "Encountered null array at position %s when merging data", i);
            if (arrays[i].columns() != cols) {
                throw new IllegalStateException("Cannot merge 2d arrays with different numbers of columns (firstNCols="
                                + cols + ", ithNCols=" + arrays[i].columns() + ")");
            }

            temp[i] = arrays[i];

            if (masks != null && masks[i] != null && masks[i] != null) {
                hasMasks = true;
            }
        }

        INDArray out = Nd4j.specialConcat(0, temp);
        INDArray outMask = null;
        if (hasMasks) {
            outMask = DataSetUtil.mergePerOutputMasks2d(out.shape(), arrays, masks);
        }

        return new Pair<>(out, outMask);
    }


    public static INDArray mergePerOutputMasks2d(long[] outShape, INDArray[][] arrays, INDArray[][] masks, int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(arrays, masks, inOutIdx);
        return mergePerOutputMasks2d(outShape, p.getFirst(), p.getSecond());
    }

    /**
     * @deprecated USe {@link #mergeMasks2d(long[], INDArray[], INDArray[])}
     */
    @Deprecated
    public static INDArray mergePerOutputMasks2d(long[] outShape, INDArray[] arrays, INDArray[] masks) {
        return mergeMasks2d(outShape, arrays, masks);
    }

    public static INDArray mergeMasks2d(long[] outShape, INDArray[] arrays, INDArray[] masks) {
        val numExamplesPerArr = new long[arrays.length];
        for (int i = 0; i < numExamplesPerArr.length; i++) {
            numExamplesPerArr[i] = arrays[i].size(0);
        }

        INDArray outMask = Nd4j.ones(arrays[0].dataType(), outShape); //Initialize to 'all present' (1s)

        int rowsSoFar = 0;
        for (int i = 0; i < masks.length; i++) {
            long thisRows = numExamplesPerArr[i]; //Mask itself may be null -> all present, but may include multiple examples
            if (masks[i] == null) {
                continue;
            }

            outMask.put(new INDArrayIndex[] {NDArrayIndex.interval(rowsSoFar, rowsSoFar + thisRows),
                            NDArrayIndex.all()}, masks[i]);
            rowsSoFar += thisRows;
        }
        return outMask;
    }

    public static INDArray mergeMasks4d(INDArray[] featuresOrLabels, INDArray[] masks) {
        long[] outShape = null;
        long mbCountNoMask = 0;
        for (int i = 0; i < masks.length; i++) {
            if(masks[i] == null) {
                mbCountNoMask += featuresOrLabels[i].size(0);
                continue;
            }
            if(masks[i].rank() != 4)
                throw new IllegalStateException("Cannot merge mask arrays: expected mask array of rank 4. Got mask array of rank " + masks[i].rank()
                        + " with shape " + Arrays.toString(masks[i].shape()));
            if(outShape == null)
                outShape = masks[i].shape().clone();
            else {
                INDArray m = masks[i];
                if(m.size(1) != outShape[1] || m.size(2) != outShape[2] || m.size(3) != outShape[3]){
                    throw new IllegalStateException("Mismatched mask shapes: masks should have same depth/height/width for all examples." +
                            " Prior examples had shape [mb," + masks[1] + "," + masks[2] + "," + masks[3] + "], next example has shape " +
                            Arrays.toString(m.shape()));
                }
                outShape[0] += m.size(0);
            }
        }

        if(outShape == null)
            return null;    //No masks to merge

        outShape[0] += mbCountNoMask;

        INDArray outMask = Nd4j.ones(outShape); //Initialize to 'all present' (1s)

        int exSoFar = 0;
        for (int i = 0; i < masks.length; i++) {
            if (masks[i] == null) {
                exSoFar += featuresOrLabels[i].size(0);
                continue;
            }
            long nEx = masks[i].size(0);

            outMask.put(new INDArrayIndex[] {NDArrayIndex.interval(exSoFar, exSoFar + nEx),
                    NDArrayIndex.all()}, masks[i]);
            exSoFar += nEx;
        }
        return outMask;
    }

    /**
     * Merge the specified time series (3d) arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @param inOutIdx Index to extract out before merging
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> mergeTimeSeries(INDArray[][] arrays, INDArray[][] masks, int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(arrays, masks, inOutIdx);
        return mergeTimeSeries(p.getFirst(), p.getSecond());
    }

    /**
     * Merge the specified time series (3d) arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> mergeTimeSeries(INDArray[] arrays, INDArray[] masks) {
        //Merge time series data, and handle masking etc for different length arrays

        //Complications with time series:
        //(a) They may have different lengths (if so: need input + output masking arrays)
        //(b) Even if they are all the same length, they may have masking arrays (if so: merge the masking arrays too)
        //(c) Furthermore: mask arrays can be per-time-step (2d) or per output (3d). Per-input masks (3d feature masks)
        //    are not supported, however

        long firstLength = arrays[0].size(2);
        long size = arrays[0].size(1);
        long maxLength = firstLength;

        boolean hasMask = false;
        int maskRank = -1;
        boolean lengthsDiffer = false;
        int totalExamples = 0;
        for (int i = 0; i < arrays.length; i++) {
            totalExamples += arrays[i].size(0);
            long thisLength = arrays[i].size(2);
            maxLength = Math.max(maxLength, thisLength);
            if (thisLength != firstLength)
                lengthsDiffer = true;
            if (masks != null && masks[i] != null && masks[i] != null) {
                maskRank = masks[i].rank();
                hasMask = true;
            }

            if (arrays[i].size(1) != size) {
                throw new IllegalStateException(
                                "Cannot merge time series with different size for dimension 1 (first shape: "
                                                + Arrays.toString(arrays[0].shape()) + ", " + i + "th shape: "
                                                + Arrays.toString(arrays[i].shape()));
            }
        }

        boolean needMask = hasMask || lengthsDiffer;
        INDArray arr = Nd4j.create(arrays[0].dataType(), totalExamples, size, maxLength);
        INDArray mask = (needMask && maskRank != 3 ? Nd4j.ones(arrays[0].dataType(), totalExamples, maxLength) : null);

        //Now, merge the time series (and if necessary, mask arrays):
        int examplesSoFar = 0;
        if (!lengthsDiffer && !needMask) {
            //Simplest case: same length, no mask arrays
            for (int i = 0; i < arrays.length; i++) {
                long thisNExamples = arrays[i].size(0);
                arr.put(new INDArrayIndex[] {NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                NDArrayIndex.all(), NDArrayIndex.all()}, arrays[i]);
                examplesSoFar += thisNExamples;
            }
            return new Pair<>(arr, null);
        } else {
            //Either different length, or have mask arrays (or, both)
            if ((lengthsDiffer && !hasMask) || maskRank == 2) {
                //Standard per-example masking required
                for (int i = 0; i < arrays.length; i++) {
                    INDArray a = arrays[i];
                    long thisNExamples = a.size(0);
                    long thisLength = a.size(2);
                    arr.put(new INDArrayIndex[] {NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                    NDArrayIndex.all(), NDArrayIndex.interval(0, thisLength)}, a);

                    if (masks != null && masks[i] != null && masks[i] != null) {
                        INDArray origMask = masks[i];
                        long maskLength = origMask.size(1);
                        mask.put(new INDArrayIndex[] {
                                        NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                        NDArrayIndex.interval(0, maskLength)}, origMask);
                        if (maskLength < maxLength) {
                            //Set end mask array to zero...
                            mask.put(new INDArrayIndex[] {
                                            NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                            NDArrayIndex.interval(maskLength, maxLength)},
                                            Nd4j.zeros(thisNExamples, maxLength - maskLength));
                        }
                    } else {
                        if (thisLength < maxLength) {
                            //Mask the end
                            mask.put(new INDArrayIndex[] {
                                            NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                            NDArrayIndex.interval(thisLength, maxLength)},
                                            Nd4j.zeros(thisNExamples, maxLength - thisLength));
                        }
                    }

                    examplesSoFar += thisNExamples;
                }
            } else if (maskRank == 3) {
                //Per output masking required. May also be variable length
                mask = Nd4j.create(arr.dataType(), arr.shape());
                for (int i = 0; i < arrays.length; i++) {
                    INDArray m = masks[i];
                    INDArray a = arrays[i];
                    long thisNExamples = a.size(0);
                    long thisLength = a.size(2);
                    arr.put(new INDArrayIndex[] {NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                    NDArrayIndex.all(), NDArrayIndex.interval(0, thisLength)}, a);

                    if (m == null) {
                        //This mask is null -> equivalent to "all present"
                        mask.get(NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                        NDArrayIndex.all(), NDArrayIndex.interval(0, thisLength)).assign(1);
                    } else {
                        mask.put(new INDArrayIndex[] {
                                        NDArrayIndex.interval(examplesSoFar, examplesSoFar + thisNExamples),
                                        NDArrayIndex.all(), NDArrayIndex.interval(0, thisLength)}, m);
                    }

                    examplesSoFar += thisNExamples;
                }
            } else {
                throw new UnsupportedOperationException("Cannot merge time series with mask rank " + maskRank);
            }
        }

        return new Pair<>(arr, mask);
    }

    /**
     * Merge the specified 4d arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @param inOutIdx Index to extract out before merging
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> merge4d(INDArray[][] arrays, INDArray[][] masks, int inOutIdx) {
        Pair<INDArray[], INDArray[]> p = selectColumnFromMDSData(arrays, masks, inOutIdx);
        return merge4d(p.getFirst(), p.getSecond());
    }

    /**
     * Merge the specified 4d arrays and masks. See {@link #mergeFeatures(INDArray[], INDArray[])}
     * and {@link #mergeLabels(INDArray[], INDArray[])}
     *
     * @param arrays   Arrays to merge
     * @param masks    Mask arrays to merge
     * @return Merged arrays and mask
     */
    public static Pair<INDArray, INDArray> merge4d(INDArray[] arrays, INDArray[] masks) {
        //4d -> images. In principle: could have 2d mask arrays (per-example masks)

        int nExamples = 0;
        long[] shape = arrays[0].shape();
        INDArray[] temp = new INDArray[arrays.length];
        boolean hasMasks = false;
        int maskRank = -1;
        for (int i = 0; i < arrays.length; i++) {
            Preconditions.checkNotNull(arrays[i], "Encountered null array when merging data at position %s", i);
            nExamples += arrays[i].size(0);
            long[] thisShape = arrays[i].shape();
            if (thisShape.length != 4) {
                throw new IllegalStateException("Cannot merge 4d arrays with non 4d arrays");
            }
            for (int j = 1; j < 4; j++) {
                if (thisShape[j] != shape[j])
                    throw new IllegalStateException(
                                    "Cannot merge 4d arrays with different shape (other than # examples): "
                                                    + " data[0].shape = " + Arrays.toString(shape) + ", data[" + i
                                                    + "].shape = " + Arrays.toString(thisShape));
            }

            temp[i] = arrays[i];
            if (masks != null && masks[i] != null ) {
                hasMasks = true;
                maskRank = masks[i].rank();
            }
        }

        INDArray out = Nd4j.specialConcat(0, temp);
        INDArray outMask = null;
        if (hasMasks) {
            if(maskRank == 2) {
                outMask = DataSetUtil.mergeMasks2d(out.shape(), arrays, masks);
            } else if(maskRank == 4){
                outMask = DataSetUtil.mergeMasks4d(arrays, masks);
            }
        }

        return new Pair<>(out, outMask);
    }
}
