updated the code a lot