package com.google.research.bleth.simulator;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Table;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/** A container for all the agents, representing their locations, either real or estimated. */
public class Board {
    private int rowNum;
    private int colNum;
    private ArrayList<Agent>[][] matrix;

    public Board(int rows, int cols) {
        rowNum = rows;
        colNum = cols;
        matrix = new ArrayList[rowNum][colNum];
        for (int row = 0; row < rows; row++) {
            Arrays.setAll(matrix[row], ArrayList::new);
        }
    }

    private boolean isLocationValid(Location location) {
        Preconditions.checkNotNull(location);
        boolean isRowValid = 0 <= location.row && location.row < rowNum;
        boolean isColValid = 0 <= location.col && location.col < colNum;
        return isColValid && isRowValid;
    }

    /**
     * Create and return an immutable list of all the agents which located on specific location.
     * @param location is the requested location to retrieve all the agents that located on it.
     * @return an immutable copy of the list of all agents which located on location.
     */
    public List<Agent> getAgentsOnLocation(Location location) {
        if (!isLocationValid(location)) {
            throw new IllegalArgumentException("Invalid Location");
        }
        return ImmutableList.copyOf(matrix[location.row][location.col]);
    }

    /**
     * Place an agent on board if the given location is valid.
     * @param newLocation is the location where the agent will be placed.
     * @param agent is the Agent which placed on the board.
     */
    public void placeAgent(Location newLocation, Agent agent) {
        Preconditions.checkNotNull(agent);
        if (!isLocationValid(newLocation)) {
            throw new IllegalArgumentException("Invalid Location");
        }
        matrix[newLocation.row][newLocation.col].add(agent);
    }

    /**
     * Remove an agent from its current location on board and place it on new location.
     * @param oldLocation is the agent's current location.
     * @param newLocation is the location where the agent will be placed.
     * @param agent is the Agent which moves from oldLocation to newLocation.
     */
    public void moveAgent(Location oldLocation, Location newLocation, Agent agent) {
        Preconditions.checkNotNull(agent);
        if (!isLocationValid(oldLocation) || !isLocationValid(newLocation)) {
            throw new IllegalArgumentException("Invalid Location");
        }
        matrix[oldLocation.row][oldLocation.col].remove(agent);
        placeAgent(newLocation, agent);
    }
}

@@ -117,29 +117,26 @@
 collect.
-ImmutableList
+ArrayTable
 ;%0Aimport
@@ -162,21 +162,29 @@
 collect.
-Table
+ImmutableList
 ;%0Aimport
@@ -222,21 +222,55 @@
 va.util.
-Array
+List;%0Aimport java.util.stream.Collector
 s;%0Aimpor
@@ -281,20 +281,32 @@
 va.util.
-Li
 st
+ream.IntStream
 ;%0A%0A/** A
@@ -424,24 +424,30 @@
     private 
+final 
 int rowNum;%0A
@@ -458,16 +458,22 @@
 private 
+final 
 int colN
@@ -487,16 +487,45 @@
  private
+ ArrayTable%3CInteger, Integer,
  ArrayLi
@@ -537,12 +537,9 @@
 ent%3E
-%5B%5D%5B%5D
+%3E
  mat
@@ -544,16 +544,17 @@
 atrix;%0A%0A
+%0A
     publ
@@ -650,142 +650,363 @@
 ix =
- new
  Array
-List%5BrowNum%5D%5BcolNum%5D;%0A        for (int row = 0; row %3C rows; row++) %7B%0A            Arrays.setAll(matrix%5Brow%5D, ArrayList::new);
+Table.create(IntStream.range(0, rows).boxed().collect(Collectors.toList()),%0A                                    IntStream.range(0, cols).boxed().collect(Collectors.toList()));%0A        for (int row = 0; row %3C rows; row++) %7B%0A            for (int col = 0; col %3C cols; col++) %7B%0A                matrix.set(row, col, new ArrayList%3CAgent%3E());%0A            %7D
 %0A   
@@ -1837,17 +1837,21 @@
 f(matrix
-%5B
+.get(
 location
@@ -1854,18 +1854,18 @@
 tion.row
-%5D%5B
+, 
 location
@@ -1868,17 +1868,17 @@
 tion.col
-%5D
+)
 );%0A    %7D
@@ -2336,17 +2336,21 @@
   matrix
-%5B
+.get(
 newLocat
@@ -2356,18 +2356,18 @@
 tion.row
-%5D%5B
+, 
 newLocat
@@ -2373,17 +2373,17 @@
 tion.col
-%5D
+)
 .add(age
@@ -3007,17 +3007,21 @@
   matrix
-%5B
+.get(
 oldLocat
@@ -3031,10 +3031,10 @@
 .row
-%5D%5B
+, 
 oldL
@@ -3048,9 +3048,9 @@
 .col
-%5D
+)
 .rem

Change the matrix from a 2D array to an ArrayTable.