import gym
import numpy as np
from simple_pid import PID 
import matplotlib.pyplot as plt
env = gym.make('CartPole-v0')
# PID testing with gym
pars=[0.3,0.01,0.1]
setpoint = 0;
pid = PID(pars[0], pars[1], pars[2], setpoint)
#pid.output_limits = (0, 1) # Limit output to the range of action space
pid.sample_time = 0.01 # new output will only be calculated when sample_time seconds has passed
observation = env.reset() # reset the simulation environment
simulation_time = 199; # set the simulation time
log_err = np.zeros(simulation_time) # log the error signal
abs_force_sum = 0
# main loop
for t in range(simulation_time):
	env.render() # render the simulation
	cartpos = observation[0] # position of the cart
	angle = observation[2]*(180/np.pi) # angle of the pole
	err = setpoint - angle; # error calculation
	log_err[t] = err; # log the error for visualization
	force_k = pid(err);
	abs_force_sum += np.absolute(force_k)
	observation, reward, done, info = env.step(force_k) # do the action, get the measurements
	if done: # if Pole Angle is more than +-12 deg or Cart Position is more than +-2.4 (center of the cart reaches the edge of the display) the simulation ends
		break
abs_error_sum = np.sum(np.absolute(log_err))
print(abs_error_sum)
print(abs_force_sum)
plt.plot(log_err[1:t])
plt.show()