/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.apache.arrow.vector;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.vector.compare.ApproxEqualsVisitor;
import org.apache.arrow.vector.compare.Range;
import org.apache.arrow.vector.compare.VectorEqualsVisitor;
import org.apache.arrow.vector.compare.VectorValueEqualizer;
import org.apache.arrow.vector.compare.util.ValueEpsilonEqualizers;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.apache.arrow.vector.util.TransferPair;

/**
 * Holder for a set of vectors to be loaded/unloaded. A VectorSchemaRoot is a container that can
 * hold batches, batches flow through VectorSchemaRoot as part of a pipeline. Note this is different
 * from other implementations (i.e. in C++ and Python, a RecordBatch is a collection of equal-length
 * vector instances and was created each time for a new batch).
 *
 * <p>The recommended usage for VectorSchemaRoot is creating a single VectorSchemaRoot based on the
 * known schema and populated data over and over into the same VectorSchemaRoot in a stream of
 * batches rather than create a new VectorSchemaRoot instance each time (see Flight or
 * ArrowFileWriter for better understanding). Thus at any one point a VectorSchemaRoot may have data
 * or may have no data (say it was transferred downstream or not yet populated).
 */
public class VectorSchemaRoot implements AutoCloseable {

  private Schema schema;
  private int rowCount;
  private final List<FieldVector> fieldVectors;
  private final Map<Field, FieldVector> fieldVectorsMap = new LinkedHashMap<>();

  /** Constructs new instance containing each of the vectors. */
  public VectorSchemaRoot(Iterable<FieldVector> vectors) {
    this(
        StreamSupport.stream(vectors.spliterator(), false)
            .map(t -> t.getField())
            .collect(Collectors.toList()),
        StreamSupport.stream(vectors.spliterator(), false).collect(Collectors.toList()));
  }

  /** Constructs a new instance containing the children of parent but not the parent itself. */
  public VectorSchemaRoot(FieldVector parent) {
    this(parent.getField().getChildren(), parent.getChildrenFromFields(), parent.getValueCount());
  }

  /**
   * Constructs a new instance.
   *
   * @param fields The types of each vector.
   * @param fieldVectors The data vectors (must be equal in size to <code>fields</code>.
   */
  public VectorSchemaRoot(List<Field> fields, List<FieldVector> fieldVectors) {
    this(
        new Schema(fields),
        fieldVectors,
        fieldVectors.size() == 0 ? 0 : fieldVectors.get(0).getValueCount());
  }

  /**
   * Constructs a new instance.
   *
   * @param fields The types of each vector.
   * @param fieldVectors The data vectors (must be equal in size to <code>fields</code>.
   * @param rowCount The number of rows contained.
   */
  public VectorSchemaRoot(List<Field> fields, List<FieldVector> fieldVectors, int rowCount) {
    this(new Schema(fields), fieldVectors, rowCount);
  }

  /**
   * Constructs a new instance.
   *
   * @param schema The schema for the vectors.
   * @param fieldVectors The data vectors.
   * @param rowCount The number of rows
   */
  public VectorSchemaRoot(Schema schema, List<FieldVector> fieldVectors, int rowCount) {
    if (schema.getFields().size() != fieldVectors.size()) {
      throw new IllegalArgumentException(
          "Fields must match field vectors. Found "
              + fieldVectors.size()
              + " vectors and "
              + schema.getFields().size()
              + " fields");
    }
    this.schema = schema;
    this.rowCount = rowCount;
    this.fieldVectors = fieldVectors;
    for (int i = 0; i < schema.getFields().size(); ++i) {
      Field field = schema.getFields().get(i);
      FieldVector vector = fieldVectors.get(i);
      fieldVectorsMap.put(field, vector);
    }
  }

  /** Creates a new set of empty vectors corresponding to the given schema. */
  public static VectorSchemaRoot create(Schema schema, BufferAllocator allocator) {
    List<FieldVector> fieldVectors = new ArrayList<>(schema.getFields().size());
    for (Field field : schema.getFields()) {
      FieldVector vector = field.createVector(allocator);
      fieldVectors.add(vector);
    }
    if (fieldVectors.size() != schema.getFields().size()) {
      throw new IllegalArgumentException(
          "The root vector did not create the right number of children. found "
              + fieldVectors.size()
              + " expected "
              + schema.getFields().size());
    }
    return new VectorSchemaRoot(schema, fieldVectors, 0);
  }

  /** Constructs a new instance from vectors. */
  public static VectorSchemaRoot of(FieldVector... vectors) {
    return new VectorSchemaRoot(Arrays.stream(vectors).collect(Collectors.toList()));
  }

  /**
   * Do an adaptive allocation of each vector for memory purposes. Sizes will be based on previously
   * defined initial allocation for each vector (and subsequent size learned).
   */
  public void allocateNew() {
    for (FieldVector v : fieldVectors) {
      v.allocateNew();
    }
    rowCount = 0;
  }

  /**
   * Release all the memory for each vector held in this root. This DOES NOT remove vectors from the
   * container.
   */
  public void clear() {
    for (FieldVector v : fieldVectors) {
      v.clear();
    }
    rowCount = 0;
  }

  public List<FieldVector> getFieldVectors() {
    return Collections.unmodifiableList(fieldVectors);
  }

  /**
   * gets a vector by name.
   *
   * <p>if name occurs multiple times this returns the first inserted entry for name
   */
  public FieldVector getVector(String name) {
    for (Map.Entry<Field, FieldVector> entry : fieldVectorsMap.entrySet()) {
      if (entry.getKey().getName().equals(name)) {
        return entry.getValue();
      }
    }
    return null;
  }

  public FieldVector getVector(Field field) {
    return fieldVectorsMap.get(field);
  }

  public FieldVector getVector(int index) {
    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
    return fieldVectors.get(index);
  }

  /**
   * Add vector to the record batch, producing a new VectorSchemaRoot.
   *
   * @param index field index
   * @param vector vector to be added.
   * @return out VectorSchemaRoot with vector added
   */
  public VectorSchemaRoot addVector(int index, FieldVector vector) {
    Preconditions.checkNotNull(vector);
    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
    List<FieldVector> newVectors = new ArrayList<>();
    for (int i = 0; i < fieldVectors.size(); i++) {
      if (i == index) {
        newVectors.add(vector);
      }
      newVectors.add(fieldVectors.get(i));
    }
    return new VectorSchemaRoot(newVectors);
  }

  /**
   * Remove vector from the record batch, producing a new VectorSchemaRoot.
   *
   * @param index field index
   * @return out VectorSchemaRoot with vector removed
   */
  public VectorSchemaRoot removeVector(int index) {
    Preconditions.checkArgument(index >= 0 && index < fieldVectors.size());
    List<FieldVector> newVectors = new ArrayList<>();
    for (int i = 0; i < fieldVectors.size(); i++) {
      if (i != index) {
        newVectors.add(fieldVectors.get(i));
      }
    }
    return new VectorSchemaRoot(newVectors);
  }

  public Schema getSchema() {
    return schema;
  }

  public int getRowCount() {
    return rowCount;
  }

  /**
   * Set the row count of all the vectors in this container. Also sets the value count for each root
   * level contained FieldVector.
   *
   * @param rowCount Number of records.
   */
  public void setRowCount(int rowCount) {
    this.rowCount = rowCount;
    for (FieldVector v : fieldVectors) {
      v.setValueCount(rowCount);
    }
  }

  @Override
  public void close() {
    try {
      AutoCloseables.close(fieldVectors);
    } catch (RuntimeException ex) {
      throw ex;
    } catch (Exception ex) {
      // should never happen since FieldVector.close() doesn't throw IOException
      throw new RuntimeException(ex);
    }
  }

  private void printRow(StringBuilder sb, List<Object> row) {
    boolean first = true;
    for (Object v : row) {
      if (first) {
        first = false;
      } else {
        sb.append("\t");
      }
      sb.append(v);
    }
    sb.append("\n");
  }

  /** Returns a tab separated value of vectors (based on their java object representation). */
  public String contentToTSVString() {
    StringBuilder sb = new StringBuilder();
    List<Object> row = new ArrayList<>(schema.getFields().size());
    for (Field field : schema.getFields()) {
      row.add(field.getName());
    }
    printRow(sb, row);
    for (int i = 0; i < rowCount; i++) {
      row.clear();
      for (FieldVector v : fieldVectors) {
        row.add(v.getObject(i));
      }
      printRow(sb, row);
    }
    return sb.toString();
  }

  /**
   * Synchronizes the schema from the current vectors. In some cases, the schema and the actual
   * vector structure may be different. This can be caused by a promoted writer (For details, please
   * see {@link org.apache.arrow.vector.complex.impl.PromotableWriter}). For example, when writing
   * different types of data to a {@link org.apache.arrow.vector.complex.ListVector} may lead to
   * such a case. When this happens, this method should be called to bring the schema and vector
   * structure in a synchronized state.
   *
   * @return true if the schema is updated, false otherwise.
   */
  public boolean syncSchema() {
    List<Field> oldFields = this.schema.getFields();
    List<Field> newFields =
        this.fieldVectors.stream().map(ValueVector::getField).collect(Collectors.toList());
    if (!oldFields.equals(newFields)) {
      this.schema = new Schema(newFields);
      return true;
    }
    return false;
  }

  /**
   * Slice this root from desired index.
   *
   * @param index start position of the slice
   * @return the sliced root
   */
  public VectorSchemaRoot slice(int index) {
    return slice(index, this.rowCount - index);
  }

  /**
   * Slice this root at desired index and length.
   *
   * @param index start position of the slice
   * @param length length of the slice
   * @return the sliced root
   */
  public VectorSchemaRoot slice(int index, int length) {
    Preconditions.checkArgument(index >= 0, "expecting non-negative index");
    Preconditions.checkArgument(length >= 0, "expecting non-negative length");
    Preconditions.checkArgument(index + length <= rowCount, "index + length should <= rowCount");

    List<FieldVector> sliceVectors =
        fieldVectors.stream()
            .map(
                v -> {
                  TransferPair transferPair = v.getTransferPair(v.getAllocator());
                  transferPair.splitAndTransfer(index, length);
                  return (FieldVector) transferPair.getTo();
                })
            .collect(Collectors.toList());

    return new VectorSchemaRoot(sliceVectors);
  }

  /** Determine if two VectorSchemaRoots are exactly equal. */
  public boolean equals(VectorSchemaRoot other) {
    if (other == null) {
      return false;
    }

    if (!this.schema.equals(other.schema)) {
      return false;
    }

    if (this.rowCount != other.rowCount) {
      return false;
    }

    for (int i = 0; i < fieldVectors.size(); i++) {
      FieldVector vector = fieldVectors.get(i);
      FieldVector otherVector = other.fieldVectors.get(i);
      if (!VectorEqualsVisitor.vectorEquals(vector, otherVector)) {
        return false;
      }
    }

    return true;
  }

  /**
   * Determine if two VectorSchemaRoots are approximately equal using the given functions to
   * calculate difference between float/double values. Note that approx equals are in regards to
   * floating point values, other values are comparing to exactly equals.
   *
   * @param floatDiffFunction function to calculate difference between float values.
   * @param doubleDiffFunction function to calculate difference between double values.
   */
  public boolean approxEquals(
      VectorSchemaRoot other,
      VectorValueEqualizer<Float4Vector> floatDiffFunction,
      VectorValueEqualizer<Float8Vector> doubleDiffFunction) {

    Preconditions.checkNotNull(floatDiffFunction);
    Preconditions.checkNotNull(doubleDiffFunction);

    if (other == null) {
      return false;
    }

    if (!this.schema.equals(other.schema)) {
      return false;
    }

    if (this.rowCount != other.rowCount) {
      return false;
    }

    Range range = new Range(0, 0, 0);
    for (int i = 0; i < fieldVectors.size(); i++) {
      FieldVector vector = fieldVectors.get(i);
      FieldVector otherVector = other.fieldVectors.get(i);
      if (vector.getValueCount() != otherVector.getValueCount()) {
        return false;
      }
      ApproxEqualsVisitor visitor =
          new ApproxEqualsVisitor(vector, otherVector, floatDiffFunction, doubleDiffFunction);
      range.setLength(vector.getValueCount());
      if (!visitor.rangeEquals(range)) {
        return false;
      }
    }

    return true;
  }

  /**
   * Determine if two VectorSchemaRoots are approximately equal using default functions to calculate
   * difference between float/double values.
   */
  public boolean approxEquals(VectorSchemaRoot other) {
    VectorValueEqualizer<Float4Vector> floatDiffFunction =
        new ValueEpsilonEqualizers.Float4EpsilonEqualizer(
            ApproxEqualsVisitor.DEFAULT_FLOAT_EPSILON);
    VectorValueEqualizer<Float8Vector> doubleDiffFunction =
        new ValueEpsilonEqualizers.Float8EpsilonEqualizer(
            ApproxEqualsVisitor.DEFAULT_DOUBLE_EPSILON);
    return approxEquals(other, floatDiffFunction, doubleDiffFunction);
  }
}
