Fixed a small bug ... still debugging