package com.google.research.bleth.simulator;

import static com.google.common.truth.Truth.assertThat;

import com.google.appengine.api.datastore.dev.LocalDatastoreService;
import com.google.appengine.tools.development.testing.LocalDatastoreServiceTestConfig;
import com.google.appengine.tools.development.testing.LocalServiceTestHelper;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.junit.MockitoJUnitRunner;

@RunWith(MockitoJUnitRunner.class)
public class TracingSimulationIT {

    private final LocalServiceTestHelper helper =
            new LocalServiceTestHelper(new LocalDatastoreServiceTestConfig()
                    .setAutoIdAllocationPolicy(LocalDatastoreService.AutoIdAllocationPolicy.SCATTERED));

    private static final IMovementStrategy MOVE_UP = new UpMovementStrategy();
    private static final IMovementStrategy STATIONARY = new StationaryMovementStrategy();

    private static final int BOARD_DIMENSION_EQUALS_TWO = 2;
    private static final int MAX_ROUNDS_EQUALS_TWO = 2;
    private static final int NUMBER_OF_BEACONS_EQUALS_ONE = 1;
    private static final int NUMBER_OF_OBSERVERS_EQUALS_ONE = 1;
    private static final double TRANSMISSION_THRESHOLD_RADIUS_EQUALS_ONE = 1.0;
    private static final int AWAKENESS_CYCLE_EQUALS_TWO = 2;
    private static final int AWAKENESS_DURATION_EQUALS_ONE = 1;
    private static final AwakenessStrategyFactory.Type FIXES_AWAKENESS_STRATEGY_TYPE = AwakenessStrategyFactory.Type.FIXED;

    @Before
    public void setUp() {
        helper.setUp();
    }

    @Test
    public void runSimulationSingleRoundVerifyAgentsLocations() {
        // Create new simulation.
        AbstractSimulation simulation = new TracingSimulation.Builder()
                .setMaxNumberOfRounds(MAX_ROUNDS_EQUALS_TWO)
                .setRowNum(BOARD_DIMENSION_EQUALS_TWO)
                .setColNum(BOARD_DIMENSION_EQUALS_TWO)
                .setBeaconsNum(NUMBER_OF_BEACONS_EQUALS_ONE)
                .setObserversNum(NUMBER_OF_OBSERVERS_EQUALS_ONE)
                .setTransmissionThresholdRadius(TRANSMISSION_THRESHOLD_RADIUS_EQUALS_ONE)
                .setBeaconMovementStrategy(MOVE_UP)
                .setObserverMovementStrategy(STATIONARY)
                .setAwakenessCycle(AWAKENESS_CYCLE_EQUALS_TWO)
                .setAwakenessDuration(AWAKENESS_DURATION_EQUALS_ONE)
                .setAwakenessStrategyType(FIXES_AWAKENESS_STRATEGY_TYPE)
                .build();

        // Find agents and calculate their expected location after a single round.
        Location beaconInitialLocation = simulation.beacons.get(0).getLocation();
        Location observerInitialLocation = simulation.observers.get(0).getLocation();
        Location beaconExpectedLocation = predictLocationAfterMoveUp(beaconInitialLocation);
        Location observerExpectedLocation = observerInitialLocation;

        // Run single-rounded simulation and find agents again.
        simulation.run();
        Location beaconActualLocation = simulation.beacons.get(0).getLocation();
        Location observerActualLocation = simulation.observers.get(0).getLocation();

        assertThat(beaconActualLocation).isEqualTo(beaconExpectedLocation);
        assertThat(observerActualLocation).isEqualTo(observerExpectedLocation);
    }

    private Location predictLocationAfterMoveUp(Location location) {
        if (location.row() == 0) {
            return location;
        }
        return Location.create(location.row() - 1, location.col());
    }
}

@@ -2818,32 +2818,35 @@
 nInitialLocation
+, 1
 );%0A        Locat
@@ -3391,85 +3391,31 @@
 tion
-) %7B%0A        if (location.row() == 0) %7B%0A            return location;%0A        %7D
+, int numberOfRounds) %7B
 %0A   
@@ -3442,16 +3442,28 @@
 .create(
+Math.max(0, 
 location
@@ -3471,17 +3471,31 @@
 row() - 
-1
+numberOfRounds)
 , locati

Make predictLocationAfterMoveUp more generic.