       updateScreen();
     }
     else {
+      
       // Figure out which colour counter has been dropped
-      char counter;
-      if ((btns & 0b1111111) == col) counter = 'O'; 
-      else counter = 'X';
-      // Has human chosen a colour yet?
-      if (human == '-') {
-        // Human has now chosen a colour and robot chooses the other
-        human = counter;
-        if (human == 'X') {
-          humanColour = strip.Color(32, 32, 0);
-          robot = 'O'; 
-          robotColour = strip.Color(0, 0, 64);
-        }
-        else {
-          humanColour = strip.Color(0, 0, 64);
-          robot = 'X';
-          robotColour = strip.Color(32, 32, 0);
-        }
-      }
-      // Has human played with robot's colour?
-      if (counter != human) {
+      byte counter;
+      if ((btns & 0b1111111) == col) counter = blue; else counter = yellow;
+      
+      // Have player colours been chosen yet?
+      if (playerColour[human] == 0 && playerColour[robot] == 0) {
+        // Colours have now been chosen
+        playerColour[human] = counter;
+        playerColour[robot] = 1 - counter;
+      }
+      
+      // Has the correct colour token been dropped?
+      if (counter != playerColour[whosTurn]) {
         strcpy(screen[0], "Cheat!");
         updateScreen();
       }
       else {
-        // Human has played correct colour. Figure out which column
+        
+        // The correct colour has bee dropped. Figure out which column
         for (byte i=0; i<=6; i++) {
           strip.setPixelColor(i, 0);
           if (bit(i) == col) {


void loop() {
  char c;
  byte best;
  int currentScore;
  unsigned int btnsNow = 0;
  // Read all buttons in the 7 cols x 2 rows matrix
  for (byte row=0; row<=1; row++) {
    for (byte col=0; col<=6; col++) {

          }
        }
        // Drop the counter
        dropCounter(board, token[whosTurn], c);
        displayBoard(board);
        updateScreen();
        
        // Has the game been won?
        currentScore = score(board, 'O', 'X');
        if (currentScore == 1000) {
          strcpy(screen[0], "You Won!");
          updateScreen();
          while (true) {}
        }
        else if (currentScore == -1000) {
          strcpy(screen[0], "You lost!");
          updateScreen();
          while (true) {}
        }
        else {
          
          // The game continues
          if (whosTurn == human) {
            
            // Indicate the move the human made
            strip.setPixelColor(c, neoPixelColour[playerColour[human]]);
            strip.show();

            // The robot takes a go
            whosTurn = robot;
            strcpy(screen[0], "My go:");
            updateScreen();
            
            // Robot decides its move
            currentScore = bestMove(board, 'X', 'O', robotMove, 4);
            digitalWrite(LED, LOW);
            
            // Robot indicates its chosen move and waits for counter to drop
            strip.setPixelColor(robotMove, neoPixelColour[playerColour[robot]]);
            strip.show();
          }
          else {
            
            // it is still the robot's go
            // Has the counter been dropped into the robot's chosen column?
            if (c != robotMove) {
              strcpy(screen[0], "Cheat!");
              updateScreen();
            }
            else {
              // Clear the led indicating the robot's move
              strip.setPixelColor(robotMove, 0);
              strip.show();
              
              // Humans's turn
              whosTurn = human;
              strcpy(screen[0], "Your go:");
              updateScreen();
            }
          }
        }
      }
    btns = NO_BUTTONS;
    }
  }
}

void displayBoard(char board[]) {
  byte i = 0;
  for(byte row = 0; row <= 5; row++) {
    for(byte col = 0; col <= 6; col++) {
