Just some more bug fixing