Main code (needs tweaking)