working now

need to add more test cases