#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// -------- OLED --------
#define SCREEN_WIDTH   128
#define SCREEN_HEIGHT  64
#define SCREEN_ADDRESS 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1, 1000000UL, 1000000UL);

// -------- RGB LED (active LOW) --------
#define PIN_RED    17
#define PIN_GREEN  16
#define PIN_BLUE   25
inline void ledsOff(){ digitalWrite(PIN_RED,HIGH); digitalWrite(PIN_GREEN,HIGH); digitalWrite(PIN_BLUE,HIGH); }
inline void ledRed(bool on){ digitalWrite(PIN_RED, on?LOW:HIGH); }
inline void ledGreen(bool on){ digitalWrite(PIN_GREEN, on?LOW:HIGH); }
inline void ledBlue(bool on){ digitalWrite(PIN_BLUE, on?LOW:HIGH); }

// -------- Fast IO fallbacks --------
#ifndef digitalWriteFast
  #define digitalWriteFast(pin,val) digitalWrite((pin),(val))
#endif
#ifndef digitalReadFast
  #define digitalReadFast(pin)      digitalRead((pin))
#endif

// -------- Touch config (your routine) --------
#define N_TOUCH   6
#define THRESHOLD 5
int  touch_pins[N_TOUCH] = {3, 4, 2, 27, 1, 26};
int  touch_values[N_TOUCH] = {0,0,0,0,0,0};
bool pin_touched_now[N_TOUCH]  = {false,false,false,false,false,false};
bool pin_touched_past[N_TOUCH] = {false,false,false,false,false,false};

// Map pad *indices* to roles (adjust to your layout)
#define START_IDX 4   // start pad index
#define STOP_IDX  3   // stop pad index
#define MODEL_IDX 0   // cycle Pooled -> Male -> Female
#define DEBUG_IDX 1   // dump CSV when on RESULTS (and global hook if you want)

// -------- Models (avoid struct in signatures) --------
enum SexModel : uint8_t { SX_POOLED=0, SX_MALE=1, SX_FEMALE=2 };
SexModel currentModel = SX_POOLED;

// coefficients (age = a + b*x + c*x^2), x=RT(ms)
const float A_POOLED = -35.392393f, B_POOLED = 0.1930446f,   C_POOLED = -0.00007400686f;
const float A_MALE   = -58.123019f, B_MALE   = 0.2640664f,   C_MALE   = -0.000116241846f;
const float A_FEMALE = -31.603679f, B_FEMALE = 0.176111015f, C_FEMALE = -0.000063725667f;

float predictAgeFromRT(unsigned long rt_ms){
  float x = (float)rt_ms;
  float a,b,c;
  switch (currentModel){
    case SX_MALE:   a=A_MALE;   b=B_MALE;   c=C_MALE;   break;
    case SX_FEMALE: a=A_FEMALE; b=B_FEMALE; c=C_FEMALE; break;
    default:        a=A_POOLED; b=B_POOLED; c=C_POOLED; break;
  }
  float age = a + b*x + c*x*x;
  if (age < 5) age = 5; if (age > 110) age = 110;
  return age;
}
const char* modelName(){
  switch (currentModel){ case SX_MALE: return "Male"; case SX_FEMALE: return "Female"; default: return "Pooled"; }
}

// -------- History (RAM) --------
struct Trial { unsigned long rt_ms; float age_pred; uint8_t model_id; };
const int MAX_TRIALS = 200;
Trial history_[MAX_TRIALS];
int   trialCount = 0;

void saveTrial(unsigned long rt, float age){
  Trial t{rt, age, (uint8_t)currentModel};
  history_[trialCount % MAX_TRIALS] = t;
  trialCount++;
}
void dumpHistoryCSV(){
  Serial.println("\ntrial,rt_ms,age_pred,model");
  int n = (trialCount < MAX_TRIALS) ? trialCount : MAX_TRIALS;
  int start = trialCount - n;
  for (int i=0;i<n;i++){
    const Trial& t = history_[(start+i)%MAX_TRIALS];
    const char* m = (t.model_id==SX_MALE? "male" : (t.model_id==SX_FEMALE? "female":"pooled"));
    Serial.print(i+1); Serial.print(','); Serial.print(t.rt_ms);
    Serial.print(','); Serial.print(t.age_pred,3);
    Serial.print(','); Serial.println(m);
  }
}
void showStatsSmall(){
  int n=(trialCount<MAX_TRIALS)?trialCount:MAX_TRIALS; if(n==0) return;
  int start=trialCount-n;
  float sRT=0,sAge=0; for(int i=0;i<n;i++){ const Trial&t=history_[(start+i)%MAX_TRIALS]; sRT+=t.rt_ms; sAge+=t.age_pred; }
  float avgRT=sRT/n, avgAge=sAge/n;
  display.setCursor(0,52); display.print("Trials "); display.print(n);
  display.setCursor(64,52); display.print("AvgRT "); display.print(avgRT,0); display.print("ms");
  display.setCursor(0,60); display.print("AvgAge "); display.print(avgAge,1);
}

// -------- Touch routines (your code) --------
void update_touch(){
  int t, t_max=200, p;
  for (int i=0;i<N_TOUCH;i++){
    p = touch_pins[i];
    pinMode(p, OUTPUT);
    digitalWriteFast(p, LOW);
    delayMicroseconds(25);
    noInterrupts();
    pinMode(p, INPUT_PULLUP);
    t = 0;
    while (!digitalReadFast(p) && t < t_max) { t++; }
    touch_values[i] = t;
    interrupts();
    pin_touched_past[i] = pin_touched_now[i];
    pin_touched_now[i]  = (touch_values[i] > THRESHOLD);
  }
}
inline bool justPressed(int idx){ return  pin_touched_now[idx] && !pin_touched_past[idx]; }
inline bool isHeld(int idx)     { return  pin_touched_now[idx]; }
void print_touch(){
  for(int i=0;i<N_TOUCH;i++){ Serial.print(touch_values[i]); Serial.print('\t'); }
  Serial.println();
}

// -------- UI helpers --------
void bannerModel(){
  display.fillRect(0,0,SCREEN_WIDTH,10,SSD1306_BLACK);
  display.setCursor(0,0); display.print("Model: "); display.print(modelName());
  display.display();
}
void showIdle(){
  display.clearDisplay(); display.setTextSize(1); display.setTextColor(SSD1306_WHITE);
  bannerModel();
  display.setCursor(0,14); display.print("Touch START to begin");
  display.setCursor(0,26); display.print("(MODEL pad cycles)");
  display.setCursor(0,38); display.print("(DEBUG pad dumps CSV)");
  display.display();
}
void showArmed(){
  display.clearDisplay(); bannerModel();
  display.setCursor(0,14); display.print("Armed... wait");
  display.setCursor(0,26); display.print("False start -> idle");
  display.display();
}
void showGo(){
  display.clearDisplay(); bannerModel();
  display.setCursor(0,14); display.print("GO! Touch STOP now!");
  display.display();
}
void showResult(unsigned long rt, float age){
  display.clearDisplay(); bannerModel();
  display.setCursor(0,14); display.print("Reaction: "); display.print(rt); display.print(" ms");
  display.setCursor(0,26); display.print("Pred Age: "); display.print(age,2); display.print(" y");
  showStatsSmall();
  display.display();
}

// -------- State machine --------
enum State { IDLE, WAITING, GO, RESULTS };
State state = IDLE;
unsigned long tArmMs=0, tGoMs=0, randomWaitMs=0;

void setup(){
  pinMode(PIN_RED,OUTPUT); pinMode(PIN_GREEN,OUTPUT); pinMode(PIN_BLUE,OUTPUT); ledsOff();
  display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS);
  display.clearDisplay(); display.setTextSize(1); display.setTextColor(SSD1306_WHITE);
  Serial.begin(115200);
  delay(150);
  randomSeed(analogRead(A0) + micros());
  showIdle();
}

void loop(){
  update_touch();
  print_touch(); // comment out when tuned

  // global: cycle model
  if (justPressed(MODEL_IDX)){
    currentModel = (SexModel)((currentModel + 1) % 3);
    ledsOff(); ledBlue(true);
    bannerModel();
    delay(200);
    ledsOff();
  }
  // global: CSV dump (anytime you like, esp. on RESULTS)
  if (justPressed(DEBUG_IDX)){
    dumpHistoryCSV();
  }

  switch (state){
    case IDLE:
      if (justPressed(START_IDX)){
        ledRed(true);
        randomWaitMs = random(1000,3000);
        tArmMs = millis();
        showArmed();
        state = WAITING;
      }
      break;

    case WAITING:
      if (isHeld(STOP_IDX)){ // false start
        ledsOff(); showIdle(); state = IDLE;
        break;
      }
      if (millis() - tArmMs >= randomWaitMs){
        ledRed(false); ledGreen(true);
        tGoMs = millis();
        showGo();
        state = GO;
      }
      break;

    case GO:
      if (justPressed(STOP_IDX)){
        ledGreen(false);
        unsigned long rt = millis() - tGoMs;
        float age = predictAgeFromRT(rt);
        saveTrial(rt, age);
        showResult(rt, age);
        ledBlue(true);
        state = RESULTS;
      }
      break;

    case RESULTS:
      if (justPressed(START_IDX)){
        ledsOff(); ledRed(true);
        randomWaitMs = random(1000,3000);
        tArmMs = millis();
        showArmed();
        state = WAITING;
      }
      break;
  }

  delay(40);
}