
int button = 8;     
int greenled = 2;
int yellowled = 5;
int choice=5, select, yellow;

void setup() {
  pinMode(greenled, OUTPUT);   
  pinMode(yellowled, OUTPUT);    
  pinMode(button, INPUT);
  randomSeed(analogRead(5));
}

void loop(){
  // roll the die
  select = random(1,10);
  // compare to threshold
  if (select<choice){
    digitalWrite(yellowled,HIGH);
    yellow=1;
  }
  else{
    digitalWrite(greenled,HIGH);
    yellow=0;
  }
  // wait and see if we get a reward
  for (int t=0;t<100;t++){
    if (digitalRead(button) == HIGH) { 
      // rewarded, encourage current behavior
      if (yellow==1 and choice<=9){
        choice++;
      }
      if (yellow==0 and choice>0){
        choice--;
      }
      // wait for the user to release the button
      while (digitalRead(button) == HIGH) {
        delay(20);  
      }
    }
    // give the user time to react
    delay(20);
  }
  // switch off both LEDs
  digitalWrite(yellowled,LOW);
  digitalWrite(greenled,LOW);
  delay(200);
}
