Got rid of bugs and added tests