Clean and restructure code