Fixed python implementation of energy calculation (SLOW)