clean up, and add posible feature