More code modfications