ENH: improve script

... additional improvements + fixes.