Reworked most of the code