still need to add remove individual number

PhoneBook.java

@@ -1,47 +1,52 @@
  
-import java.util.TreeMap;
+import java.util.*;
 /**
  * Created by leon on 1/23/18.
  */
 public class PhoneBook {
-    private TreeMap<String, String> phoneBook;
-    private TreeMap<String, String> reversePhoneBook;
-    
+    private TreeMap<String, ArrayList<String>> phoneBook;
     public PhoneBook() {
-        phoneBook = new TreeMap<String, String>();
-        reversePhoneBook = new TreeMap<String, String>();
+        phoneBook = new TreeMap<String, ArrayList<String>>(); 
     }
     
     public void add(String name, String phoneNumber) {
-        phoneBook.put(name, phoneNumber);
-        reversePhoneBook.put(phoneNumber, name);
+        ArrayList<String> input = phoneBook.containsKey(name) ? 
+                                  phoneBook.get(name)
+                                  :                         
+                                  new ArrayList<String>();
+        input.add(phoneNumber);
+        phoneBook.put(name, input);
     }
     
-    public void remove(String name) {
-        String phoneNumber = phoneBook.get(name);
+    public void removeEntry(String name) {
         phoneBook.remove(name);
-        reversePhoneBook.remove(phoneNumber);
     }
     
     public String lookup(String name) {
-        return phoneBook.get(name) != null ? 
-                                   phoneBook.get(name)
-                                   : 
-                                   "N/A";
+        return phoneBook.get(name)
+                        .toString()
+                        .replaceAll("[\\[\\]]", "");
     }
     
     public String reverseLookup(String phoneNumber) {
-        return reversePhoneBook.get(phoneNumber) != null ? 
-                                                 reversePhoneBook.get(phoneNumber)
-                                                 :
-                                                 "N/A";
+        String output = "";
+        Set<String> keys = phoneBook.keySet();
+        for (String key : keys) {
+            if (phoneBook.get(key).contains(phoneNumber)) {
+                output = key;
+            }
+        }
+        return output.equals("") ? "No associated name." : output;
     }
     
-    public TreeMap getPhoneBook() {
-        return phoneBook; 
+    public String display() {
+       return phoneBook.toString().replaceAll("[{}=,]", "")
+                                  .replaceAll("\\[", " - ")
+                                  .replaceAll("\\]", "\n");
+       
     }
     
-    public TreeMap getReversePhoneBook() {
-        return reversePhoneBook;
+    public TreeMap getPhoneBook() {
+        return phoneBook; 
     }
 }

PhoneBookTest.java

@@ -33,17 +33,15 @@ public class PhoneBookTest
        assertEquals(expectedNumber, actual);
     }
     
-    @Test
-    public void testRemovePhoneBookEntry_And_Lookup()
+    @Test(expected = NullPointerException.class)
+    public void testRemoveEntry_And_Lookup()
     {
        PhoneBook test = new PhoneBook();
        String name = "Jenny";
-       String wrongName = "Tommy Tutone";
        String number = "(302)-867-5309";
        test.add(name, number);
-       
-       String actual = test.lookup(wrongName); 
-       assertEquals("N/A", actual);
+       test.removeEntry(name);
+       test.lookup(name);   
     }
     
     @Test
@@ -57,5 +55,14 @@ public class PhoneBookTest
        String actual = test.reverseLookup(phoneNumber);
        assertEquals(expectedName, actual);
      
-    }   
+    } 
+    
+    @Test
+    public void testDisplay() {
+       PhoneBook test = new PhoneBook();
+       test.add("hey", "302");
+       test.add("yo", "521");
+       test.add("hey", "498");
+       System.out.println(test.display());
+    }
 }