using POMDPs
using POMDPDistributions
using DiscreteValueIteration
using MCTS

#using POMDPToolbox
#using SARSOP
#using QMDP

using PyPlot
using PyCall
@pyimport matplotlib.patches as patch

# TODO:
# - observations/pomdp
# - add tSinceLastRecalibration to BotState, implement as POMDP with noisy observations
# - add recalibration action, with some cost
# - might penalize for placing target when not high enough belief?

### conventions
# - world locations are (1,1) in bottom-left, (+x,+y) in top-right

### RUNTIME CONFIG
#WORLD_CONFIG = :small
#WORLD_CONFIG = :medium
WORLD_CONFIG = :large
assert( (WORLD_CONFIG == :small) || (WORLD_CONFIG == :medium) || (WORLD_CONFIG == :large) )
# certainty configs, useful for debugging, testing, comparisons
#CERTAINTY_CONFIG = :certain
#CERTAINTY_CONFIG = :uncertainMovement_uncertainTargetInteraction
#CERTAINTY_CONFIG = :uncertainMovement_certainTargetInteraction
CERTAINTY_CONFIG = :uncertainMovement_uncertainTargetInteraction_decreasedCertainty
#CERTAINTY_CONFIG = :uncertainMovement_certainTargetInteraction_decreasedCertainty
if CERTAINTY_CONFIG == :certain
  PROBABILITY_MOVEMENT_SUCCESS = 1.0 # probability of transitioning to the desired state when moving/turning in a specific direction
  PROBABILITY_TARGET_INTERACTION_SUCCESS = 1.0 # probability of successfully picking up/placing target object
elseif CERTAINTY_CONFIG == :uncertainMovement_uncertainTargetInteraction
  PROBABILITY_MOVEMENT_SUCCESS = 0.9
  PROBABILITY_TARGET_INTERACTION_SUCCESS = 0.9
elseif CERTAINTY_CONFIG == :uncertainMovement_certainTargetInteraction # useful for eliminating increased rewards for picking/placing target repeatedly due to world-induced failure
  PROBABILITY_MOVEMENT_SUCCESS = 0.9
  PROBABILITY_TARGET_INTERACTION_SUCCESS = 1.0
elseif CERTAINTY_CONFIG == :uncertainMovement_uncertainTargetInteraction_decreasedCertainty
  PROBABILITY_MOVEMENT_SUCCESS = 0.7
  PROBABILITY_TARGET_INTERACTION_SUCCESS = 0.7
elseif CERTAINTY_CONFIG == :uncertainMovement_certainTargetInteraction_decreasedCertainty
  PROBABILITY_MOVEMENT_SUCCESS = 0.7
  PROBABILITY_TARGET_INTERACTION_SUCCESS = 1.0
end
# solver config
RUN_DISCRETE_VALUE_ITERATION = true
RUN_SARSOP = false

### config constants
if WORLD_CONFIG == :small
  BASE_LOCATION = (5, 5) # location of base where targets originate
  WORLD_SIZE_X = 5 # number of blocks in world x dimension
  WORLD_SIZE_Y = 5 # number of blocks in world y dimension
  NUM_TOTAL_TARGETS = 4 # number of total targets the bot must place
  NUM_TOTAL_OBSTACLES = 3 # number of total obstacles the bot must avoid
elseif WORLD_CONFIG == :medium
  BASE_LOCATION = (1, 1) # location of base where targets originate
  WORLD_SIZE_X = 10 # number of blocks in world x dimension
  WORLD_SIZE_Y = 10 # number of blocks in world y dimension
  NUM_TOTAL_TARGETS = 6 # number of total targets the bot must place
  NUM_TOTAL_OBSTACLES = 6 # number of total obstacles the bot must avoid
elseif WORLD_CONFIG == :large
  BASE_LOCATION = (20, 10) # location of base where targets originate
  WORLD_SIZE_X = 20 # number of blocks in world x dimension
  WORLD_SIZE_Y = 15 # number of blocks in world y dimension
  NUM_TOTAL_TARGETS = 15 # number of total targets the bot must place
  NUM_TOTAL_OBSTACLES = 6 # number of total obstacles the bot must avoid
end

### functions for generating specific constants
# create dict with facingDirectionSymbol => Int64, for use in sub2ind
function createFacingDirectionSymbolsIndexDict(facingDirectionSymbols)
  indexDict = Dict{Symbol, Int64}()
  for (i, symbol) in enumerate(facingDirectionSymbols)
    indexDict[symbol] = i
  end
  return indexDict
end

### constants
ACTION_SYMBOLS = [:moveForward, :turnLeft, :turnRight, :pickUpTarget, :placeTarget] # all possible robot actions (as symbols)
NUM_ACTIONS = length(ACTION_SYMBOLS) # total number of possible actions
FACING_DIRECTION_SYMBOLS = [:up, :down, :left, :right] # all possible facing directions (as symbols)
FACING_DIRECTION_SYMBOLS_INDEX_DICT = createFacingDirectionSymbolsIndexDict(FACING_DIRECTION_SYMBOLS)
FACING_DIRECTION_ADDITION_VECTORS = Dict(:up => (0, +1), :down => (0, -1), :left => (-1, 0), :right => (+1, 0)) # dictionary for addition vectors based on facing directions; useful for computing new location if move forward when facing a particular direction
NUM_FACING_DIRECTIONS = length(FACING_DIRECTION_SYMBOLS) # total number of facing directions
NUM_MOVE_FORWARD_POSSIBLE_NEW_LOCATIONS = 4 # total number of new locations we can end up in when moving forward
NUM_TOTAL_STATES = (WORLD_SIZE_X+2) * (WORLD_SIZE_Y+2) * NUM_FACING_DIRECTIONS * (NUM_TOTAL_TARGETS+1) * 2 # *2 for each bitwise state variable, NUM_TOTAL_TARGETS+1 to allow 0 state, SIZE+2 to allow for invalid states on either side for each dimension for value iteration etc
OBSTACLE_ERROR_PENALTY_DEFAULT = -100 # penalty incurred if receive error at this default obstacle
REWARD_PICKUP_TARGET = 100.0 # reward for picking up a target from the base location
REWARD_PLACE_TARGET = 1000.0 # reward for successfully placing a target at its desired location
PENALTY_INVALID_STATE = -100 # penalty for entering invalid state (used for value iteration)
PENALTY_MOVEMENT = -1 # penalty for movement
PENALTY_TURN = -1 # penalty for turning/changing direction
PENALTY_OUT_OF_BOUNDS = -50 # penalty for moving out-of-bounds
PENALTY_INVALID_PICKUPTARGET_NOTBASE = -30 # penalty for trying to pick up a target when not at the base location
PENALTY_INVALID_PICKUPTARGET_ALREADYHOLDING = -100 # penalty for trying to pick up a target when already holding one
PENALTY_INVALID_PLACETARGET = -50 # penalty for trying to place a target when aren't actually holding one
PENALTY_WRONG_PLACETARGET = -100 # penalty for placing a target at the wrong location
PENALTY_STANDING_ON_PAST_TARGET_LOCATION = 0 # penalty if standing on past target location in current state, not necessarily meaning to get here. in practice, this is penalty if turning on location of target just placed
PENALTY_MOVING_TO_PAST_TARGET_LOCATION = -100 # penalty for walking on a target location where target has already been placed
DISCOUNT_FACTOR = 0.9 # discount factor

### BotState model
type BotState <: POMDPs.State
  x::Int64 # x position (ground-truth)
  y::Int64 # y position (ground-truth)
  facingDirection::Symbol # direction robot is facing
  numTargetsRemaining::Int64 # number of targets not yet completed/placed
  isHoldingTarget::Bool # true if currently holding a target, false otherwise
end

### Action type
type BotAction <: POMDPs.Action
  actionSymbol::Symbol
end

### Observation type
type BotObservation <: POMDPs.Observation
  xObs::Int64 # x position (observed, noisy)
  yObs::Int64 # y position (observed, noisy)
end

### Targets - target specification
type TargetSpecification
  x::Int64 # x position
  y::Int64 # y position
  orderIndex::Int64 # 1-indexed index of this target's place in overall order targets must be placed
end

### Obstacles - obstacle specification
type ObstacleSpecification
  x::Int64 # x position
  y::Int64 # y position
  #errorFraction::Float64 # fraction that this obstacle causes error penalty
  errorPenalty::Float64 # penalty incurred by error at this obstacle
end

### World model
# - See initial constants sections for comments on specific variables
type BotWorld <: POMDP
  WORLD_CONFIG::Symbol
  WORLD_SIZE_X::Int64
  WORLD_SIZE_Y::Int64
  BASE_LOCATION::Tuple{Int64, Int64}
  NUM_TOTAL_STATES::Int64
  NUM_MAX_STATE_NEIGHBORS::Int64
  STATES_INDEXES_MAP::Dict{BotState, Int64} # map from state to state index
  NUM_TOTAL_TARGETS::Int64
  ACTION_SYMBOLS::Array{Symbol}
  FACING_DIRECTION_SYMBOLS::Array{Symbol}
  FACING_DIRECTION_ADDITION_VECTORS::Dict{Symbol, Tuple{Int64, Int64}}
  FACING_DIRECTION_SYMBOLS_INDEX_DICT::Dict{Symbol, Int64}
  NUM_MOVE_FORWARD_POSSIBLE_NEW_LOCATIONS::Int64
  TARGET_SPECIFICATIONS::Vector{TargetSpecification} # list of targets for this pomdp world
  OBSTACLE_SPECIFICATIONS::Vector{ObstacleSpecification} # list of obstacles for this pomdp world
  REWARD_STATES_ACTIONS_VALUES::Vector{Tuple{BotState, BotAction, Float64}}
  PROBABILITY_MOVEMENT_SUCCESS::Float64
  PROBABILITY_TARGET_INTERACTION_SUCCESS::Float64
  REWARD_TARGET_PICKUP::Float64
  REWARD_TARGET_PLACEMENT::Float64
  PENALTY_INVALID_STATE::Float64
  PENALTY_MOVEMENT::Float64
  PENALTY_TURN::Float64
  PENALTY_OUT_OF_BOUNDS::Float64
  PENALTY_INVALID_PICKUPTARGET_NOTBASE::Float64
  PENALTY_INVALID_PICKUPTARGET_ALREADYHOLDING::Float64
  PENALTY_INVALID_PLACETARGET::Float64
  PENALTY_WRONG_PLACETARGET::Float64
  PENALTY_STANDING_ON_PAST_TARGET_LOCATION::Float64
  PENALTY_MOVING_TO_PAST_TARGET_LOCATION::Float64
  DISCOUNT_FACTOR::Float64
end

### StateSpace type
type StateSpace <: POMDPs.AbstractSpace
  states::Vector{BotState}
end
### StateSpace multiple-dispatch for POMDPs.jl
function POMDPs.states(pomdp::BotWorld)
  #states, _ = createStates(pomdp)
  #return StateSpace(states)
  return StateSpace(createStates(pomdp))
end
### StateSpace domain - return iterator over state space for POMDPs.jl
function POMDPs.domain(stateSpace::StateSpace)
  return stateSpace.states
end
### allow uniform sampling (in-place) of state space
function POMDPs.rand!(rng::AbstractRNG, state::BotState, stateSpace::StateSpace)
  #sampledState = stateSpace.states[rand(rng, 1:end)]
  copy!(state, stateSpace.states[rand(rng, 1:end)])
  return state
end
### actually create states
function createStates(pomdp::BotWorld)
  states, stateIndexesMap = createStates(pomdp.WORLD_SIZE_X, pomdp.WORLD_SIZE_Y, pomdp.FACING_DIRECTION_SYMBOLS, pomdp.NUM_TOTAL_TARGETS, pomdp.NUM_TOTAL_STATES)
  return states
  ## instantiate vector
  #states = BotState[]
  #stateIndexesMap = Dict{BotState, Int64}()
  #currentStateIndex = 1
  ## for each x location
  #for x=1:pomdp.WORLD_SIZE_X
  #  # for each y location
  #  for y=1:pomdp.WORLD_SIZE_Y
  #    # for each facing direction
  #    for facingDirection in pomdp.FACING_DIRECTION_SYMBOLS
  #      # for each numTargetsRemaining
  #      for numTargetsRemaining=0:pomdp.NUM_TOTAL_TARGETS
  #        # for each isHoldingTarget
  #        for isHoldingTarget in [true, false]
  #          # create state
  #          state = BotState(x, y, facingDirection, numTargetsRemaining, isHoldingTarget)
  #          # append state
  #          push!(states, state)
  #          #push!(states, BotState(x, y, facingDirection, numTargetsRemaining, isHoldingTarget))
  #          # get state index
  #          stateIndexesMap[state] = currentStateIndex
  #          # increment currentStateIndex
  #          currentStateIndex += 1
  #        end
  #      end
  #    end
  #  end
  #end
  ## assert
  #assert( length(states) == pomdp.NUM_TOTAL_STATES )
  #assert( length(stateIndexesMap) == pomdp.NUM_TOTAL_STATES )
  ## print
  #@printf("Total number of states: %d\n", pomdp.NUM_TOTAL_STATES)
  ## return
  #return states, stateIndexesMap
end
# NOTE: also creates invalid-location states around the edge
function createStates(worldSizeX, worldSizeY, facingDirectionSymbols, numTotalTargets, numTotalStates)
  # instantiate vector
  states = BotState[]
  stateIndexesMap = Dict{BotState, Int64}()
  currentStateIndex = 1
  # for each x location
  for x=0:worldSizeX+1
    # for each y location
    for y=0:worldSizeY+1
      # for each facing direction
      for facingDirection in facingDirectionSymbols
        # for each numTargetsRemaining
        for numTargetsRemaining=0:numTotalTargets
          # for each isHoldingTarget
          for isHoldingTarget in [true, false]
            # create state
            state = BotState(x, y, facingDirection, numTargetsRemaining, isHoldingTarget)
            # append state
            push!(states, state)
            #push!(states, BotState(x, y, facingDirection, numTargetsRemaining, isHoldingTarget))
            # get state index
            stateIndexesMap[state] = currentStateIndex
            # increment currentStateIndex
            currentStateIndex += 1
          end
        end
      end
    end
  end
  # assert
  assert( length(states) == numTotalStates )
  assert( length(stateIndexesMap) == numTotalStates )
  # print
  @printf("Total number of states: %d\n", numTotalStates)
  # return
  return states, stateIndexesMap
end
### helper function for MCTS
Base.hash(state::BotState, h::UInt64 = zero(UInt64)) = hash(state.x, hash(state.y, hash(state.facingDirection, hash(state.numTargetsRemaining, hash(state.isHoldingTarget, h)))))
### parse a passed-in state into a tuple of its components
function parseBotState(state::BotState)
  return state.x, state.y, state.facingDirection, state.numTargetsRemaining, state.isHoldingTarget
end
### get initial state
function getInitialBotState(pomdp::BotWorld)
  return BotState(pomdp.BASE_LOCATION[1], pomdp.BASE_LOCATION[2], :up, pomdp.NUM_TOTAL_TARGETS, false)
end
### get default state
function getDefaultBotState(pomdp::BotWorld)
  return BotState(1, 1, :ANY_DIRECTION, pomdp.NUM_TOTAL_TARGETS, false)
end
### get invalid state, which can't be equal to any valid states
function getInvalidBotState(pomdp::BotWorld)
  return BotState(-1, -1, :ANY_DIRECTION, -1, false)
end
### make copy of inputState, place into outputState in-place
function Base.copy!(outputState::BotState, inputState::BotState)
  outputState.x = inputState.x
  outputState.y = inputState.y
  outputState.facingDirection = inputState.facingDirection
  outputState.numTargetsRemaining = inputState.numTargetsRemaining
  outputState.isHoldingTarget = inputState.isHoldingTarget
  return outputState
end
### get max number of state neighbors at any given state
# using actions [:moveForward, :turnLeft, :turnRight, :pickUpTarget, :placeTarget]
function getMaxNumStateNeighbors()
  # - can move forward - get forward, same, or forward-left, or forward-right (4)
  # - can face new direction by turning, or face same direction (3)
  # - can have new number of targets remaining, or same (2)
  # - can be holding or not holding target (2)
  # - since we can only pick one action at a time, we take the max of the above options
  return 4
end
### determine if two states are equal, optionally ignoring specific params
function isEqualStates(s1::BotState, s2::BotState; ignoreDirection=false)
  if (ignoreDirection == false) && (s1.facingDirection != s2.facingDirection)
    return false
  end
  return (s1.x == s2.x) && (s1.y == s2.y) && (s1.numTargetsRemaining == s2.numTargetsRemaining) && (s1.isHoldingTarget == s2.isHoldingTarget)
  #if (s1.x != s2.x)
  #  return false
  #elseif (s1.y != s2.y)
  #  return false
  #elseif (s1.numTargetsRemaining != s2.numTargetsRemaining)
  #  return false
  #elseif (s1.facingDirection != s2.facingDirection) && (ignoreDirection == false)
  #  return false
  #elseif (s1.isHoldingTarget != s2.isHoldingTarget)
  #  return false
  #else
  #  return true
  #end
end
### determine if state and TargetSpecification are equal in location
function isEqualLocations(state::BotState, targetSpecification::TargetSpecification)
  return (state.x == targetSpecification.x) && (state.y == targetSpecification.y)
end
### determine if state and ObstacleSpecification are equal in location
function isEqualLocations(state::BotState, obstacleSpecification::ObstacleSpecification)
  return (state.x == obstacleSpecification.x) && (state.y == obstacleSpecification.y)
end
### initializer function
function POMDPs.create_state(pomdp::BotWorld)
  return getInitialBotState(pomdp)
end
### helper function for MCTS
Base.isequal(s1::BotState, s2::BotState) = isEqualStates(s1, s2)
### determine if passed-in state is at base location state
function isAtBaseLocation(pomdp::BotWorld, state::BotState)
  return ( (state.x == pomdp.BASE_LOCATION[1]) && (state.y == pomdp.BASE_LOCATION[2]) )
end
### determine if passed-in state is a valid state
function isValidState(pomdp::BotWorld, state::BotState)

  # if out-of-bounds state
  if isOutOfBounds(pomdp, state)
    return false
  end

  # if no targets remaining but holding a target
  if (state.numTargetsRemaining == 0) && (state.isHoldingTarget == true)
    return false
  end

  # otherwise, valid
  return true
end
### determine if passed-in state is a terminal state
function isTerminalState(state::BotState)
  return (state.numTargetsRemaining == 0) #&& (isHoldingTarget == false)# && (isPlacingTarget == false)
end
function MCTS.isterminal(pomdp::BotWorld, state::BotState)
  return isTerminalState(state)
end
### determine if passed-in state has a location that is out-of-bounds
function isOutOfBounds(pomdp::BotWorld, state::BotState)
  return ( (state.x <= 0) || (state.x > pomdp.WORLD_SIZE_X) || (state.y <= 0) || (state.y > pomdp.WORLD_SIZE_Y) )
end
### get intended/expected next state if successfully execute action from prev state
### - NOTE: does not check for validity of this action in this state, just assumes it can happen as commanded (important: will not check if at target location for picking up/placing target; do this elsewhere)
function getIntendedNextState(pomdp::BotWorld, currentState::BotState, action::BotAction)

  # copy currentState to nextState
  nextState = getDefaultBotState(pomdp)
  copy!(nextState, currentState)

  # parse currentState
  currentX, currentY, currentFacingDirection, currentNumTargetsRemaining, currentIsHoldingTarget = parseBotState(currentState)
  # get action symbol from action
  actionSymbol = action.actionSymbol

  # if moveForward
  if actionSymbol == :moveForward
    # get new location if move forward one step in direction we are currently facing
    dx, dy = pomdp.FACING_DIRECTION_ADDITION_VECTORS[currentFacingDirection]
    nextState.x += dx
    nextState.y += dy

  # elseif turnLeft
  elseif actionSymbol == :turnLeft
    if currentFacingDirection == :up
      nextState.facingDirection = :left
    elseif currentFacingDirection == :left
      nextState.facingDirection = :down
    elseif currentFacingDirection == :down
      nextState.facingDirection = :right
    elseif currentFacingDirection == :right
      nextState.facingDirection = :up
    else
      @printf("getIntendedNextState: error with actionSymbol=%s, currentFacingDirection=%s\n", actionSymbol, currentFacingDirection)
      error("getIntendedNextState: invalid parameters")
    end
  # elseif turnRight
  elseif actionSymbol == :turnRight
    if currentFacingDirection == :up
      nextState.facingDirection = :right
    elseif currentFacingDirection == :right
      nextState.facingDirection = :down
    elseif currentFacingDirection == :down
      nextState.facingDirection = :left
    elseif currentFacingDirection == :left
      nextState.facingDirection = :up
    else
      @printf("getIntendedNextState: error with actionSymbol=%s, currentFacingDirection=%s\n", actionSymbol, currentFacingDirection)
      error("getIntendedNextState: invalid parameters")
    end

  # elseif pickUpTarget - NOTE: not checked for validity
  elseif actionSymbol == :pickUpTarget
    nextState.isHoldingTarget = true

  # elseif :placeTarget - NOTE: not checked for validity
  elseif actionSymbol == :placeTarget
    nextState.isHoldingTarget = false
    nextState.numTargetsRemaining -= 1

  # otherwise, invalid actionSymbol
  else
    @printf("getIntendedNextState: error, invalid actionSymbol=%s\n", actionSymbol)
    error("getIntendedNextState: invalid parameters")
  end

  # return
  return nextState
end
### Returns possible new locations if moving forward from a current state
### - First is current state (to work with transitions function format)
### - Second is intended state (moving forward one step)
### - Third is moving forward one, then left one (relative to currentFacingDirection)
### - Fourth is moving forward one, then right one (relative to currentFacingDirection)
function getMoveForwardPossibleNewLocations(pomdp::BotWorld, state::BotState)

  # instantiate
  newLocations = Vector{BotState}()
  for i=1:pomdp.NUM_MOVE_FORWARD_POSSIBLE_NEW_LOCATIONS
    push!(newLocations, getInvalidBotState(pomdp))
  end

  # copy current state into first slot
  copy!(newLocations[1], state)

  # get intended forward state
  intendedNextState = getIntendedNextState(pomdp, state, BotAction(:moveForward))
  # copy into second slot
  copy!(newLocations[2], intendedNextState)

  # get forward-left and forward-right states
  currentFacingDirection = state.facingDirection
  copy!(newLocations[3], intendedNextState)
  copy!(newLocations[4], intendedNextState)
  if currentFacingDirection == :up
    newLocations[3].x -= 1
    newLocations[4].x += 1
  elseif currentFacingDirection == :left
    newLocations[3].y -= 1
    newLocations[4].y += 1
  elseif currentFacingDirection == :down
    newLocations[3].x += 1
    newLocations[4].x -= 1
  elseif currentFacingDirection == :right
    newLocations[3].y += 1
    newLocations[4].y -= 1
  else
    error("getMoveForwardPossibleNewLocations: invalid currentFacingDirection")
  end

  # return
  return newLocations
end
### return state that is nearest to outOfBoundsState, but in bounds
function moveToNearestInBoundsLocation(pomdp::BotWorld, outOfBoundsState::BotState)
  inBoundsState = getInvalidBotState(pomdp)
  copy!(inBoundsState, outOfBoundsState)
  if inBoundsState.x < 1
    inBoundsState.x = 1
  elseif inBoundsState.x > pomdp.WORLD_SIZE_X
    inBoundsState.x = pomdp.WORLD_SIZE_X
  end
  if inBoundsState.y < 1
    inBoundsState.y = 1
  elseif inBoundsState.y > pomdp.WORLD_SIZE_Y
    inBoundsState.y = pomdp.WORLD_SIZE_Y
  end
  assert( !isOutOfBounds(pomdp, inBoundsState) )
  return inBoundsState
end

### ActionSpace type
type ActionSpace <: AbstractSpace
  actions::Vector{BotAction}
end
### ActionSpace multiple-dispatch for POMDPs.jl
function POMDPs.actions(pomdp::BotWorld)
  return ActionSpace(createActions(pomdp))
end
function POMDPs.actions(pomdp::BotWorld, state::BotState, actionSpace::ActionSpace=POMDPs.actions(pomdp))
  return actionSpace
end
### ActionSpace domain - return iterator over action space for POMDPs.jl
function POMDPs.domain(actionSpace::ActionSpace)
  return actionSpace.actions
end
### allow uniform sampling (in-place) of action space
function POMDPs.rand!(rng::AbstractRNG, action::BotAction, actionSpace::ActionSpace)
  #sampledAction = actionSpace.actions[rand(rng, 1:end)]
  copy!(action, actionSpace.actions[rand(rng, 1:end)])
  return action
end
### actually create actions
function createActions(pomdp::BotWorld)
  # instantiate vector
  actions = BotAction[]
  # create BotActions from symbols
  for actionSymbol in pomdp.ACTION_SYMBOLS
    push!(actions, BotAction(actionSymbol))
  end
  # print
  @printf("Total number of actions: %d\n", length(actions))
  # return
  return actions
end
### initializer function
function POMDPs.create_action(pomdp::BotWorld)
  return BotAction(:moveForward)
end
### determine if two actions are equivalent
function isEqualActions(a1::BotAction, a2::BotAction)
  return (a1.actionSymbol == a2.actionSymbol)
end
### make copy of inputAction, place into outputAction in-place
function Base.copy!(outputAction::BotAction, inputAction::BotAction)
  outputAction.actionSymbol = inputAction.actionSymbol
  return outputAction
end

### Transition distribution
type BotWorldDistribution <: POMDPs.AbstractDistribution
  neighbors::Vector{BotState} # the states s' in the distribution
  probabilities::Vector{Float64} # array of probabilities for each corresponding state s'
  categorical::POMDPDistributions.Categorical # comes from POMDPDistributions.jl and is used for sampling
end
### To reduce memory allocation, the POMDPs.jl interface defines some initialization functions that return initial
### types to be filled later. This function returns the distribution type filled with some values. We don't care
### what the distribution container has in it, because it will be modified at each call to the transition function.
function POMDPs.create_transition_distribution(pomdp::BotWorld)
  #neighbors = [BotState(1, 1, :ANY_DIRECTION, NUM_TOTAL_TARGETS, false, false) for i=1:NUM_MAX_STATE_NEIGHBORS]
  neighbors = [getDefaultBotState(pomdp) for i=1:pomdp.NUM_MAX_STATE_NEIGHBORS]
  probabilities = zeros(pomdp.NUM_MAX_STATE_NEIGHBORS) + 1.0/pomdp.NUM_MAX_STATE_NEIGHBORS
  categorical = POMDPDistributions.Categorical(pomdp.NUM_MAX_STATE_NEIGHBORS)
  return BotWorldDistribution(neighbors, probabilities, categorical)
end
### return iterator over discrete states in distribution (neighbors array in our distribution)
function POMDPs.domain(distribution::BotWorldDistribution)
  return distribution.neighbors
end
### implement PDF/PMF function from POMDPS.jl - for discrete distribution, pdf function returns probability of
### drawing the state s from the distribution d
function POMDPs.pdf(distribution::BotWorldDistribution, targetState::BotState)
  for (i, candidateState) in enumerate(distribution.neighbors)
    #if candidateState == targetState
    if isEqualStates(candidateState, targetState)
      return distribution.probabilities[i]
    end
  end
  return 0.0
end
### implement sampling function that can draw samples from our distribution
function POMDPs.rand!(rng::AbstractRNG, state::BotState, distribution::BotWorldDistribution)
  set_prob!(distribution.categorical, distribution.probabilities) # fill the Categorical distribbution with our state probabilities
  #sampledNeighbor = distribution.neighbors
  copy!(state, distribution.neighbors[rand(rng, distribution.categorical)]) # sample a neighbor state according to the distribution
  return state
end
### implement transition function
function POMDPs.transition(pomdp::BotWorld, currentState::BotState, action::BotAction, distribution::BotWorldDistribution=create_transition_distribution(pomdp))

  ### SETUP

  # get reference to neighbors from given distribution variable
  neighbors = distribution.neighbors
  probabilities = distribution.probabilities

  # fill with all zeros, copy current state into first neighbor for convenience
  fill!(probabilities, 0.0)
  copy!(neighbors[1], currentState)
  # copy empty states into other neighbors to make sure we don't have duplicates of currentState/etc, as this causes problems with pdf
  for i=2:length(neighbors)
    copy!(neighbors[i], getInvalidBotState(pomdp))
  end
  #assert(isEqualStates(neighbors[1], currentState))

  ### PARSE CURRENTSTATE AND ACTION FOR CONVENIENCE
  currentX, currentY, currentFacingDirection, currentNumTargetsRemaining, currentIsHoldingTarget = parseBotState(currentState)
  actionSymbol = action.actionSymbol
  #currentX = currentState.x
  #currentY = currentState.y
  #currentFacingDirection = currentState.facingDirection
  #currentNumTargetsRemaining = currentState.numTargetsRemaining
  #currentIsHoldingTarget = currentState.isHoldingTarget
  #currentIsPlacingTarget = currentState.isPlacingTarget

  ### GET INTENDEDNEXTSTATE
  intendedNextState = getIntendedNextState(pomdp, currentState, action)
  ## place into second slot
  #copy!(neighbors[2], intendedNextState)

  ### HANDLE BEING IN OUT-OF-BOUNDS STATE - return to nearest in-bounds state immediately (incur penalty in rewards function already)
  if isOutOfBounds(pomdp, currentState)
    copy!(neighbors[1], moveToNearestInBoundsLocation(pomdp, currentState))
    probabilities[1] = 1.0
    return distribution
  end

  ### HANDLE BEING IN INVALIDSTATE FROM START - stay in invalid state forever
  if !isValidState(pomdp, currentState)
    probabilities[1] = 1.0
    return distribution
  end

  ### HANDLE DETERMINISTIC TRANSITIONS INDEPENDENT OF ACTION

  ## if currently placing target, then transition deterministically to next state where we are no longer placing, are no longer holding, and have one fewer targets to worry about
  #if (currentIsHoldingTarget == true) && (currentIsPlacingTarget == true)
  #  neighbors[1].isHoldingTarget = false
  #  neighbors[1].numTargetsRemaining -= 1
  #  probabilities[1] = 1.0
  #  return distributions
  #end

  ### handle terminal state case first
  if isTerminalState(currentState)
    # only transition to same state deterministically
    #fill!(probabilities, 0.0)
    probabilities[1] = 1.0
    #copy!(neighbors[1], currentState)
    return distribution # when sample distributions in future, will only get the state in neighbors[1] - our terminal state
  end

  ### set neighbors/probabilities based on actions
#ACTION_SYMBOLS = [:moveForward, :turnLeft, :turnRight, :pickUpTarget, :placeTarget] # all possible robot actions (as symbols)

  # if moveForward
  if actionSymbol == :moveForward

    # get appropriate probability p
    pNextState = pomdp.PROBABILITY_MOVEMENT_SUCCESS

    ## if intended next state will take us out of bounds
    #if isOutOfBounds(pomdp, intendedNextState)

    # get all moveForwardPossibleNewLocations
    # - first/same state already in neighbors[1]
    # - intended state is in newLocations[2]
    # - front-left and front-right in newLocations[3 and 4]
    moveForwardPossibleNewLocations = getMoveForwardPossibleNewLocations(pomdp, currentState)

    # copy intendedNextState into second neighbor
    assert( isEqualStates(intendedNextState, moveForwardPossibleNewLocations[2]) )
    copy!(neighbors[2], moveForwardPossibleNewLocations[2])

    # copy 3rd and 4th into 3rd and 4th neighbors
    copy!(neighbors[3], moveForwardPossibleNewLocations[3])
    copy!(neighbors[4], moveForwardPossibleNewLocations[4])

    # 2nd probability is intended; divide the rest equally
    pOther = (1-pNextState)/3
    probabilities[1] = pOther
    probabilities[2] = pNextState
    probabilities[3] = pOther
    probabilities[4] = pOther

    ## if nextState is still in bounds
    ##if !isOutOfBounds(pomdp, intendedNextState)
    #  # copy intendedNextState into second neighbor
    #  copy!(neighbors[2], intendedNextState)
    #  # stay in same state wth prob=(1-p), move to new state with prob=(p)
    #  probabilities[1] = 1-pNextState
    #  probabilities[2] = pNextState

    ## otherwise, stay in same state deterministically
    #else
    #  probabilities[1] = 1.0
    #end

    ## with probability PROBABILITY_MOVEMENT_SUCCESS
    #if rand() <= pomdp.PROBABILITY_MOVEMENT_SUCCESS

    #  # move forward one step
    #  #nextState = getIntendedNextState(pomdp, currentState, action)
    #  #(newX, newY) = getForwardLocationGivenFacingDirection(pomdp, currentStateX, currentStateY, currentFacingDirection)
    #  # move only if nextState still in bounds
    #  if !isOutOfBounds(pomdp, intendedNextState)
    #    copy!(neighbors[1], intendedNextState)
    #    #neighbors[1].x = nextState.x #newX
    #    #neighbors[1].y = nextState.y #newY
    #  end
    #end

    ## all probability goes to being in state one step ahead (or not, if movement failed), with everything else fixed
    #probabilities[1] = 1.0

  # if turnLeft/turnRight
  elseif actionSymbol == :turnLeft || actionSymbol == :turnRight

    # get appropriate probability p
    pNextState = pomdp.PROBABILITY_MOVEMENT_SUCCESS

    # copy intendedNextState into second neighbor
    copy!(neighbors[2], intendedNextState)
    # stay in same state with prob=(1-p), move to new state with prob=(p)
    probabilities[1] = 1-pNextState
    probabilities[2] = pNextState

    ## with probability PROBABILITY_MOVEMENT_SUCCESS
    #if rand() <= pomdp.PROBABILITY_MOVEMENT_SUCCESS

    #  # face new intended direction successfully
    #  #nextState = getIntendedNextState(pomdp, currentState, action)
    #  #copy!(neighbors[1], nextState)
    #  copy!(neighbors[1], intendedNextState)#getIntendedNextState(podmp, currentState, action))
    #  #newFacingDirection = getIntendedNewFacingDirection(currentFacingDirection, action)
    #  #neighbors[1].facingDirection = newFacingDirection

    #end

    ## all probability goes to being in this new state (or same state if movement failed)
    #probabilities[1] = 1.0

  # if pickUpTarget
  elseif actionSymbol == :pickUpTarget

    # get appropriate probability p
    pNextState = pomdp.PROBABILITY_TARGET_INTERACTION_SUCCESS

    # if at base location, and not already holding target
    if (isAtBaseLocation(pomdp, currentState)) && (currentIsHoldingTarget == false) && (currentNumTargetsRemaining > 0)

      # copy intendedNextState into second neighbor
      copy!(neighbors[2], intendedNextState)
      # stay in same state with prob=(1-p), move to new state with prob=(p)
      probabilities[1] = 1-pNextState
      probabilities[2] = pNextState

    # otherwise, invalid action, and just deterministically stay in same state
    else
      probabilities[1] = 1.0
    end

    ## with probability PROBABILITY_TARGET_INTERACTION_SUCCESS
    #if rand() <= pomdp.PROBABILITY_TARGET_INTERACTION_SUCCESS

    #  # if at base location, and not already holding target
    #  if (isAtBaseLocation(pomdp, currentState)) && (currentIsHoldingTarget == false) && (currentNumTargetsRemaining > 0)
    #  #if ((currentX, currentY) == pomdp.BASE_LOCATION) && (currentIsHoldingTarget == false)

    #    # pick up target
    #    #copy!(neighbors[1], getIntendedNextState(pomdp, currentState, action))
    #    copy!(neighbors[1], intendedNextState)
    #    #neighbors[1].isHoldingTarget = true

    #  end
    #end

    ## all probability goes to being in this new state (or same state if interaction failed)
    #probabilities[1] = 1.0

  # if placeTarget
  elseif actionSymbol == :placeTarget

    # get appropriate probability p
    pNextState = pomdp.PROBABILITY_TARGET_INTERACTION_SUCCESS

    # get past, current, future targets
    pastTargetSpecifications, currentTargetSpecification, futureTargetSpecifications = getTargetSpecificationsPartitionedByPlaced(pomdp, currentState)

    # if already holding target and at least one target remaining for us to hold, and if correct location
    if (currentIsHoldingTarget == true) && (currentNumTargetsRemaining > 0) && (isEqualLocations(currentState, currentTargetSpecification))

      # copy intendedNextState into second neighbor
      copy!(neighbors[2], intendedNextState)
      # stay in same state with prob=(1-p), move to new state with prob=(p)
      probabilities[1] = 1-pNextState
      probabilities[2] = pNextState

    # otherwise, invalid action, and just deterministically stay in same state
    else
      probabilities[1] = 1.0
    end
    ## with probability PROBABILITY_TARGET_INTERACTION_SUCCESS
    #if rand() <= pomdp.PROBABILITY_TARGET_INTERACTION_SUCCESS

    #  # if already holding target and at least one target remaining for us to hold
    #  if (currentIsHoldingTarget == true) && (currentNumTargetsRemaining > 0) #&& (currentIsPlacingTarget == false)

    #    # place target
    #    copy!(neighbors[1], intendedNextState)
    #    #neighbors[1].isHoldingTarget = false
    #    #neighbors[1].numTargetsRemaining -= 1

    #  end
    #end

    ## all probability goes to being in this new state (or same state if interaction failed)
    #probabilities[1] = 1.0

  # otherwise, invalid action
  else
    # print and error
    @printf("POMDPs.transition: got invalid action. currentState=%s, action=%s\n", currentState, action)
    error("POMDPs.transition: got invalid action.")
  end

  # return full distribution, containing neighbors and probabilities
  return distribution
end

### Targets and obstacles - create targets and obstacles list
function createTargetAndObstacleSpecifications(worldConfig, worldSizeX, worldSizeY, numTotalTargets, numTotalObstacles)

  # print
  @printf("Building targets and obstacles for worldConfig:%s. Size:(%d, %d). Targets:%d. Obstacles:%d\n", worldConfig, worldSizeX, worldSizeY, numTotalTargets, numTotalObstacles)

  # instantiate
  targetSpecifications = Vector{TargetSpecification}() # list of targets for this pomdp world
  obstacleSpecifications = Vector{ObstacleSpecification}() # list of obstacles for this pomdp world

  # if :small world
  if WORLD_CONFIG == :small

    # config
    #BASE_LOCATION = (5, 5) # location of base where targets originate
    #WORLD_SIZE_X = 5 # number of blocks in world x dimension
    #WORLD_SIZE_Y = 5 # number of blocks in world y dimension
    #NUM_TOTAL_TARGETS = 4 # number of total targets the bot must place
    #NUM_TOTAL_OBSTACLES = 3 # number of total obstacles the bot must avoid

    # targets
    push!(targetSpecifications, TargetSpecification(3, 3, 1))
    push!(targetSpecifications, TargetSpecification(3, 4, 2))
    push!(targetSpecifications, TargetSpecification(4, 3, 3))
    push!(targetSpecifications, TargetSpecification(4, 4, 4))
    # obstacles
    push!(obstacleSpecifications, ObstacleSpecification(2, 4, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(2, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(3, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))

  # if :medium world
  elseif WORLD_CONFIG == :medium

    # config
    #BASE_LOCATION = (1, 1) # location of base where targets originate
    #WORLD_SIZE_X = 10 # number of blocks in world x dimension
    #WORLD_SIZE_Y = 10 # number of blocks in world y dimension
    #NUM_TOTAL_TARGETS = 6 # number of total targets the bot must place
    #NUM_TOTAL_OBSTACLES = 6 # number of total obstacles the bot must avoid

    # targets
    push!(targetSpecifications, TargetSpecification(5, 4, 1))
    push!(targetSpecifications, TargetSpecification(6, 4, 2))
    push!(targetSpecifications, TargetSpecification(7, 5, 3))
    push!(targetSpecifications, TargetSpecification(6, 6, 4))
    push!(targetSpecifications, TargetSpecification(5, 6, 5))
    push!(targetSpecifications, TargetSpecification(4, 5, 6))

    # obstacles
    push!(obstacleSpecifications, ObstacleSpecification(1, 4, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(2, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(4, 1, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(7, 2, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(8, 7, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(9, 7, OBSTACLE_ERROR_PENALTY_DEFAULT))

  elseif WORLD_CONFIG == :large

    # config
    #BASE_LOCATION = (20, 10) # location of base where targets originate
    #WORLD_SIZE_X = 20 # number of blocks in world x dimension
    #WORLD_SIZE_Y = 15 # number of blocks in world y dimension
    #NUM_TOTAL_TARGETS = 10 # number of total targets the bot must place
    #NUM_TOTAL_OBSTACLES = 10 # number of total obstacles the bot must avoid

    # targets
    push!(targetSpecifications, TargetSpecification(6, 7, 1))
    push!(targetSpecifications, TargetSpecification(7, 6, 2))
    push!(targetSpecifications, TargetSpecification(7, 8, 3))
    push!(targetSpecifications, TargetSpecification(8, 5, 4))
    push!(targetSpecifications, TargetSpecification(8, 9, 5))
    push!(targetSpecifications, TargetSpecification(9, 6, 6))
    push!(targetSpecifications, TargetSpecification(9, 8, 7))
    push!(targetSpecifications, TargetSpecification(10, 7, 8))
    push!(targetSpecifications, TargetSpecification(11, 6, 9))
    push!(targetSpecifications, TargetSpecification(11, 8, 10))
    push!(targetSpecifications, TargetSpecification(12, 5, 11))
    push!(targetSpecifications, TargetSpecification(12, 9, 12))
    push!(targetSpecifications, TargetSpecification(13, 6, 13))
    push!(targetSpecifications, TargetSpecification(13, 8, 14))
    push!(targetSpecifications, TargetSpecification(14, 7, 15))

    # obstacles
    push!(obstacleSpecifications, ObstacleSpecification(5, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(5, 9, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(10, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(10, 9, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(14, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
    push!(obstacleSpecifications, ObstacleSpecification(14, 9, OBSTACLE_ERROR_PENALTY_DEFAULT))

  # otherwise
  else
    # throw error - undefined scenario
    error("Undefined world for creating targets and obstacles.")
  end

  # assert
  assert( length(targetSpecifications) == numTotalTargets ) # must have appropriate total number of targets/specifications
  for (i, targetSpecification) in enumerate(targetSpecifications) # must be in order, or at least have one target specification for each order index
    assert( targetSpecification.orderIndex == i )
  end
  assert( length(obstacleSpecifications) == numTotalObstacles ) # must have appropriate total number of obstacles/specifications

  # print
  #@printf("Total number of targets: %d, total number of obstacles: %d\n", numTotalTargets, numTotalObstacles)

  # return
  return targetSpecifications, obstacleSpecifications
end
  ## create targets and obstacles for the (1,1) world with 1 target
  #if worldSizeX == 1 && worldSizeY == 2 && numTotalTargets == 1
  #  # target 1: (1, 2)
  #  push!(targetSpecifications, TargetSpecification(1, 2, 1))

  ## create targets and obstacles for the (2,2) world with 1 target
  #elseif worldSizeX == 2 && worldSizeY == 2 && numTotalTargets == 1
  #  # target 1: (1, 2)
  #  push!(targetSpecifications, TargetSpecification(1, 2, 1))

  ## create targets and obstacles for the (2,2) world with 2 targets
  #elseif worldSizeX == 2 && worldSizeY == 2 && numTotalTargets == 2
  #  # target 1: (1, 2)
  #  push!(targetSpecifications, TargetSpecification(1, 2, 1))
  #  # target 2: (2, 2)
  #  push!(targetSpecifications, TargetSpecification(2, 2, 2))

  ## create targets and obstacles for the (3,3) world with 2 targets
  #elseif worldSizeX == 3 && worldSizeY == 3 && numTotalTargets == 2
  #  # target 1: (1, 3)
  #  push!(targetSpecifications, TargetSpecification(1, 3, 1))
  #  # target 2: (3, 3)
  #  push!(targetSpecifications, TargetSpecification(3, 3, 2))

  ## create targets and obstacles for the (5,5) world with 2 targets
  #elseif worldSizeX == 5 && worldSizeY == 5 && numTotalTargets == 4 && numTotalObstacles == 3
  #  # targets
  #  push!(targetSpecifications, TargetSpecification(4, 4, 1))
  #  push!(targetSpecifications, TargetSpecification(4, 3, 2))
  #  push!(targetSpecifications, TargetSpecification(3, 4, 3))
  #  push!(targetSpecifications, TargetSpecification(3, 3, 4))
  #  # obstacles
  #  push!(obstacleSpecifications, ObstacleSpecification(1, 3, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(1, 4, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(2, 4, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #

  ## create targets and obstacles for the (5,5) world with 2 targets
  #elseif worldSizeX == 7 && worldSizeY == 7 && numTotalTargets == 2
  #  # target 1: (1, 7)
  #  push!(targetSpecifications, TargetSpecification(1, 7, 1))
  #  # target 2: (7, 7)
  #  push!(targetSpecifications, TargetSpecification(7, 7, 2))

  ## create targets and obstacles for the (10,10) world with 2 targets
  #elseif worldSizeX == 10 && worldSizeY == 10 && numTotalTargets == 2
  #  # target 1: (1, 10)
  #  push!(targetSpecifications, TargetSpecification(2, 10, 1))
  #  # target 10: (10, 10)
  #  push!(targetSpecifications, TargetSpecification(10, 10, 2))

  ## create targets and obstacles for the (10,10) world with 6 targets and 4 obstacles
  #elseif worldSizeX == 10 && worldSizeY == 10 && numTotalTargets == 6 && numTotalObstacles == 6
  #  # target 1: (5, 4)
  #  push!(targetSpecifications, TargetSpecification(5, 4, 1))
  #  # target 2: (6, 4)
  #  push!(targetSpecifications, TargetSpecification(6, 4, 2))
  #  # target 3: (7, 5)
  #  push!(targetSpecifications, TargetSpecification(7, 5, 3))
  #  # target 4: (6, 6)
  #  push!(targetSpecifications, TargetSpecification(6, 6, 4))
  #  # target 5: (5, 6)
  #  push!(targetSpecifications, TargetSpecification(5, 6, 5))
  #  # target 6: (4, 5)
  #  push!(targetSpecifications, TargetSpecification(4, 5, 6))

  #  # obstacles
  #  push!(obstacleSpecifications, ObstacleSpecification(1, 4, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(2, 5, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(4, 1, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(7, 2, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(8, 7, OBSTACLE_ERROR_PENALTY_DEFAULT))
  #  push!(obstacleSpecifications, ObstacleSpecification(9, 7, OBSTACLE_ERROR_PENALTY_DEFAULT))

### partition targetSpecifications as past/alreadyPlaced, present/currentTarget, and future/futureTargets based on state
function getTargetSpecificationsPartitionedByPlaced(pomdp::BotWorld, state::BotState)

  # init
  past = Vector{TargetSpecification}()
  current = Vector{TargetSpecification}()
  future = Vector{TargetSpecification}()

  # get number of targets already placed according to state
  numAlreadyPlacedTargets = pomdp.NUM_TOTAL_TARGETS - state.numTargetsRemaining

  # iterate
  for targetSpecification in pomdp.TARGET_SPECIFICATIONS
    # if already placed
    if targetSpecification.orderIndex <= numAlreadyPlacedTargets
      # add to past list
      push!(past, targetSpecification)
    # else if current target
    elseif targetSpecification.orderIndex == (numAlreadyPlacedTargets + 1)
      # add to current list
      push!(current, targetSpecification)
    # otherwise, future target
    else
      # add to future list
      push!(future, targetSpecification)
    end
  end

  # check
  assert( length(past) + length(current) + length(future) == pomdp.NUM_TOTAL_TARGETS )
  assert( length(current) == 1 )

  # return
  return past, current[1], future
end
#### get targetSpecification for state's current target
#function getCurrentTargetSpecification(pomdp::BotWorld, state::BotState)
#  # get number of targets already placed according to state
#  numAlreadyPlacedTargets = pomdp.NUM_TOTAL_TARGETS - state.numTargetsRemaining
#
#  # get the next target (that hasn't been placed)
#  for targetSpecification in pomdp.TARGET_SPECIFICATIONS
#    if targetSpecification.orderIndex == (numAlreadyPlacedTargets + 1)
#      return targetSpecification
#    end
#  end
#
#  # otherwise (if already placed all targets, for example), return empty
#  return nothing
#end
#### get vector of targetSpecifications that have already been placed according to state
#function getAlreadyPlacedTargetSpecifications(pomdp::BotWorld, state::BotState)
#
#  # init
#  alreadyPlacedTargetSpecifications = Vector{TargetSpecification}()
#
#  # get number of targets already placed according to state
#  numAlreadyPlacedTargets = pomdp.NUM_TOTAL_TARGETS - state.numTargetsRemaining
#
#  # get first numAlreadyPlacedTargets targetSpecifications
#  for targetSpecification in pomdp.TARGET_SPECIFICATIONS
#    if targetSpecification.orderIndex <= numAlreadyPlacedTargets
#      push!(alreadyPlacedTargetSpecifications, targetSpecification)
#    end
#  end
#
#  # return
#  return alreadyPlacedTargetSpecifications
#end
#### determine if passed-in state is at current target location state (doesn't check if holding target; do this elsewhere)
#function isAtCurrentTargetLocation(pomdp::BotWorld, state::BotState)
#
#  # get current state's current target
#  currentTargetSpecification = getCurrentTargetSpecification(pomdp, state)
#
#  # return true if at currentTargetSpecification's location
#  return (state.x == currentTargetSpecification.x) && (state.y == currentTargetSpecification.y)
#end

  #### - note that rewardStates we loop through must have holdingTarget=true meaning they are placeTarget location, not pickUpTarget location
  ## if at base location, then not at target location
  #if isAtBaseLocation(pomdp, state)
  #  return false
  #end
  ## for each rewardStatesActionsValues
  #for (rewardState, rewardAction, rewardValue) in rewardStatesActionsValues
  #  # check if currentState has rewardState location and numTargetsRemaining
  #  if (state.x == rewardState.x) && (state.y == rewardState.y) && (state.numTargetsRemaining == rewardState.numTargetsRemaining) && (rewardState.isHoldingTarget == true)
  #    return true
  #  end
  #end
  ## otherwise, return false
  #return false
#end

### Rewards - rewardStatesActionsValues
function createRewards(worldSizeX, worldSizeY, numTotalTargets, baseLocation, targetSpecifications, rewardPickUpTarget, rewardPlaceTarget)

  # instantiate
  rewardStatesActionsValues = Vector{Tuple{BotState, BotAction, Float64}}() # when in BotState and take BotAction, receive reward
  baseLocationX = baseLocation[1]
  baseLocationY = baseLocation[2]
  actionPickUpTarget = BotAction(:pickUpTarget)
  actionPlaceTarget = BotAction(:placeTarget)

  # for each targetSpecification
  for (i, targetSpecification) in enumerate(targetSpecifications)

    # compute number of targets that must be remaining BEFORE PLACING THIS TARGET AT ITS TARGET LOCATION based on targetSpecification's order index
    numTargetsRemaining = numTotalTargets - targetSpecification.orderIndex + 1

    # create rewardStateActionValue tuple for picking up this target from base
    push!(rewardStatesActionsValues, (BotState(baseLocationX, baseLocationY, :ANY_DIRECTION, numTargetsRemaining, false), actionPickUpTarget, rewardPickUpTarget))
    # create rewardStateActionValue tuple for placing at target location
    push!(rewardStatesActionsValues, (BotState(targetSpecification.x, targetSpecification.y, :ANY_DIRECTION, numTargetsRemaining, true), actionPlaceTarget, rewardPlaceTarget))

  end

  ## create rewardStateActionValue for ending state
  #push!(rewardStatesActionsValues, (BotState(baseLocationX, baseLocationY, :ANY_DIRECTION, 0, false)

  # assert
  assert( length(rewardStatesActionsValues) == numTotalTargets*2 ) # *2 since count for each pickup and placement

  # print
  @printf("Total number of reward states: %d\n", length(rewardStatesActionsValues))

  # return
  return rewardStatesActionsValues

end

### Rewards - R(s,a) model
function POMDPs.reward(pomdp::BotWorld, state::BotState, action::BotAction)

  ## get reward model
  #rewardStatesActionsValues = pomdp.REWARD_STATES_ACTIONS_VALUES
  # instantiate reward
  r = 0.0
  # get expected next state
  intendedNextState = getIntendedNextState(pomdp, state, action)
  # get actionSymbol
  actionSymbol = action.actionSymbol

  # if invalid state, then return penalty
  if !isValidState(pomdp, state)
    #@printf("Got invalid state: %s\n", state)
    return pomdp.PENALTY_INVALID_STATE
  end

  # if terminal state already, then no more reward
  if isTerminalState(state)
    return 0.0
  end

  # for each reward state-action-reward tuple
  for (rewardState, rewardAction, rewardValue) in pomdp.REWARD_STATES_ACTIONS_VALUES
    # if currently in rewardState and taking rewardAction
    if isEqualStates(state, rewardState, ignoreDirection=true) && isEqualActions(action, rewardAction)
      ## TODO: REMOVE
      #@printf("In reward state: state=%s, rewardState=%s\n", state, rewardState)
      # rewarded with rewardValue
      r += rewardValue
    end
  end

  # but penalty for walking on target location if already placed target there
  # - note: if always penalize regardless of placing there yet or not, then can get stuck in middle spinning around!
  pastTargetSpecifications, currentTargetSpecification, futureTargetSpecification = getTargetSpecificationsPartitionedByPlaced(pomdp, state)
  for pastTargetSpecification in pastTargetSpecifications
    # penalty if current state is on target location that has already been placed
    if isEqualLocations(state, pastTargetSpecification)
      r += pomdp.PENALTY_STANDING_ON_PAST_TARGET_LOCATION
    end
    # penalty if trying to move to target location that has already been placed. turning is ok, since we likely just placed this target, and might need to avoid moving forward into an enclosed area/trap
    if isEqualLocations(intendedNextState, pastTargetSpecification) && (actionSymbol != :turnLeft && actionSymbol != :turnRight)
      r += pomdp.PENALTY_MOVING_TO_PAST_TARGET_LOCATION
    end
  end
  #for targetSpecification in pomdp.TARGET_SPECIFICATIONS
  #  if isEqualLocations(state, targetSpecification)
  #    r += pomdp.PENALTY_MOVING_TO_TARGET_LOCATION
  #  end
  #  if isEqualLocations(intendedNextState, targetSpecification)
  #    r += pomdp.PENALTY_MOVING_TO_TARGET_LOCATION
  #  end
  #end

  # TODO: also penalty for ending on target too...
  # penalty for being on or trying to move onto obstacle location
  for obstacleSpecification in pomdp.OBSTACLE_SPECIFICATIONS
    # TODO: comment this out again?
    if isEqualLocations(state, obstacleSpecification)
      r += obstacleSpecification.errorPenalty
    end
    if isEqualLocations(intendedNextState, obstacleSpecification)
      r += obstacleSpecification.errorPenalty
    end
  end

  # penalty for movement
  if actionSymbol == :moveForward
    r += pomdp.PENALTY_MOVEMENT
  end

  # penalty for turn
  if (actionSymbol == :turnLeft) || (actionSymbol == :turnRight)
    r += pomdp.PENALTY_TURN
  end

  # penalty for trying to go out of bounds
  if isOutOfBounds(pomdp, intendedNextState)
    r += pomdp.PENALTY_OUT_OF_BOUNDS
  end

  # penalty for trying to pick up a target when not at the base location, or if already have one
  if actionSymbol == :pickUpTarget
    # if not at base location
    if !isAtBaseLocation(pomdp, state)
      r += pomdp.PENALTY_INVALID_PICKUPTARGET_NOTBASE
    end
    # if picking up when already have one
    if state.isHoldingTarget == true
      r += pomdp.PENALTY_INVALID_PICKUPTARGET_ALREADYHOLDING
    end
  end

  # penalty for placing a target at an incorrect location, or if don't have a target already
  if actionSymbol == :placeTarget
    # get partitioned targets based on current state
    pastTargetSpecifications, currentTargetSpecification, futureTargetSpecification = getTargetSpecificationsPartitionedByPlaced(pomdp, state)
    # if at incorrect location for current target
    if !isEqualLocations(state, currentTargetSpecification)
    #if !isAtTargetLocation(pomdp, state)
      r += pomdp.PENALTY_WRONG_PLACETARGET
    end
    # if placing but don't actually have a target on us
    if state.isHoldingTarget == false
      r += pomdp.PENALTY_INVALID_PLACETARGET
    end
  end

  # return
  return r
end

### observation initializer function
function POMDPs.create_observation(pomdp::BotWorld)
  baseLocation = pomdp.BASE_LOCATION
  return BotObservation(baseLocation[1], baseLocation[2])
end
### ObservationSpace
type BotObservationSpace <: POMDPs.AbstractSpace
  observations::Vector{BotObservation}
end
### functions returning observation space
function POMDPs.observations(pomdp::BotWorld)
  return BotObservationSpace([POMDPs.create_observation(pomdp)])
end
function POMDPs.observations(pomdp::BotWorld, state::BotState, observationSpace::BotObservationSpace)
  return observationSpace
end
### function returning iterator over observation space
function POMDPs.domain(observationSpace::BotObservationSpace)
  return observationSpace.observations
end
### observation distribution type
type BotObservationDistribution <: POMDPs.AbstractDistribution
  probabilities::Vector{Float64}
  observations::Vector{BotObservation}
end
### observation distribution initializer
function POMDPs.create_observation_distribution(pomdp::BotWorld)
  return BotObservationDistribution([1.0], [POMDPs.create_observation(pomdp)])
end
### observation pdf
function POMDPs.pdf(distribution::BotObservationDistribution, observation::BotObservation)
  return distribution.probabilities[1]
end
### sample from observation distribution
function POMDPs.rand!(rng::AbstractRNG, observation::BotObservation, distribution::BotObservationDistribution)
  categorical = Categorical(distribution.probabilities)
  #observation.xObs = distribution.observations[sampledObservationIndex]sampledObservation.xObs
  sampledObservationIndex = rand(rng, categorical)
  sampledObservation = distribution.observations[sampledObservationIndex]
  observation.xObs = sampledObservation.xObs
  observation.yObs = sampledObservation.yObs
  return observation
end
### observation model
function POMDPs.observation(pomdp::BotWorld, state::BotState, action::BotAction, distribution::BotObservationDistribution=POMDPs.create_observation_distribution(pomdp))
  probabilities = distribution.probabilities
  observations = distribution.observations
  observations[1] = BotObservation(state.x, state.y)
  probabilities[1] = 1.0
  return distribution
end

### BELIEFS
### initialization function
function POMDPs.create_belief(pomdp::BotWorld)
  @printf("length(POMDPs.domain(POMDPS.states(pomdp))) == %d\n", length(POMDPs.domain(POMDPs.states(pomdp))))
  return POMDPToolbox.DiscreteBelief(length(POMDPs.domain(POMDPs.states(pomdp))))
end
# initial belief is same as create_belief
function POMDPs.initial_belief(pomdp::BotWorld)
  @printf("length(POMDPs.domain(POMDPS.states(pomdp))) == %d\n", length(POMDPs.domain(POMDPs.states(pomdp))))
  return POMDPToolbox.DiscreteBelief(length(POMDPs.domain(POMDPs.states(pomdp))))
  #return POMDPToolbox.DiscreteBelief(1)
end

### miscellaneous functions
function POMDPs.n_states(pomdp::BotWorld)
  return pomdp.NUM_TOTAL_STATES
end
function POMDPs.n_actions(pomdp::BotWorld)
  return length(pomdp.ACTION_SYMBOLS)
end
function POMDPs.n_observations(pomdp::BotWorld)
  return 1
end
function POMDPs.discount(pomdp::BotWorld)
  return pomdp.DISCOUNT_FACTOR
end
# have to make modificatoins so that no 0 indexes can appear
function POMDPs.index(pomdp::BotWorld, state::BotState)
  #println(state)
  #return sub2ind( (pomdp.WORLD_SIZE_X, pomdp.WORLD_SIZE_Y, length(pomdp.FACING_DIRECTION_SYMBOLS), pomdp.NUM_TOTAL_TARGETS+1, 2), state.x, state.y, pomdp.FACING_DIRECTION_SYMBOLS_INDEX_DICT[state.facingDirection], state.numTargetsRemaining+1, Int(state.isHoldingTarget+1) )
  return pomdp.STATES_INDEXES_MAP[state]
end

### run

# print
println("Begin")
timeStart = time()

## create states
#states = createStates()

# create state space - has array .states
#stateSpace = createStateSpace()

# do any dynamic setup needed
NUM_MAX_STATE_NEIGHBORS = getMaxNumStateNeighbors()

# create state indexes map based on world params
_, statesIndexesMap = createStates(WORLD_SIZE_X, WORLD_SIZE_Y, FACING_DIRECTION_SYMBOLS, NUM_TOTAL_TARGETS, NUM_TOTAL_STATES)

# create targets and obstacles based on world params
targetSpecifications, obstacleSpecifications = createTargetAndObstacleSpecifications(WORLD_CONFIG, WORLD_SIZE_X, WORLD_SIZE_Y, NUM_TOTAL_TARGETS, NUM_TOTAL_OBSTACLES)
# create rewards based on world params
rewardStatesActionsValues = createRewards(WORLD_SIZE_X, WORLD_SIZE_Y, NUM_TOTAL_TARGETS, BASE_LOCATION, targetSpecifications, REWARD_PICKUP_TARGET, REWARD_PLACE_TARGET)
#rewardStates, rewardValues = createRewards()

# additional asserts
assert( (WORLD_SIZE_X > 0) && (WORLD_SIZE_Y > 0) )
assert( NUM_MAX_STATE_NEIGHBORS > 0 )
assert( (length(FACING_DIRECTION_SYMBOLS) == length(FACING_DIRECTION_ADDITION_VECTORS)) && (length(FACING_DIRECTION_SYMBOLS) == length(FACING_DIRECTION_SYMBOLS_INDEX_DICT)) )
assert( NUM_MOVE_FORWARD_POSSIBLE_NEW_LOCATIONS == 4 )
assert( (0 <= PROBABILITY_MOVEMENT_SUCCESS) && (PROBABILITY_MOVEMENT_SUCCESS <= 1) )
assert( (0 <= PROBABILITY_TARGET_INTERACTION_SUCCESS) && (PROBABILITY_TARGET_INTERACTION_SUCCESS <= 1) )
assert( REWARD_PICKUP_TARGET > 0 )
assert( REWARD_PLACE_TARGET > 0 )
assert( PENALTY_INVALID_STATE < 0 )
assert( PENALTY_MOVEMENT < 0 )
assert( PENALTY_TURN < 0 )
assert( PENALTY_OUT_OF_BOUNDS < 0 )
assert( PENALTY_INVALID_PICKUPTARGET_NOTBASE < 0 )
assert( PENALTY_INVALID_PICKUPTARGET_ALREADYHOLDING < 0 )
assert( PENALTY_INVALID_PLACETARGET < 0 )
assert( PENALTY_WRONG_PLACETARGET < 0 )
assert( PENALTY_STANDING_ON_PAST_TARGET_LOCATION <= 0 )
assert( PENALTY_MOVING_TO_PAST_TARGET_LOCATION < 0 )
assert( (0 <= DISCOUNT_FACTOR) && (DISCOUNT_FACTOR <= 1) )

# instantiate world
pomdp = BotWorld(
          WORLD_CONFIG,
          WORLD_SIZE_X,
          WORLD_SIZE_Y,
          BASE_LOCATION,
          NUM_TOTAL_STATES,
          NUM_MAX_STATE_NEIGHBORS,
          statesIndexesMap,
          NUM_TOTAL_TARGETS,
          ACTION_SYMBOLS,
          FACING_DIRECTION_SYMBOLS,
          FACING_DIRECTION_ADDITION_VECTORS,
          FACING_DIRECTION_SYMBOLS_INDEX_DICT,
          NUM_MOVE_FORWARD_POSSIBLE_NEW_LOCATIONS,
          targetSpecifications,
          obstacleSpecifications,
          rewardStatesActionsValues,
          #rewardStates,
          #rewardValues,
          PROBABILITY_MOVEMENT_SUCCESS,
          PROBABILITY_TARGET_INTERACTION_SUCCESS,
          REWARD_PICKUP_TARGET,
          REWARD_PLACE_TARGET,
          PENALTY_INVALID_STATE,
          PENALTY_MOVEMENT,
          PENALTY_TURN,
          PENALTY_OUT_OF_BOUNDS,
          PENALTY_INVALID_PICKUPTARGET_NOTBASE,
          PENALTY_INVALID_PICKUPTARGET_ALREADYHOLDING,
          PENALTY_INVALID_PLACETARGET,
          PENALTY_WRONG_PLACETARGET,
          PENALTY_STANDING_ON_PAST_TARGET_LOCATION,
          PENALTY_MOVING_TO_PAST_TARGET_LOCATION,
          DISCOUNT_FACTOR)

### DISCRETE VALUE ITERATION

if RUN_DISCRETE_VALUE_ITERATION
  # initialize the policy by passing in problem
  policy = DiscreteValueIteration.ValueIterationPolicy(pomdp)
  # initialize solver
  solver = DiscreteValueIteration.ValueIterationSolver(max_iterations=200, belres=1e-3)
  # solve for an optimal policy
  DiscreteValueIteration.solve(solver, pomdp, policy, verbose=true)
end

#### MCTS
#solver = MCTS.MCTSSolver(n_iterations=1000, depth=100, exploration_constant=50.0)
#policy = MCTS.MCTSPolicy(solver, pomdp)

### SARSOP
if RUN_SARSOP

  # initialize policy; argument is name of policy file
  policy = POMDPPolicy("bot.policy")
  # create .pomdpx file - format which the SARSOP backend reads in
  pomdpxFile = POMDPFile(pomdp, "bot.pomdpx")
  # initialize solver
  solver = SARSOPSolver()
  # solve
  solve(solver, pomdpxFile, policy)
  # print
  println("Done solving sarsop")
  println(size(SARSOP.alphas(policy)))

  # sarsop init
  observation = create_observation(pomdp)
  belief = initial_belief(pomdp)
  updater = DiscreteUpdater(pomdp)
  action_map = POMDPs.domain(POMDPs.actions(pomdp))

end

#### QMDP
## initializer solver
#solver = QMDP.QMDPSolver(max_iterations=100, tolerance=1e-3)
## initialize policy
#policy = QMDP.create_policy(solver, pomdp)
## solve
#QMDP.solve(solver, pomdp, policy, verbose=true)

# simulation setup
println("Beginning simulation")
rng = MersenneTwister(round(Int, rand()*1000))
currentState = getInitialBotState(pomdp)
nextState = getInitialBotState(pomdp)
iterations = 1
rewardTotal = 0.0

# visualization setup
include("Visualizer.jl")
visualizer = Visualizer(pomdp)
visualizer.init()
visualizerDraw(visualizer)
firstIteration = true

# loop simulation
while currentState.numTargetsRemaining > 0 #&& iterations < 100

  # get optimal action from currentState according to policy
  if RUN_DISCRETE_VALUE_ITERATION
    action = DiscreteValueIteration.action(policy, currentState)
  elseif RUN_SARSOP
    # get action from SARSOP policy
    ai = SARSOP.action(policy, belief)
    action = action_map[ai]
  end

  # compute reward
  reward = POMDPs.reward(pomdp, currentState, action)
  rewardTotal += reward

  # transition to next state
  transitionDistribution = POMDPs.transition(pomdp, currentState, action)
  POMDPs.rand!(rng, nextState, transitionDistribution)

  # sample new observation
  if RUN_SARSOP
    observationDistribution = POMDPs.observation(pomdp, currentState, action)
    rand!(rng, observation, observationDistribution)
    # update belief
    belief = update(updater, belief, action, observation)
  end

  # print
  #@printf("s=%s, a=%s, s'=%s, r=%0.2f\n", currentState, action, nextState, reward)
  if RUN_DISCRETE_VALUE_ITERATION
    @printf("s=%s, a=%s, s'=%s, r=%0.2f\n", currentState, action, nextState, reward)
  elseif RUN_SARSOP
    @printf("s=%s, a=%s, s'=%s, r=%0.2f, o=%s\n", currentState, action, nextState, reward, observation)
  end

  # draw states
  visualizerDraw(visualizer, currentState, action, nextState)
  PyPlot.pause(0.03)

  # if first iteration, then stop here
  if firstIteration == true
    firstIteration = false
    userInput = readline(STDIN)
  end

  # copy next state into current state
  copy!(currentState, nextState)

  # increment iterations
  iterations += 1

end
@printf("RewardTotal=%0.2f\n", rewardTotal)

# finish
timeFinish = time() - timeStart
@printf("Done. Took %0.2f seconds\n", timeFinish)
userInput = readline(STDIN)


using Astellarn
using Test

# First test perft on known positions for common bugs
# https://www.chessprogramming.org/Perft_Results

@testset "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" begin
    fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
    b = Astellarn.importfen(fen)
    @test cancastlekingside(b) == true
    @test cancastlekingside(b, BLACK) == true
    @test cancastlequeenside(b) == true
    @test cancastlequeenside(b, BLACK) == true
    @test ischeck(b) == false
    @test b.enpass == zero(UInt8)
    @test perft(b, 1) == 20
    @test perft(b, 2) == 400
    @test perft(b, 3) == 8902
    @test perft(b, 4) == 197281
    @test perft(b, 5) == 4865609
    @test perft(b, 6) == 119060324
end

@testset "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq -" begin
    fen = "r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq -"
    b = importfen(fen)
    @test cancastlekingside(b) == true
    @test cancastlekingside(b, BLACK) == true
    @test cancastlequeenside(b) == true
    @test cancastlequeenside(b, BLACK) == true
    @test ischeck(b) == false
    @test b.enpass == zero(UInt8)
    @test perft(b, 1) == 48
    @test perft(b, 2) == 2039
    @test perft(b, 3) == 97862
    @test perft(b, 4) == 4085603
    @test perft(b, 5) == 193690690
    #@test perft(b, 6) == 8031647685
end

@testset "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - -" begin
    fen = "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - -"
    b = importfen(fen)
    @test cancastlekingside(b) == false
    @test cancastlekingside(b, BLACK) == false
    @test cancastlequeenside(b) == false
    @test cancastlequeenside(b, BLACK) == false
    @test ischeck(b) == false
    @test b.enpass == zero(UInt8)
    @test perft(b, 1) == 14
    @test perft(b, 2) == 191
    @test perft(b, 3) == 2812
    @test perft(b, 4) == 43238
    @test perft(b, 5) == 674624
    @test perft(b, 6) == 11030083
end

@testset "r3k2r/Pppp1ppp/1b3nbN/nP6/BBP1P3/q4N2/Pp1P2PP/R2Q1RK1 w kq - 0 1" begin
    fen1 = "r3k2r/Pppp1ppp/1b3nbN/nP6/BBP1P3/q4N2/Pp1P2PP/R2Q1RK1 w kq - 0 1"
    fen2 = "r2q1rk1/pP1p2pp/Q4n2/bbp1p3/Np6/1B3NBn/pPPP1PPP/R3K2R b KQ - 0 1"
    b1 = importfen(fen1)
    b2 = importfen(fen2)
    @test cancastlekingside(b1, BLACK) == true
    @test cancastlequeenside(b1, BLACK) == true
    @test cancastlekingside(b1, WHITE) == false
    @test cancastlequeenside(b1, WHITE) == false
    @test cancastlekingside(b2, WHITE) == true
    @test cancastlequeenside(b2, WHITE) == true
    @test cancastlekingside(b2, BLACK) == false
    @test cancastlequeenside(b2, BLACK) == false
    @test perft(b1, 1) == perft(b2, 1) == 6
    @test perft(b1, 2) == perft(b2, 2) == 264
    @test perft(b1, 3) == perft(b2, 3) == 9467
    @test perft(b1, 4) == perft(b2, 4) == 422333
    @test perft(b1, 5) == perft(b2, 5) == 15833292
end

@testset "SEE k7/8/6b1/3p4/4r3/3PQ3/5N2/K7 w - - 0 1" begin
    fen = "k7/8/6b1/3p4/4r3/3PQ3/5N2/K7 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(21, 28)
    @test static_exchange_evaluator(b, move, t) == true
    b.turn = !b.turn
    move = Move(28, 20)
    @test static_exchange_evaluator(b, move, t) == true
    move = Move(28, 36)
    @test static_exchange_evaluator(b, move, t) == false
end

@testset "SEE 7k/3r4/3q4/8/3p4/2BK4/8/8 w - - 0 1" begin
    fen = "7k/3r4/3q4/8/3p4/2BK4/8/8 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(22, 29)
    @test static_exchange_evaluator(b, move, t) == false
    fen = "3r3k/8/8/2p1p3/3p4/2BK4/2N5/8 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(22, 29)
    @test static_exchange_evaluator(b, move, t) == false
    fen = "3r3k/3r4/3q4/8/3p4/2BK4/1QN5/8 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(22, 29)
    @test static_exchange_evaluator(b, move, t) == true
    fen = "3r3k/8/8/2p1p3/3n4/2BKP3/8/8 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(22, 29)
    @test static_exchange_evaluator(b, move, t) == true
    fen = "k7/8/8/8/1rrpK3/3R4/8/8 w - - 0 1"
    b = importfen(fen)
    t = 0
    move = Move(21, 29)
    @test static_exchange_evaluator(b, move, t) == false
end

@testset "Wider SEE tests" begin
    fen_data = [
    ("4R3/2r3p1/5bk1/1p1r3p/p2PR1P1/P1BK1P2/1P6/8 b - -", Move(33, 26), 0),
    ("4R3/2r3p1/5bk1/1p1r1p1p/p2PR1P1/P1BK1P2/1P6/8 b - -", Move(33, 26), 0),
    ("4r1k1/5pp1/nbp4p/1p2p2q/1P2P1b1/1BP2N1P/1B2QPPK/3R4 b - -", Move(26, 19), 0),
    ("2r1r1k1/pp1bppbp/3p1np1/q3P3/2P2P2/1P2B3/P1N1B1PP/2RQ1RK1 b - -", Move(45, 36), SEE_VALUES[1]),
    ("7r/5qpk/p1Qp1b1p/3r3n/BB3p2/5p2/P1P2P2/4RK1R w - -", Move(4, 60), 0),
    ("6rr/6pk/p1Qp1b1p/2n5/1B3p2/5p2/P1P2P2/4RK1R w - -", Move(4, 60), -SEE_VALUES[4]),
    ("7r/5qpk/2Qp1b1p/1N1r3n/BB3p2/5p2/P1P2P2/4RK1R w - -", Move(4, 60), -SEE_VALUES[4]),
    ("6RR/4bP2/8/8/5r2/3K4/5p2/4k3 w - -", Move(51, 59, __QUEEN_PROMO), SEE_VALUES[3] - SEE_VALUES[1]),
    ("6RR/4bP2/8/8/5r2/3K4/5p2/4k3 w - -", Move(51, 59, __KNIGHT_PROMO), SEE_VALUES[2] - SEE_VALUES[1]),
    ("7R/5P2/8/8/8/3K2r1/5p2/4k3 w - -", Move(51, 59, __QUEEN_PROMO), SEE_VALUES[5] - SEE_VALUES[1]),
    ("7R/5P2/8/8/8/3K2r1/5p2/4k3 w - -", Move(51, 59, __BISHOP_PROMO), SEE_VALUES[3] - SEE_VALUES[1]),
    ("7R/4bP2/8/8/1q6/3K4/5p2/4k3 w - -", Move(51, 59, __ROOK_PROMO), -SEE_VALUES[1])
    ]
    for fd in fen_data
        @test static_exchange_evaluator(importfen(fd[1]), fd[2], fd[3]) == true
    end
end

@testset "Symmetry eval" begin
    fen = "2q5/b2k3P/4pn2/3pr3/3PR3/4PN2/B2K3p/2Q5 w - - 0 1"
    b = importfen(fen)
    pktable = PawnKingTable()
    eval_white = evaluate(b, pktable)
    b.turn = !b.turn
    eval_black = evaluate(b, pktable)
    @test abs(eval_black) == abs(eval_white)

    fen1 = "k7/P6n/KP5r/7b/7p/8/8/5RBN w - - 0 1"
    b = importfen(fen1)
    pktable = PawnKingTable()
    eval_white = evaluate(b, pktable)
    ei, ea = initEvalInfo(b)
    scorewk = evaluate_knights(b, ei, ea)
    scorewb = evaluate_bishops(b, ei, ea)
    scorewr = evaluate_rooks(b, ei, ea)
    scorewq = evaluate_queens(b, ei, ea)
    scorewki = evaluate_kings(b, ea)
    scorewpin = evaluate_pins(b)
    scorewspace = evaluate_space(b, ea)
    scorewthreat = evaluate_threats(b, ei, ea)
    scorewpass = evaluate_passed(b, ea)
    scorewinit = evaluate_initiative(b, ea.passed, 400)
    fen2 = "5rbn/8/8/7P/7B/kp5R/p6N/K7 b - - 0 1"
    b = importfen(fen2)
    eval_black = evaluate(b, pktable)
    ei, ea = initEvalInfo(b)
    scorebk = evaluate_knights(b, ei, ea)
    scorebb = evaluate_bishops(b, ei, ea)
    scorebr = evaluate_rooks(b, ei, ea)
    scorebq = evaluate_queens(b, ei, ea)
    scorebki = evaluate_kings(b, ea)
    scorebpin = evaluate_pins(b)
    scorebspace = evaluate_space(b, ea)
    scorebthreat = evaluate_threats(b, ei, ea)
    scorebpass = evaluate_passed(b, ea)
    scorebinit = evaluate_initiative(b, ea.passed, 400)

    @test scorebk == -scorewk
    @test scorebb == -scorewb
    @test scorebr == -scorewr
    @test scorebq == -scorewq
    @test scorebki == -scorewki
    @test scorebpin == -scorewpin
    @test scorebspace == -scorewspace
    @test scorebthreat == -scorewthreat
    @test scorebpass == -scorewpass
    @test scorebinit == -scorewinit
    @test eval_black == -eval_white
end


function ShowFGENWindow(fgen, fgen_conf, rev_state_dict, refresh_cnt, base)
	CImGui.Begin("GFG3015 Function Generator")
	# DRAW DISPLAY INFO

	CImGui.Text("Func:"), CImGui.SameLine(), CImGui.Text(fgen_conf.crt_func)
	CImGui.Text("Freq:")
	CImGui.SameLine(), CImGui.Text(fgen_conf.freq)
	CImGui.SameLine(), CImGui.Text("Hz")
	CImGui.Text("Ampl:")
	CImGui.SameLine(), CImGui.Text(fgen_conf.amplit)
	CImGui.SameLine(), CImGui.Text(fgen_conf.amplit_unit)
	CImGui.Text("Offs:")
	CImGui.SameLine(), CImGui.Text(fgen_conf.offs)
	CImGui.SameLine(), CImGui.Text(fgen_conf.amplit_unit)
	CImGui.Text("Duty:")
	CImGui.SameLine(), CImGui.Text(fgen_conf.duty)
	CImGui.SameLine(), CImGui.Text("%")

	CImGui.Text("EDIT ")
	@c CImGui.InputDouble("", &fgen_conf.crt_value, 0.01, 1.0, "%.4f")

	CImGui.Button(" FUNC ") && begin
		# toggle between the three functions
		toggle_func!(fgen, fgen_conf)
	end
	CImGui.SameLine(),CImGui.Button(" AMPL ") && begin
		set_amplit(fgen, fgen_conf.crt_chan, fgen_conf.crt_value)
	end
	CImGui.SameLine(),CImGui.Button(" DUTY ") && begin
		set_duty(fgen, fgen_conf.crt_chan, fgen_conf.crt_value)
	end

	CImGui.Button(" FREQ ") && begin
		# toggle between the three functions
		set_freq(fgen, fgen_conf.crt_chan, fgen_conf.crt_value)
	end
	CImGui.SameLine(),CImGui.Button("OFFSET") && begin
		set_offs(fgen, fgen_conf.crt_chan, fgen_conf.crt_value)
	end
	CImGui.SameLine(),CImGui.Button("MOD ON") && begin
		# don't need
	end

	CImGui.SameLine()

	# Draw REFRESH
	CImGui.SameLine()
	fgen_conf.refresh_btn, pressed = draw_toggle_button(fgen_conf.refresh_btn)
	fgen_conf.refresh_btn.state && @async update_fgen_conf!(fgen_conf, fgen, refresh_cnt, base)

	CImGui.End()
	return fgen_conf
end


include("../src/OnePhase.jl")

#using JLD, OnePhase, NLPModels, CUTEst
using JLD, NLPModels, CUTEst
USE_IPOPT = false
if USE_IPOPT
  using Ipopt
end
include("ipopt.jl")
include("plots.jl")
include("summary.jl")
include("process_results.jl")
include("run_and_store.jl")

function if_mkdir(dir::String)
  if !isdir(dir)
     mkdir(dir)
  end
end;


"""
  Model(expr, parts::Vector{FormulaTerm})

Representation of a multipart model composed of a vector of FormulaTerm's.
It is inspired by the [https://github.com/matthieugomez/FixedEffectModels.jl](FixedEffectModels.jl) package
"""
struct Model
  expr
  parts::Vector{FormulaTerm}
end

"""
    @model(formula1, formula2, ...)

Parses a sequence of comma-separated FormulaTerm's representing a multipart model.
"""
macro model(args...)
  Expr(:call, :model_helper, (esc(Base.Meta.quot(a)) for a in args)...)
end

function model_helper(args...)
  expr = args
  parts = FormulaTerm[]
  for part in args
    push!(parts, @eval(@formula($(part.args[2]) ~ $(part.args[3]))))
  end
  Model(expr,parts)
end

function Base.show(io::IO, m::Model)
  print(io, "$(length(m.parts))-part model: [")
  notfirst = false
  for p in m.parts
    if notfirst
      print(io, ", ")
    end
    print(io, p)
    notfirst = true
  end
  print(io, "]")
  nothing
end

function mergerhsterms(a, b)
  terms = union(a, b)

  ina = findall((in)(a), terms)
  inb = findall((in)(b), terms)

  (terms...,), ina, inb
end

"maps inzero and inpos to ModelMatrix columns (important with factor variables)"
function mapins(inzero, inpos, appliedschema)
  mappedinzero = Int[]
  mappedinpos = Int[]
  from = 1
  to = 0
  for t=1:length(appliedschema)
    to = from + width(appliedschema[t]) - 1
    if t in inzero
      [push!(mappedinzero, i) for i in from:to]
    end
    if t in inpos
      [push!(mappedinpos, i) for i in from:to]
    end
    from = to + 1
  end
  mappedinzero, mappedinpos
end

"""
  getrhsterms(m, lhs)

Returns the right-hand-side Terms of the part of model `m` with left-hand-side
symbol `lhs`.
"""
function getrhsterms(m::Model, lhs::Symbol)
  ix = findfirst(p->p.lhs==term(lhs),m.parts)
  @assert ix > 0 "The model is missing a formula with $lhs on its left-hand-side."
  f = m.parts[ix]
  f.rhs
end

StatsModels.missing_omit(data::T, formula::TupleTerm) where T<:ColumnTable =
    missing_omit(NamedTuple{tuple(termvars(formula)...)}(data))

function StatsModels.modelcols(trms::TupleTerm, data, counts::AbstractMatrix;
    model::Type{M}=DCR, contrasts=Dict{Symbol,Any}()) where M

  cols = columntable(data)
  cols, nonmissing = missing_omit(cols, trms)
  s = schema(trms, cols, contrasts)
  as = apply_schema(trms, s, M)
  modelcols(as, nonmissing, cols, counts)
end

function StatsModels.modelcols(as::TupleTerm, nonmissing, cols, counts::AbstractMatrix)

  covars = modelcols(MatrixTerm(as), cols)

  if !all(nonmissing)
    counts = counts[nonmissing,:]
  end

  covars, counts, as
end

"Similar to StatsModels.TableRegressionModel, but also holds a counts matrix and Model"
struct TableCountsRegressionModel{M,D,C} <: RegressionModel
  model::M    # actual fitted model (result of fit())
  data::D     # data table used
  counts::C   # counts matrix used
  f::Model    # @model specification
  schema      # applied schema
  sprojdir::Union{Nothing,Symbol}   # lhs variable in forward regression
end
TableCountsRegressionModel(model, data, counts, f, schema) = TableCountsRegressionModel(model, data, counts, f, schema, nothing)

function Base.show(io::IO, model::TableCountsRegressionModel)
    println(io, typeof(model))
    println(io)
    println(io, model.f)
end