improve class structure