clean up state, make functions only take classes map