Added some fixes / improvements