Fixing code , work in progress