Few fixes/cleanup