Working code. Should implement other options