diff --git a/Client.java b/Client.java
index 5baa331..a22160c 100755
--- a/Client.java
+++ b/Client.java
@@ -272,6 +272,19 @@ public class Client {
 			userInput = stdIn.readLine();
 			output.println(userInput);
 			
+			int diff = 0;
+			in.read(b, 0, 256);
+			temp = String.valueOf(b).trim();
+			System.out.print(temp);
+			while(diff < 1 || diff > 3) {
+				userInput = stdIn.readLine();
+				diff = Integer.parseInt(userInput);
+				if(diff < 1 || diff > 3) {
+					System.out.print("Invalid difficulty.\n" + temp);
+				}
+			}
+			output.println(userInput);
+			
 			while(!end) {
 				char[] buffer = new char[256];
 				in.read(buffer, 0, 256);
diff --git a/Server.cpp b/Server.cpp
index ef0d679..bf84d43 100644
--- a/Server.cpp
+++ b/Server.cpp
@@ -118,18 +118,15 @@ int main(int argc, char *argv[])
 	cout << "OK" << endl;
 	bzero(info,256); //Resets info back to normal
 	
-	/*//Becca- once your AI code is working, use this code:
-	
-	string difficulty_select = "Choose Difficulty\n1. Easy\n2. Medium\n4. Hard\nEnter choice: \n";
+	string difficulty_select = "Choose Difficulty\n1. Easy\n2. Medium\n3. Hard\nEnter choice: \n";
 	write(newsockfd, difficulty_select.c_str(), difficulty_select.length());
 	n = read(newsockfd,info,255);
-	istringstream convert(info);
-	convert >> choice_difficulty;
+	istringstream con(info);
+	int choice_difficulty = 1;
+	con >> choice_difficulty;
 	cout << "OK" << endl;
 	bzero(info,256);
 	
-	*/
-	
 	while(true) {
 		 
 		 while(gameOver != true)
@@ -162,15 +159,12 @@ int main(int argc, char *argv[])
 			
 			while(e.getBoard()->getTurn() == 'X' )
 			{
-				e.easyAI();
-				/*Becca- once you finish your AI, uncomment this out and remove the line directly above this
 				if(choice_difficulty == 1)
-					easyAI();
+					e.AI(3);
 				else if(choice_difficulty == 2)
-					mediumAI();
+					e.AI(5);
 				else
-					hardAI();
-				*/
+					e.AI(7);
 			}
 			
 			gameOver = e.getBoard()->isGameOver();