Replaced map.flatten with flat_map

Minor refactor: replaced map.flatten with flat_map