Fixing some bug. Still one bug (tested)