Fix and re-arrange production code