Main class code cleanup