Restructure of code