Better code and bug fix :)