Slightly changed code structure and stuff