Added method, cleaned up code