Added most logic. Nearly there but still a bug.