Added logic; still buggy.