Code simplification + a few fixes