Revert "Refactor for simplicity"

This reverts commit 8ca8567abc0fa3dc55b89ea5b5a0ea1b57e2162a.