Fix bugs (coming close to working :))