package jrex.ui.model;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Table;
import com.google.common.collect.Tables;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import javax.swing.table.AbstractTableModel;
import jrex.ui.NonUniqueKeyException;

/**
 * A table model that can be pivoted.
 *
 * @author Mark George <mark.george@otago.ac.nz>
 */
public class PivotTableModel extends AbstractTableModel {

   private Table<Integer, Integer, String> table = HashBasedTable.create();
   private Map<Integer, String> columnHeadings = new LinkedHashMap<>();
   private Map<String, Integer> headingIndexes = new HashMap<>();

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

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

   public void setColumnHeadings(String[] headings) {
      columnHeadings.clear();
      for (int i = 0; i < headings.length; i++) {
         columnHeadings.put(i + 1, headings[i]);
         headingIndexes.put(headings[i], i+1);
      }
   }

   @Override
   public String getColumnName(int columnIndex) {
      if (columnHeadings.isEmpty()) {
         // no explicit headings, so use column index
         return String.valueOf(columnIndex + 1);
      } else {
         return columnHeadings.get(columnIndex + 1);
      }
   }

   public Integer getHeadingIndex(String heading) {
      return headingIndexes.get(heading);
   }

   @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 List<String> getHeadings() {
      if (columnHeadings.isEmpty()) {
         for (int i=0; i < table.columnKeySet().size(); i++) {
            columnHeadings.put(i+1, String.valueOf(i+1));
            headingIndexes.put(String.valueOf(i+1), i+1);
         }
      }
      return new ArrayList(columnHeadings.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();

      for (Map.Entry<Integer, String> col : columnHeadings.entrySet()) {
         for (Integer row : table.rowKeySet()) {
            allMap.put(col.getValue(), table.get(row, col.getKey() - 1));
         }
      }

      // 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 pivoted table ###//
      int numHeaders = preservedColumns.size() + 2;
      String[] headings = new String[numHeaders];
      int col = 0;

      // preserved column headers
      for (String preservedColumn : preservedColumns) {
         headings[col] = preservedColumn;
         col++;
      }

      headings[col] = "Key";
      headings[col + 1] = "Value";

      setColumnHeadings(headings);

      int row = 0;
      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 = 0;
         pCol++;
      }

      table = pivoted;
      fireTableStructureChanged();
   }

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

      for (Map.Entry<Integer, String> col : columnHeadings.entrySet()) {
         for (Integer row : table.rowKeySet()) {
            allMap.put(col.getValue(), table.get(row, col.getKey() - 1));
         }
      }

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

      
      for (Integer rowIdx : table.rowKeySet()) {
         List<String> columnVals = new ArrayList<>();

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

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

         if(unique.contains(key.toString(), allMap.get(headerColumn).get(rowIdx))) {
            throw new NonUniqueKeyException("There is already a value for [" + allMap.get(headerColumn).get(rowIdx) + "] for " + columnVals +".\n\nAborting pivot since data will be lost.");
         }
         
         unique.put(key.toString(), allMap.get(headerColumn).get(rowIdx), allMap.get(valueColumn).get(rowIdx));

      }

      // 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> colSet = row.getValue().keySet();
            setColumnHeadings(colSet.toArray(new String[colSet.size()]));
         }

         int colIdx = 0;

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

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

      table = pivoted;
      fireTableStructureChanged();
   }

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

}