Fixing mistakes and cleaning the implementation