Fix class cast problem

other cleanup