package jrex;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Table;
import com.google.common.collect.Tables;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import javax.swing.table.AbstractTableModel;

/**
 *
 * @author Mark George <mark.george@otago.ac.nz>
 */
public class PivotTableModel extends AbstractTableModel {

	private Table<Integer, Integer, String> table = HashBasedTable.create();

	@Override
	public int getRowCount() {
		return table.rowKeySet().size();
	}

	@Override
	public int getColumnCount() {
		return table.columnKeySet().size();
	}

	@Override
	public String getColumnName(int columnIndex) {
		return String.valueOf(columnIndex + 1);
	}

	@Override
	public Class<?> getColumnClass(int columnIndex) {
		return Object.class;
	}

	@Override
	public boolean isCellEditable(int rowIndex, int columnIndex) {
		return false;
	}

	@Override
	public Object getValueAt(int rowIndex, int columnIndex) {
		return table.get(rowIndex, columnIndex);
	}

	@Override
	public void setValueAt(Object aValue, int rowIndex, int columnIndex) {
		table.put(rowIndex, columnIndex, aValue.toString());
		fireTableCellUpdated(rowIndex, columnIndex);
	}

	public void addRow(String[] values) {
		int rowCount = getRowCount();
		for (int i = 0; i < values.length; i++) {
			table.put(rowCount, i, values[i]);
		}
		fireTableRowsInserted(rowCount, rowCount);
	}

	public Collection<String> getHeaders() {
		return table.row(0).values();
	}

	public void transpose() {
		table = Tables.transpose(table);
		fireTableStructureChanged();
	}

	public void pivotWideToLong(List<String> preservedColumns, List<String> valueColumns) {
		Table<Integer, Integer, String> pivoted = HashBasedTable.create();

		//### create multimaps so that we can easily get the data by header ###//
		ListMultimap<String, String> allMap = ArrayListMultimap.create();

		Map<Integer, String> headers = table.row(0);

		for (Map.Entry<Integer, String> col : headers.entrySet()) {
			for (Integer row : table.rowKeySet()) {
				if (!row.equals(0)) {
					allMap.put(col.getValue(), table.get(row, col.getKey()));
				}
			}
		}

		// create multimap for preserved columns only
		ListMultimap<String, String> preservedMap = ArrayListMultimap.create();
		for (String preservedColumn : preservedColumns) {
			preservedMap.putAll(preservedColumn, allMap.get(preservedColumn));
		}

		// create multimap for value columns only
		ListMultimap<String, String> valueMap = ArrayListMultimap.create();
		for (String valueColumn : valueColumns) {
			valueMap.putAll(valueColumn, allMap.get(valueColumn));
		}

		//### add headers to privoted table ###//
		int col = 0;

		// preserved column headers
		for (String preservedColumn : preservedColumns) {
			pivoted.put(0, col, preservedColumn);
			col++;
		}

		// value column headers (key value pairs)
		pivoted.put(0, col, "Key");
		pivoted.put(0, col + 1, "Value");

		int row = 1; // headers are row 0, so start at 1
		int pCol = 0;
		int pColSize = preservedColumns.size();
		int index = 0;  // original row index

		//### add preserved columns ###/
		for (String preservedCol : preservedColumns) {
			for (String preservedVal : preservedMap.get(preservedCol)) {
				for (String valueColumn : valueColumns) {
					// write preserved column
					pivoted.put(row, pCol, preservedVal);

					// write value columns
					List<String> values = new ArrayList<>(valueMap.get(valueColumn));
					String value = values.get(index);

					pivoted.put(row, pColSize, valueColumn);
					pivoted.put(row++, pColSize + 1, value);
				}
				index++;
			}
			index = 0;
			row = 1;
			pCol++;
		}

		table = pivoted;
		fireTableStructureChanged();
	}

	public void pivotLongToWide(List<String> repeatingColumns, String headerColumn, String valueColumn) {
		//### create multimaps so that we can easily get the data by header ###//
		ListMultimap<String, String> allMap = ArrayListMultimap.create();

		Map<Integer, String> headers = table.row(0);

		for (Map.Entry<Integer, String> col : headers.entrySet()) {
			for (Integer row : table.rowKeySet()) {
				if (!row.equals(0)) {
					allMap.put(col.getValue(), table.get(row, col.getKey()));
				}
			}
		}

		Table<String, String, String> unique = HashBasedTable.create();

		for (Integer rowIdx : table.rowKeySet()) {

			if (rowIdx == 0) {
				continue; // header row so ignore
			}

			StringBuilder key = new StringBuilder();
			for (String repeatingColumn : repeatingColumns) {
				key.append(allMap.get(repeatingColumn).get(rowIdx - 1));
			}

			for (String repeatingColumn : repeatingColumns) {
				unique.put(key.toString(), repeatingColumn, allMap.get(repeatingColumn).get(rowIdx - 1));
			}

			unique.put(key.toString(), allMap.get(headerColumn).get(rowIdx - 1), allMap.get(valueColumn).get(rowIdx - 1));

		}

		// create a new indexed table for the table model
		Table<Integer, Integer, String> pivoted = HashBasedTable.create();

		int rowIdx = 0;

		for (Map.Entry<String, Map<String, String>> row : unique.rowMap().entrySet()) {

			if (rowIdx == 0) {
				// write headers
				Set<String> colHeaders = row.getValue().keySet();

				int colIdx = 0;
				for (String header : colHeaders) {
					pivoted.put(rowIdx, colIdx, header);
					colIdx++;
				}
			}

			int colIdx = 0;

			for (Map.Entry<String, String> entry : row.getValue().entrySet()) {

				pivoted.put(rowIdx+1, colIdx, entry.getValue());
				colIdx++;
			}
			rowIdx++;
		}

		table = pivoted;
		fireTableStructureChanged();
	}

	public Table<Integer, Integer, String> getTable() {
		return table;
	}

}