Optimized code in A, moved certain functions to utils.py