Rework code & fix issues after gathering feedback