Almost working (fix the mess)