/*
 *  ******************************************************************************
 *  *
 *  *
 *  * 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.convolution;


import lombok.val;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.impl.layers.convolution.Col2Im;
import org.nd4j.linalg.api.ops.impl.layers.convolution.Im2col;
import org.nd4j.linalg.api.ops.impl.layers.convolution.Pooling2D;
import org.nd4j.linalg.api.ops.impl.layers.convolution.config.Conv2DConfig;
import org.nd4j.linalg.api.ops.impl.layers.convolution.config.PaddingMode;
import org.nd4j.linalg.api.ops.impl.layers.convolution.config.Pooling2DConfig;
import org.nd4j.linalg.factory.Nd4j;


public class Convolution {


    public enum Type {
        FULL, VALID, SAME
    }


    /**
     * Default no-arg constructor.
     */
    private Convolution() {
    }

    /**
     * @param col
     * @param stride
     * @param padding
     * @param height
     * @param width
     * @return
     */
    public static INDArray col2im(INDArray col, int[] stride, int[] padding, int height, int width) {
        return col2im(col, stride[0], stride[1], padding[0], padding[1], height, width);
    }

    /**
     * Rearrange matrix
     * columns into blocks
     *
     * @param col the column
     *            transposed image to convert
     * @param sH  stride height
     * @param sW  stride width
     * @param ph  padding height
     * @param pW  padding width
     * @param kH  height
     * @param kW  width
     * @return
     */
    public static INDArray col2im(INDArray col, int sH, int sW, int ph, int pW, int kH, int kW) {
        if (col.rank() != 6)
            throw new IllegalArgumentException("col2im input array must be rank 6");

        INDArray output = Nd4j.create(col.dataType(), new long[]{col.size(0), col.size(1), kH, kW});

        val cfg = Conv2DConfig.builder()
                .sH(sH)
                .sW(sW)
                .dH(1)
                .dW(1)
                .kH(kH)
                .kW(kW)
                .pH(ph)
                .pW(pW)
                .build();

        Col2Im col2Im = Col2Im.builder()
                .inputArrays(new INDArray[]{col})
                .outputs(new INDArray[]{output})
                .conv2DConfig(cfg)
                .build();

        Nd4j.getExecutioner().execAndReturn(col2Im);
        return col2Im.outputArguments().get(0);
    }

    public static INDArray col2im(INDArray col, INDArray z, int sH, int sW, int pH, int pW, int kH, int kW,
                                  int dH, int dW) {
        if (col.rank() != 6)
            throw new IllegalArgumentException("col2im input array must be rank 6");
        if (z.rank() != 4)
            throw new IllegalArgumentException("col2im output array must be rank 4");
        Col2Im col2Im = Col2Im.builder()
                .inputArrays(new INDArray[]{col})
                .outputs(new INDArray[]{z})
                .conv2DConfig(Conv2DConfig.builder()
                        .sH(sH)
                        .sW(sW)
                        .dH(dH)
                        .dW(dW)
                        .kH(kH)
                        .kW(kW)
                        .pH(pH)
                        .pW(pW)
                        .build())
                .build();

        Nd4j.getExecutioner().execAndReturn(col2Im);

        return z;
    }

    /**
     * @param img
     * @param kernel
     * @param stride
     * @param padding
     * @return
     */
    public static INDArray im2col(INDArray img, int[] kernel, int[] stride, int[] padding) {
        Nd4j.getCompressor().autoDecompress(img);
        return im2col(img, kernel[0], kernel[1], stride[0], stride[1], padding[0], padding[1], 0, false);
    }

    /**
     * Implement column formatted images
     *
     * @param img        the image to process
     * @param kh         the kernel height
     * @param kw         the kernel width
     * @param sy         the stride along y
     * @param sx         the stride along x
     * @param ph         the padding width
     * @param pw         the padding height
     * @param isSameMode whether to cover the whole image or not
     * @return the column formatted image
     */
    public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, boolean isSameMode) {
        return im2col(img, kh, kw, sy, sx, ph, pw, 1, 1, isSameMode);
    }

    public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, int dh, int dw, boolean isSameMode) {
        Nd4j.getCompressor().autoDecompress(img);
        //Input: NCHW format
        long outH = outputSize(img.size(2), kh, sy, ph, dh, isSameMode);
        long outW = outputSize(img.size(3), kw, sx, pw, dw, isSameMode);

        //[miniBatch,depth,kH,kW,outH,outW]
        INDArray out = Nd4j.create(img.dataType(),new long[]{img.size(0), img.size(1), kh, kw, outH, outW}, 'c');

        return im2col(img, kh, kw, sy, sx, ph, pw, dh, dw, isSameMode, out);
    }

    public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, boolean isSameMode,
                                  INDArray out) {
        Im2col im2col = Im2col.builder()
                .outputs(new INDArray[]{out})
                .inputArrays(new INDArray[]{img})
                .conv2DConfig(Conv2DConfig.builder()
                        .kH(kh)
                        .pW(pw)
                        .pH(ph)
                        .sH(sy)
                        .sW(sx)
                        .kH(kh)
                        .kW(kw)
                        .dH(1)
                        .dW(1)
                        .paddingMode(isSameMode ? PaddingMode.SAME : PaddingMode.VALID)
                        .build()).build();

        Nd4j.getExecutioner().execAndReturn(im2col);
        return im2col.outputArguments().get(0);
    }

    /**
     * Execute im2col. Note the input must be NCHW.
     * @param img the input image in NCHW
     * @param kh
     * @param kw
     * @param sy
     * @param sx
     * @param ph
     * @param pw
     * @param dH
     * @param dW
     * @param isSameMode
     * @param out
     * @return
     */
    public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, int dH, int dW, boolean isSameMode,
                                  INDArray out) {

        Im2col im2col = Im2col.builder()
                .outputs(new INDArray[]{out})
                .inputArrays(new INDArray[]{img})
                .conv2DConfig(Conv2DConfig.builder()
                        .pW(pw)
                        .pH(ph)
                        .sH(sy)
                        .sW(sx)
                        .kW(kw)
                        .kH(kh)
                        .dW(dW)
                        .dH(dH)
                        .paddingMode(isSameMode ? PaddingMode.SAME : PaddingMode.VALID)
                        .build()).build();

        Nd4j.getExecutioner().execAndReturn(im2col);
        return im2col.outputArguments().get(0);
    }

    /**
     * Pooling 2d implementation
     *
     * @param img
     * @param kh
     * @param kw
     * @param sy
     * @param sx
     * @param ph
     * @param pw
     * @param dh
     * @param dw
     * @param isSameMode
     * @param type
     * @param extra         optional argument. I.e. used in pnorm pooling.
     * @param virtualHeight
     * @param virtualWidth
     * @param out
     * @return
     */
    public static INDArray pooling2D(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw,
                                     int dh, int dw, boolean isSameMode, Pooling2D.Pooling2DType type, Pooling2D.Divisor divisor,
                                     double extra, int virtualHeight, int virtualWidth, INDArray out) {
        Pooling2D pooling = new Pooling2D(img, out, Pooling2DConfig.builder()
                .dH(dh)
                .dW(dw)
                .extra(extra)
                .kH(kh)
                .kW(kw)
                .pH(ph)
                .pW(pw)
                .paddingMode(isSameMode ? PaddingMode.SAME : PaddingMode.VALID)
                .sH(sy)
                .sW(sx)
                .type(type)
                .divisor(divisor)
                .build());
        Nd4j.getExecutioner().execAndReturn(pooling);
        return out;
    }

    /**
     * Implement column formatted images
     *
     * @param img        the image to process
     * @param kh         the kernel height
     * @param kw         the kernel width
     * @param sy         the stride along y
     * @param sx         the stride along x
     * @param ph         the padding width
     * @param pw         the padding height
     * @param pval       the padding value (not used)
     * @param isSameMode whether padding mode is 'same'
     * @return the column formatted image
     */
    public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, int pval,
                                  boolean isSameMode) {
        INDArray output = null;

        if (isSameMode) {
            int oH = (int) Math.ceil(img.size(2) * 1.f / sy);
            int oW = (int) Math.ceil(img.size(3) * 1.f / sx);

            output = Nd4j.createUninitialized(img.dataType(), new long[]{img.size(0), img.size(1), kh, kw, oH, oW}, 'c');
        } else {
            long oH = (img.size(2) - (kh + (kh - 1) * (1 - 1)) + 2 * ph) / sy + 1;
            long oW = (img.size(3) - (kw + (kw - 1) * (1 - 1)) + 2 * pw) / sx + 1;

            output = Nd4j.valueArrayOf( new long[]{img.size(0), img.size(1), kh, kw, oH, oW}, pval, img.dataType());
        }

        Im2col im2col = Im2col.builder()
                .inputArrays(new INDArray[]{img})
                .outputs(new INDArray[]{output})
                .conv2DConfig(Conv2DConfig.builder()
                        .pW(pw)
                        .pH(ph)
                        .sH(sy)
                        .sW(sx)
                        .kW(kw)
                        .kH(kh)
                        .dW(1)
                        .dH(1)
                        .paddingMode(isSameMode ? PaddingMode.SAME : PaddingMode.VALID)
                        .build()).build();

        Nd4j.getExecutioner().execAndReturn(im2col);
        return im2col.outputArguments().get(0);
    }

    /**
     * The out size for a convolution
     *
     * @param size
     * @param k
     * @param s
     * @param p
     * @param coverAll
     * @return
     */
    @Deprecated
    public static long outSize(long size, long k, long s, long p, int dilation, boolean coverAll) {
        k = effectiveKernelSize(k, dilation);

        if (coverAll)
            return (size + p * 2 - k + s - 1) / s + 1;
        else
            return (size + p * 2 - k) / s + 1;
    }

    public static long outputSize(long size, long k, long s, long p, int dilation, boolean isSameMode) {
        k = effectiveKernelSize(k, dilation);

        if (isSameMode) {
            return (int) Math.ceil(size * 1.f / s);
        } else {
            return (size - k + 2 * p) / s + 1;
        }
    }

    public static long effectiveKernelSize(long kernel, int dilation) {
        return kernel + (kernel - 1) * (dilation - 1);
    }


    /**
     * 2d convolution (aka the last 2 dimensions
     *
     * @param input  the input to op
     * @param kernel the kernel to convolve with
     * @param type
     * @return
     */
    public static INDArray conv2d(INDArray input, INDArray kernel, Type type) {
        return Nd4j.getConvolution().conv2d(input, kernel, type);
    }

    /**
     * ND Convolution
     *
     * @param input  the input to op
     * @param kernel the kerrnel to op with
     * @param type   the opType of convolution
     * @param axes   the axes to do the convolution along
     * @return the convolution of the given input and kernel
     */
    public static INDArray convn(INDArray input, INDArray kernel, Type type, int[] axes) {
        return Nd4j.getConvolution().convn(input, kernel, type, axes);
    }

    /**
     * ND Convolution
     *
     * @param input  the input to op
     * @param kernel the kernel to op with
     * @param type   the opType of convolution
     * @return the convolution of the given input and kernel
     */
    public static INDArray convn(INDArray input, INDArray kernel, Type type) {
        return Nd4j.getConvolution().convn(input, kernel, type);
    }
}
