Feature/fix bugs