Restructured classes to make more sense