Cleaned up (a bit) and added some comments and todos