// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.rankingexpression.importer.operations;

import ai.vespa.rankingexpression.importer.OrderedTensorType;
import ai.vespa.rankingexpression.importer.DimensionRenamer;
import com.yahoo.tensor.Tensor;
import com.yahoo.tensor.TensorType;
import com.yahoo.tensor.functions.TensorFunction;

import java.util.List;
import java.util.Optional;

public class ConcatV2 extends IntermediateOperation {

    private String concatDimensionName;

    public ConcatV2(String modelName, String nodeName, List<IntermediateOperation> inputs) {
        super(modelName, nodeName, inputs);
    }

    @Override
    protected OrderedTensorType lazyGetType() {
        if (!inputs.stream().map(IntermediateOperation::type).allMatch(Optional::isPresent)) {
            return null;
        }

        IntermediateOperation concatDimOp = inputs.get(inputs.size() - 1);  // ConcatV2: concat dimension is the last input
        if (!concatDimOp.getConstantValue().isPresent()) {
            throw new IllegalArgumentException("ConcatV2 in " + name + ": " +
                                               "concat dimension must be a constant.");
        }
        Tensor concatDimTensor = concatDimOp.getConstantValue().get().asTensor();
        if (concatDimTensor.type().rank() != 0) {
            throw new IllegalArgumentException("ConcatV2 in " + name + ": " +
                                               "concat dimension must be a scalar.");
        }

        OrderedTensorType aType = inputs.get(0).type().get();

        int concatDim = (int)concatDimTensor.asDouble();
        long concatDimSize = aType.dimensions().get(concatDim).size().orElse(-1L);

        for (int i = 1; i < inputs.size() - 1; ++i) {
            OrderedTensorType bType = inputs.get(i).type().get();
            if (bType.rank() != aType.rank()) {
                throw new IllegalArgumentException("ConcatV2 in " + name + ": " +
                                                   "inputs must have save rank.");
            }
            for (int j = 0; j < aType.rank(); ++j) {
                long dimSizeA = aType.dimensions().get(j).size().orElse(-1L);
                long dimSizeB = bType.dimensions().get(j).size().orElse(-1L);
                if (j == concatDim) {
                    concatDimSize += dimSizeB;
                } else if (dimSizeA != dimSizeB) {
                    throw new IllegalArgumentException("ConcatV2 in " + name + ": " +
                                                       "input dimension " + j + " differs in input tensors.");
                }
            }
        }

        OrderedTensorType.Builder typeBuilder = new OrderedTensorType.Builder();
        int dimensionIndex = 0;
        for (TensorType.Dimension dimension : aType.dimensions()) {
            if (dimensionIndex == concatDim) {
                concatDimensionName = dimension.name();
                typeBuilder.add(TensorType.Dimension.indexed(concatDimensionName, concatDimSize));
            } else {
                typeBuilder.add(dimension);
            }
            dimensionIndex++;
        }
        return typeBuilder.build();
    }

    @Override
    protected TensorFunction lazyGetFunction() {
        if (!inputs.stream().map(IntermediateOperation::function).allMatch(Optional::isPresent)) {
            return null;
        }
        TensorFunction result = inputs.get(0).function().get();
        for (int i = 1; i < inputs.size() - 1; ++i) {
            TensorFunction b = inputs.get(i).function().get();
            result = new com.yahoo.tensor.functions.Concat(result, b, concatDimensionName);
        }
        return result;
    }

    @Override
    public void addDimensionNameConstraints(DimensionRenamer renamer) {
        if (!inputs.stream().map(IntermediateOperation::type).allMatch(Optional::isPresent)) {
            return;
        }
        OrderedTensorType a = inputs.get(0).type().get();
        for (int i = 1; i < inputs.size() - 1; ++i) {
            OrderedTensorType b = inputs.get(i).type().get();
            String bDim = b.dimensions().get(i).name();
            String aDim = a.dimensions().get(i).name();
            renamer.addConstraint(aDim, bDim, DimensionRenamer::equals, this);
        }
    }

    @Override
    public void renameDimensions(DimensionRenamer renamer) {
        super.renameDimensions(renamer);
        concatDimensionName = renamer.dimensionNameOf(concatDimensionName).orElse(concatDimensionName);
   }

}

@@ -556,16 +556,54 @@
 ionName;
+%0A    private int concatDimensionIndex;
 %0A%0A    pu
@@ -1661,25 +1661,24 @@
 .get();%0A
-%0A
         
 int conc
@@ -1669,20 +1669,16 @@
         
-int 
 concatDi
@@ -1678,16 +1678,27 @@
 oncatDim
+ensionIndex
  = (int)
@@ -1786,16 +1786,27 @@
 oncatDim
+ensionIndex
 ).size()
@@ -2402,32 +2402,43 @@
  (j == concatDim
+ensionIndex
 ) %7B%0A            
@@ -2987,16 +2987,27 @@
 oncatDim
+ensionIndex
 ) %7B%0A    
@@ -4251,33 +4251,52 @@
 imensions().get(
-i
+concatDimensionIndex
 ).name();%0A      
@@ -4334,17 +4334,36 @@
 s().get(
-i
+concatDimensionIndex
 ).name()

Set correct concat dimension for rename