Code clean-up and better parallelization

Removed a lot of old code and cleaned up training and testing.