cleaned up code a bit. added more features