cleaned up code, and I think it finally works