fix(stages): Code and style improvements

Add new styling for stages
Improve some parts of the code