[TASK] Cleaned up code