/*
 * 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 tech.tablesaw.io.string;

import java.io.IOException;
import java.io.OutputStream;
import java.util.stream.IntStream;
import tech.tablesaw.columns.Column;
import tech.tablesaw.table.Relation;
import tech.tablesaw.util.StringUtils;

/**
 * A class that can pretty print a DataFrame to text for visualization in a console
 *
 * <p>Based off of
 * https://github.com/zavtech/morpheus-core/blob/master/src/main/java/com/zavtech/morpheus/reference/XDataFramePrinter.java
 * under Apache 2 license
 */
public class DataFramePrinter {

  private static final String TOO_SHORT_COLUMN_MARKER = "?";

  private final int maxRows;
  private final OutputStream stream;

  /**
   * Constructor
   *
   * @param maxRows the max rows to print
   * @param stream the print stream to write to
   */
  public DataFramePrinter(int maxRows, OutputStream stream) {
    this.maxRows = maxRows;
    this.stream = stream;
  }

  /**
   * Returns the column widths required to print the header and data
   *
   * @param headers the headers to print
   * @param data the data items to print
   * @return the required column widths
   */
  private static int[] getWidths(String[] headers, String[][] data) {
    final int[] widths = new int[headers.length];
    for (int j = 0; j < headers.length; j++) {
      final String header = headers[j];
      widths[j] = Math.max(widths[j], header != null ? header.length() : 0);
    }
    for (String[] rowValues : data) {
      for (int j = 0; j < rowValues.length; j++) {
        final String value = rowValues[j];
        widths[j] = Math.max(widths[j], value != null ? value.length() : 0);
      }
    }
    return widths;
  }

  /**
   * Returns the header template given the widths specified
   *
   * @param widths the token widths
   * @return the line format template
   */
  private static String getHeaderTemplate(int[] widths, String[] headers) {
    return IntStream.range(0, widths.length)
        .mapToObj(
            i -> {
              final int width = widths[i];
              final int length = headers[i].length();
              final int leading = (width - length) / 2;
              final int trailing = width - (length + leading);
              final StringBuilder text = new StringBuilder();
              whitespace(text, leading + 1);
              text.append("%").append(i + 1).append("$s");
              whitespace(text, trailing);
              text.append("  |");
              return text.toString();
            })
        .reduce((left, right) -> left + " " + right)
        .orElse("");
  }

  /**
   * Returns the data template given the widths specified
   *
   * @param widths the token widths
   * @return the line format template
   */
  private static String getDataTemplate(int[] widths) {
    return IntStream.range(0, widths.length)
        .mapToObj(i -> " %" + (i + 1) + "$" + widths[i] + "s  |")
        .reduce((left, right) -> left + " " + right)
        .orElse("");
  }

  /**
   * Returns a whitespace string of the length specified
   *
   * @param length the length for whitespace
   */
  private static void whitespace(StringBuilder text, int length) {
    IntStream.range(0, length).forEach(i -> text.append(" "));
  }

  /**
   * Prints the specified DataFrame to the stream bound to this printer
   *
   * @param frame the DataFrame to print
   */
  public void print(Relation frame) {
    try {
      final String[] headers = getHeaderTokens(frame);
      final String[][] data = getDataTokens(frame);
      final int[] widths = getWidths(headers, data);
      final String dataTemplate = getDataTemplate(widths);
      final String headerTemplate = getHeaderTemplate(widths, headers);
      final int totalWidth = IntStream.of(widths).map(w -> w + 5).sum() - 1;
      final int totalHeight = data.length + 1;
      int capacity = totalWidth * totalHeight;
      if (capacity < 0) {
        capacity = 0;
      }
      final StringBuilder text = new StringBuilder(capacity);
      if (frame.name() != null) {
        text.append(tableName(frame, totalWidth)).append(System.lineSeparator());
      }
      final String headerLine = String.format(headerTemplate, (Object[]) headers);
      text.append(headerLine).append(System.lineSeparator());
      for (int j = 0; j < totalWidth; j++) {
        text.append("-");
      }
      for (String[] row : data) {
        final String dataLine = String.format(dataTemplate, (Object[]) row);
        text.append(System.lineSeparator());
        text.append(dataLine);
      }
      final byte[] bytes = text.toString().getBytes();
      this.stream.write(bytes);
      this.stream.flush();
    } catch (IOException ex) {
      throw new IllegalStateException("Failed to print DataFrame", ex);
    }
  }

  private String tableName(Relation frame, int width) {
    if (frame.name().length() > width) {
      return frame.name();
    }
    int diff = width - frame.name().length();
    String result = StringUtils.repeat(" ", diff / 2) + frame.name();
    return result + StringUtils.repeat(" ", width - result.length());
  }

  /**
   * Returns the header string tokens for the frame
   *
   * @param frame the frame to create header tokens
   * @return the header tokens
   */
  private String[] getHeaderTokens(Relation frame) {
    final int colCount = frame.columnCount();
    final String[] header = new String[colCount];
    IntStream.range(0, colCount)
        .forEach(
            colIndex -> {
              header[colIndex] = frame.column(colIndex).name();
            });
    return header;
  }

  private String getDataToken(Column<?> col, int i) {
    return col.size() > i ? col.getString(i) : TOO_SHORT_COLUMN_MARKER;
  }

  /**
   * Returns the 2-D array of data tokens from the frame specified
   *
   * @param frame the DataFrame from which to create 2D array of formatted tokens
   * @return the array of data tokens
   */
  private String[][] getDataTokens(Relation frame) {
    if (frame.rowCount() == 0) return new String[0][0];
    final int rowCount = Math.min(maxRows, frame.rowCount());
    final boolean truncated = frame.rowCount() > maxRows;
    final int colCount = frame.columnCount();
    final String[][] data;
    if (truncated) {
      data = new String[rowCount + 1][colCount];
      int i;
      for (i = 0; i < Math.ceil((double) rowCount / 2); i++) {
        for (int j = 0; j < colCount; j++) {
          Column<?> col = frame.column(j);
          data[i][j] = getDataToken(col, i);
        }
      }
      for (int j = 0; j < colCount; j++) {
        data[i][j] = "...";
      }
      for (++i; i <= rowCount; i++) {
        for (int j = 0; j < colCount; j++) {
          Column<?> col = frame.column(j);
          data[i][j] = getDataToken(col, frame.rowCount() - maxRows + i - 1);
        }
      }
    } else {
      data = new String[rowCount][colCount];
      for (int i = 0; i < rowCount; i++) {
        for (int j = 0; j < colCount; j++) {
          Column<?> col = frame.column(j);
          String value = getDataToken(col, i);
          data[i][j] = value == null ? "" : value;
        }
      }
    }
    return data;
  }
}
