Code cleaning and bug fixing