testenv = Environment2D()

# Add a shape to the 2D space.
testenv.add_shapes( Line( np.asarray([-1,0]), np.asarray([1,0]), np.asarray([1,0,0,1])) )
# Set the agent's position to right of the circle.
testenv.set_position( np.asarray([0.0,2.0]) )
# Set rotation quanta in degrees.
testenv.set_rot_amount( 30 )
# Set the Forward and Right directions for the agent.
testenv.set_directions( np.asarray([0.0,-1.0]), np.asarray([1.0,0.0]) )
# Set the movement amount to 0.1 units per action.
testenv.set_move_amount( 0.4 )

#outputss = sample_lines( testenv )
outputss = sample_paths( testenv )

# N_xTxWxC
obs = np.asarray( [ [ output[0][0] for output in outputs ] for outputs in outputss ])
# N_xTxA
acs = np.asarray( [ [ to_one_hot(output[1],7) for output in outputs ] for outputs in outputss ])
# N_xTx2
poss = np.asarray([ [ output[2] for output in outputs ] for outputs in outputss ])
# N_xTx2
dirs = np.asarray([ [ output[3] for output in outputs ] for outputs in outputss ])
# N_xTxW
uvs = np.asarray( [ [ np.linspace( -1, 1, obs.shape[2] ) for output in outputs ] for outputs in outputss ])

states = np.concatenate([poss, dirs], axis=2)

N_ = obs.shape[0]
testenv = Environment2D()

# Add a shape to the 2D space.
testenv.add_shapes( Line( np.asarray([-0.5,0.5]), np.asarray([1,0.5]), np.asarray([1,0,0,1])) )
testenv.add_shapes( Line( np.asarray([-2.3,-0.5]), np.asarray([0.3,-0.5]), np.asarray([0,1,0,1])) )
# Set the agent's position to right of the circle.
testenv.set_position( np.asarray([0.0,2.0]) )
# Set rotation quanta in degrees.
testenv.set_rot_amount( 30 )
# Set the Forward and Right directions for the agent.
testenv.set_directions( np.asarray([0.0,-1.0]), np.asarray([1.0,0.0]) )
# Set the movement amount to 0.1 units per action.
testenv.set_move_amount( 0.4 )

#outputss = sample_lines( testenv )
outputss = sample_paths( testenv, episode_length=7 )

# N_xTxWxC
obs = np.asarray( [ [ output[0][0] for output in outputs ] for outputs in outputss ])
# N_xTxA
acs = np.asarray( [ [ to_one_hot(output[1],7) for output in outputs ] for outputs in outputss ])
# N_xTx2
poss = np.asarray([ [ output[2] for output in outputs ] for outputs in outputss ])
# N_xTx2
dirs = np.asarray([ [ output[3] for output in outputs ] for outputs in outputss ])
# N_xTxW
uvs = np.asarray( [ [ np.linspace( -1, 1, obs.shape[2] ) for output in outputs ] for outputs in outputss ])
# N_xTxS
states = np.concatenate([poss, dirs], axis=2)

N_ = obs.shape[0]