Revert "Aligned code with current staging"

This reverts commit 4f57f3d12b21853373344159e016a018f9b4bbbd.