Another attempt at implementing it right
Fix bugs