Rewritten most of the code