reorganised classes
minor import fix