/* Copyright 2019-2021 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

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.
=======================================================================
*/
package org.tensorflow;

import static org.tensorflow.internal.c_api.global.tensorflow.TF_AddControlInput;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_AddInput;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_AddInputList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_FinishOperation;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_FinishOperationLocked;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewOperation;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewOperationLocked;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_OperationName;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrBool;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrBoolList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrFloat;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrFloatList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrFuncName;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrInt;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrIntList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrShape;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrShapeList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrString;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrStringList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrTensor;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrTensorList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrType;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrTypeList;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrValueProto;
import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetDevice;

import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;
import org.bytedeco.javacpp.BooleanPointer;
import org.bytedeco.javacpp.BytePointer;
import org.bytedeco.javacpp.IntPointer;
import org.bytedeco.javacpp.LongPointer;
import org.bytedeco.javacpp.Pointer;
import org.bytedeco.javacpp.PointerPointer;
import org.bytedeco.javacpp.PointerScope;
import org.bytedeco.javacpp.SizeTPointer;
import org.tensorflow.Graph.Reference;
import org.tensorflow.internal.c_api.TF_Graph;
import org.tensorflow.internal.c_api.TF_Operation;
import org.tensorflow.internal.c_api.TF_OperationDescription;
import org.tensorflow.internal.c_api.TF_Output;
import org.tensorflow.internal.c_api.TF_Status;
import org.tensorflow.internal.c_api.TF_Tensor;
import org.tensorflow.ndarray.Shape;
import org.tensorflow.op.Scope;
import org.tensorflow.proto.framework.AttrValue;
import org.tensorflow.proto.framework.AttrValue.ListValue;
import org.tensorflow.proto.framework.DataType;
import org.tensorflow.proto.framework.NameAttrList;

/** An {@link OperationBuilder} for adding {@link GraphOperation}s to a {@link Graph}. */
public final class GraphOperationBuilder implements OperationBuilder {

  GraphOperationBuilder(
      Graph graph, String type, String name, Scope scope, boolean dangerousGradientBuilder) {
    this.graph = graph;
    this.scope = scope;
    this.dangerousGradientBuilder = dangerousGradientBuilder;
    try (Graph.Reference r = graph.ref()) {
      if (dangerousGradientBuilder) {
        this.unsafeNativeHandle = allocateDangerousGradient(r.nativeHandle(), type, name);
      } else {
        this.unsafeNativeHandle = allocate(r.nativeHandle(), type, name);
      }
    }
  }

  /**
   * Add the {@link GraphOperation} being built to the {@link Graph}.
   *
   * <p>The OperationBuilder is not usable after build() returns.
   */
  @Override
  public GraphOperation build() {
    scope.apply(this);
    try (Graph.Reference r = graph.ref()) {
      TF_Operation built;
      if (dangerousGradientBuilder) {
        built = finishDangerousGradient(r.nativeHandle(), unsafeNativeHandle);
      } else {
        built = finish(unsafeNativeHandle);
      }
      GraphOperation op = new GraphOperation(graph, built);
      unsafeNativeHandle = null;
      scope.onOpCreated(op);
      return op;
    }
  }

  @Override
  public GraphOperationBuilder addControlInput(Operation control) {
    graph.checkInput(control);

    if (control.env() != graph) {
      throw new IllegalArgumentException(
          "Control input " + control + " was from a different graph, can't use.");
    }

    Graph.Reference r = graph.ref();
    try {
      addControlInput(unsafeNativeHandle, ((GraphOperation) control).getUnsafeNativeHandle());
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder addInput(Output<?> input) {
    graph.checkInput(input);
    Graph.Reference r = graph.ref();
    try {
      addInput(unsafeNativeHandle, (TF_Operation) input.getUnsafeNativeHandle(), input.index());
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder addInputList(Output<?>[] inputs) {
    for (Output<?> input : inputs) {
      graph.checkInput(input);
    }

    Graph.Reference r = graph.ref();
    try {
      TF_Operation[] opHandles = new TF_Operation[inputs.length];
      int[] indices = new int[inputs.length];
      for (int i = 0; i < inputs.length; ++i) {
        opHandles[i] = (TF_Operation) inputs[i].getUnsafeNativeHandle();
        indices[i] = inputs[i].index();
      }
      addInputList(unsafeNativeHandle, opHandles, indices);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setDevice(String device) {
    Graph.Reference r = graph.ref();
    try {
      setDevice(unsafeNativeHandle, device);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, String value) {
    setAttr(name, value.getBytes(Charset.forName("UTF-8")));
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, byte[] value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrString(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, long value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrInt(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, long[] value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrIntList(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, float value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrFloat(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, float[] value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrFloatList(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, boolean value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrBool(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, boolean[] value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrBoolList(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, DataType value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrType(unsafeNativeHandle, name, value.getNumber());
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, DataType[] value) {
    int[] ctypes = new int[value.length];
    for (int i = 0; i < value.length; ++i) {
      ctypes[i] = value[i].getNumber();
    }
    Graph.Reference r = graph.ref();
    try {
      setAttrTypeList(unsafeNativeHandle, name, ctypes);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, Tensor value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrTensor(unsafeNativeHandle, name, value.asRawTensor().nativeHandle());
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, Tensor[] value) {
    TF_Tensor[] handles = new TF_Tensor[value.length];
    int idx = 0;
    for (Tensor t : value) {
      handles[idx++] = t.asRawTensor().nativeHandle();
    }
    Graph.Reference r = graph.ref();
    try {
      setAttrTensorList(unsafeNativeHandle, name, handles);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, Shape value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrShape(unsafeNativeHandle, name, value.asArray(), value.numDimensions());
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, Shape[] value) {
    int[] numDimensions = new int[value.length];
    int totalNumDimensions = 0;
    for (int idx = 0; idx < value.length; ++idx) {
      int n = value[idx].numDimensions();
      numDimensions[idx] = n;
      if (n > 0) {
        totalNumDimensions += n;
      }
    }
    // Flatten the shapes into a single array to avoid too much overhead in the
    // native part
    long[] shapes = new long[totalNumDimensions];
    int shapeIdx = 0;
    for (Shape shape : value) {
      if (shape.numDimensions() > 0) {
        for (long dim : shape.asArray()) {
          shapes[shapeIdx++] = dim;
        }
      }
    }
    Graph.Reference r = graph.ref();
    try {
      setAttrShapeList(unsafeNativeHandle, name, shapes, numDimensions);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public GraphOperationBuilder setAttr(String name, String[] value) {
    Charset utf8 = Charset.forName("UTF-8");
    byte[][] objects = new byte[value.length][];
    for (int i = 0; i < value.length; ++i) {
      objects[i] = value[i].getBytes(utf8);
    }
    Graph.Reference r = graph.ref();
    try {
      setAttrStringList(unsafeNativeHandle, name, objects);
    } finally {
      r.close();
    }
    return this;
  }

  @Override
  public OperationBuilder setAttr(String name, ConcreteFunction value) {
    graph.attachFunction(value);
    try (Reference r = graph.ref()) {
      setAttrFunctionName(unsafeNativeHandle, name, value.getDefinedName());
    }
    return this;
  }

  @Override
  public OperationBuilder setAttr(String name, ConcreteFunction[] value) {
    for (ConcreteFunction f : value) {
      graph.attachFunction(f);
    }

    try (Reference r = graph.ref()) {
      setAttrFunctionList(
          unsafeNativeHandle,
          name,
          Arrays.stream(value).map(ConcreteFunction::getDefinedName).collect(Collectors.toList()));
    }
    return this;
  }

  @Override
  public OperationBuilder setAttr(String name, AttrValue value) {
    Graph.Reference r = graph.ref();
    try {
      setAttrValue(unsafeNativeHandle, name, value);
    } finally {
      r.close();
    }
    return this;
  }

  private TF_OperationDescription unsafeNativeHandle;
  private final Graph graph;
  private final Scope scope;

  /**
   * Use builders without locking. This should only be used during custom gradient building.
   *
   * <p>The graph locks are not re-entrant, so attempting to add an op to a graph that has been
   * locked by the gradient builder will fail without this.
   */
  private final boolean dangerousGradientBuilder;

  private static void requireHandle(Pointer handle) {
    if (handle == null || handle.isNull()) {
      throw new IllegalStateException("Operation has already been built");
    }
  }

  private static void resolveOutput(TF_Operation opHandle, int index, TF_Output out) {
    if (opHandle == null || opHandle.isNull()) {
      throw new IllegalStateException("close() was called on the Graph");
    }
    out.oper(opHandle);
    out.index(index);
  }

  private static void requireTensor(TF_Tensor handle) {
    if (handle == null || handle.isNull()) {
      throw new IllegalStateException("close() has been called on the Tensor");
    }
  }

  private static TF_OperationDescription allocate(TF_Graph graphHandle, String type, String name) {
    if (graphHandle == null || graphHandle.isNull()) {
      throw new IllegalStateException("close() has been called on the Graph");
    }
    return TF_NewOperation(graphHandle, type, name);
  }

  /**
   * Use builders without locking. This should only be used during custom gradient building.
   *
   * <p>The graph locks are not re-entrant, so attempting to add an op to a graph that has been
   * locked by the gradient builder will fail without this.
   */
  private static TF_OperationDescription allocateDangerousGradient(
      TF_Graph graphHandle, String type, String name) {
    if (graphHandle == null || graphHandle.isNull()) {
      throw new IllegalStateException("close() has been called on the Graph");
    }
    return TF_NewOperationLocked(graphHandle, type, name);
  }

  private static TF_Operation finish(TF_OperationDescription handle) {
    requireHandle(handle);

    try (PointerScope scope = new PointerScope()) {
      TF_Status status = TF_Status.newStatus();
      TF_Operation op = TF_FinishOperation(handle, status);
      status.throwExceptionIfNotOK();
      return op;
    }
  }

  /**
   * Use builders without locking. This should only be used during custom gradient building.
   *
   * <p>The graph locks are not re-entrant, so attempting to add an op to a graph that has been
   * locked by the gradient builder will fail without this.
   */
  private static TF_Operation finishDangerousGradient(TF_Graph g, TF_OperationDescription handle) {
    requireHandle(handle);

    try (PointerScope scope = new PointerScope()) {
      TF_Status status = TF_Status.newStatus();
      TF_Operation op = TF_FinishOperationLocked(handle, status);
      status.throwExceptionIfNotOK();
      g.name_map().erase(TF_OperationName(op));
      return op;
    }
  }

  private static void addInput(TF_OperationDescription handle, TF_Operation opHandle, int index) {
    try (PointerScope scope = new PointerScope()) {
      TF_Output out = new TF_Output();
      resolveOutput(opHandle, index, out);
      requireHandle(handle);
      TF_AddInput(handle, out);
    }
  }

  private static void addInputList(
      TF_OperationDescription handle, TF_Operation[] opHandles, int[] indices) {
    requireHandle(handle);
    if (indices.length != opHandles.length) {
      throw new IllegalArgumentException(
          "mismatch in number of Operations ("
              + opHandles.length
              + ") and output indices ("
              + indices.length
              + ") provided");
    }

    try (PointerScope scope = new PointerScope()) {
      TF_Output o = new TF_Output(indices.length);
      for (int i = 0; i < indices.length; ++i) {
        resolveOutput(opHandles[i], indices[i], o.position(i));
      }
      TF_AddInputList(handle, o.position(0), indices.length);
    }
  }

  private static void addControlInput(TF_OperationDescription handle, TF_Operation opHandle) {
    if (opHandle == null || opHandle.isNull()) {
      throw new IllegalStateException(
          "control input is not valid, " + "perhaps the Graph containing it has been closed()?");
    }
    requireHandle(handle);
    TF_AddControlInput(handle, opHandle);
  }

  private static void setDevice(TF_OperationDescription handle, String device) {
    requireHandle(handle);
    TF_SetDevice(handle, device);
  }

  // The names of all the setAttr* family functions below correspond to the C library types, not the
  // Java library types. Roughly, setAttrFoo calls the TensorFlow C library function: TF_SetAttrFoo.

  private static void setAttrString(TF_OperationDescription handle, String name, byte[] value) {
    requireHandle(handle);
    try (PointerScope scope = new PointerScope()) {
      TF_SetAttrString(handle, name, new BytePointer(value), value.length);
    }
  }

  private static void setAttrInt(TF_OperationDescription handle, String name, long value) {
    requireHandle(handle);
    TF_SetAttrInt(handle, name, value);
  }

  private static void setAttrIntList(TF_OperationDescription handle, String name, long[] value) {
    requireHandle(handle);
    TF_SetAttrIntList(handle, name, value, value.length);
  }

  private static void setAttrFloat(TF_OperationDescription handle, String name, float value) {
    requireHandle(handle);
    TF_SetAttrFloat(handle, name, value);
  }

  private static void setAttrFloatList(TF_OperationDescription handle, String name, float[] value) {
    requireHandle(handle);
    TF_SetAttrFloatList(handle, name, value, value.length);
  }

  private static void setAttrBool(TF_OperationDescription handle, String name, boolean value) {
    requireHandle(handle);
    TF_SetAttrBool(handle, name, (byte) (value ? 1 : 0));
  }

  private static void setAttrBoolList(
      TF_OperationDescription handle, String name, boolean[] value) {
    requireHandle(handle);
    try (PointerScope scope = new PointerScope()) {
      TF_SetAttrBoolList(handle, name, new BytePointer(new BooleanPointer(value)), value.length);
    }
  }

  private static void setAttrType(TF_OperationDescription handle, String name, int type) {
    requireHandle(handle);
    TF_SetAttrType(handle, name, type);
  }

  private static void setAttrTypeList(TF_OperationDescription handle, String name, int[] type) {
    requireHandle(handle);
    TF_SetAttrTypeList(handle, name, type, type.length);
  }

  private static void setAttrTensor(
      TF_OperationDescription handle, String name, TF_Tensor tensorHandle) {
    requireHandle(handle);
    requireTensor(tensorHandle);

    try (PointerScope scope = new PointerScope()) {
      TF_Status status = TF_Status.newStatus();
      TF_SetAttrTensor(handle, name, tensorHandle, status);
      status.throwExceptionIfNotOK();
    }
  }

  private static void setAttrTensorList(
      TF_OperationDescription handle, String name, TF_Tensor[] tensorHandles) {
    requireHandle(handle);

    try (PointerScope scope = new PointerScope()) {
      PointerPointer<TF_Tensor> tensors = new PointerPointer<TF_Tensor>(tensorHandles.length);
      for (int i = 0; i < tensorHandles.length; ++i) {
        requireTensor(tensorHandles[i]);
        tensors.put(i, tensorHandles[i]);
      }

      TF_Status status = TF_Status.newStatus();
      TF_SetAttrTensorList(
          handle, new BytePointer(name), tensors.position(0), tensorHandles.length, status);
      status.throwExceptionIfNotOK();
    }
  }

  private static void setAttrShape(
      TF_OperationDescription handle, String name, long[] shape, int numDims) {
    requireHandle(handle);

    // num_dims and env->GetArrayLength(shape) are assumed to be consistent.
    // i.e., either num_dims < 0 or num_dims == env->GetArrayLength(shape).
    TF_SetAttrShape(handle, name, shape, numDims);
  }

  private static void setAttrShapeList(
      TF_OperationDescription handle, String name, long[] shapes, int[] numDims) {
    requireHandle(handle);

    try (PointerScope scope = new PointerScope()) {
      LongPointer shapesPointer = new LongPointer(shapes);
      PointerPointer<LongPointer> shapesPointers = new PointerPointer<LongPointer>(numDims.length);
      for (int i = 0; i < numDims.length; i++) {
        shapesPointers.put(i, shapesPointer);
        shapesPointer.position(shapesPointer.position() + numDims[i] * 8);
      }
      TF_SetAttrShapeList(
          handle, new BytePointer(name), shapesPointers, new IntPointer(numDims), numDims.length);
    }
  }

  private static void setAttrStringList(
      TF_OperationDescription handle, String name, byte[][] value) {
    requireHandle(handle);

    try (PointerScope scope = new PointerScope()) {
      PointerPointer<BytePointer> valuePointers = new PointerPointer<BytePointer>(value.length);
      SizeTPointer lengths = new SizeTPointer(value.length);

      for (int i = 0; i < value.length; ++i) {
        valuePointers.put(i, new BytePointer(value[i]));
        lengths.put(i, value[i].length);
      }
      TF_SetAttrStringList(handle, new BytePointer(name), valuePointers, lengths, value.length);
    }
  }

  private static void setAttrFunctionName(
      TF_OperationDescription opHandle, String attrName, String functionName) {
    requireHandle(opHandle);
    try (PointerScope scope = new PointerScope()) {
      TF_SetAttrFuncName(opHandle, attrName, functionName, functionName.length());
    }
  }

  private static void setAttrFunctionList(
      TF_OperationDescription opHandle, String attrName, List<String> functionNames) {
    AttrValue value =
        AttrValue.newBuilder()
            .setList(
                ListValue.newBuilder()
                    .addAllFunc(
                        functionNames.stream()
                            .map(x -> NameAttrList.newBuilder().setName(x).build())
                            .collect(Collectors.toList()))
                    .build())
            .build();
    setAttrValue(opHandle, attrName, value);
  }

  private static void setAttrValue(
      TF_OperationDescription opHandle, String attrName, AttrValue value) {
    requireHandle(opHandle);
    try (PointerScope scope = new PointerScope()) {
      TF_Status status = TF_Status.newStatus();
      byte[] bytes = value.toByteArray();
      TF_SetAttrValueProto(opHandle, attrName, new BytePointer(bytes), bytes.length, status);
      status.throwExceptionIfNotOK();
    }
  }
}
