/*
  Arduino Mega Peristaltic Pump Control with 1.8" ST7735S TFT (SPI)
  Professional Graphical Interface

  - Controls 3 peristaltic pumps simultaneously
  - 2 extra relays toggled anytime via B/C keys
  - Pump volumes input via 4×4 keypad
  - 1.8" TFT SPI ST7735S display UI
  - A=start, B=toggle R4, C=toggle R5, D=emergency stop
*/

#include <SPI.h>
#include <Adafruit_GFX.h>
#include <Adafruit_ST7735.h>
#include <Keypad.h>

// Pins for ST7735
#define TFT_CS   10
#define TFT_DC   9
#define TFT_RST  8
Adafruit_ST7735 tft = Adafruit_ST7735(TFT_CS, TFT_DC, TFT_RST);

// Keypad setup
const byte ROWS = 4, COLS = 4;
char keys[ROWS][COLS] = {{'1','2','3','A'},{'4','5','6','B'},{'7','8','9','C'},{'*','0','#','D'}};
byte rowPins[ROWS] = {2,3,4,5};
byte colPins[COLS] = {6,7,A0,A1};
Keypad keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// Relay pins
#define RELAY1_PIN 8
#define RELAY2_PIN 9
#define RELAY3_PIN 10
#define RELAY4_PIN 11
#define RELAY5_PIN 12

// States & params
unsigned int volumes[3] = {0,0,0};
unsigned long durations[3] = {0,0,0};
bool r4State = false, r5State = false;
bool readyToStart = false;
int inputPump = 0;

// UI constants
#define BTN_W  60
#define BTN_H  30
struct Btn { int x,y; const char *label; };
Btn btnStart = {10, 180, "Start"};
Btn btnR4    = {80, 180, "R4"};
Btn btnR5    = {150,180, "R5"};
Btn btnStop  = {220,180, "Stop"};

// Function prototypes
void drawInputScreen();
void drawMainMenu();
void updateButtons();
void drawStatusScreen(int pumpIdx, int remMl);
void runSimultaneous();
void stopAll();

void setup() {
  // Initialize SPI TFT
  tft.initR(INITR_BLACKTAB);
  tft.setRotation(1);
  tft.fillScreen(ST77XX_BLACK);

  // Initialize keypad & relays
  pinMode(RELAY1_PIN, OUTPUT);
  pinMode(RELAY2_PIN, OUTPUT);
  pinMode(RELAY3_PIN, OUTPUT);
  pinMode(RELAY4_PIN, OUTPUT);
  pinMode(RELAY5_PIN, OUTPUT);
  stopAll();

  // Initial UI: volume input
  drawInputScreen();
}

void loop() {
  char key = keypad.getKey();
  if (!key) return;

  // Toggle relays anytime B/C
  if (key=='B') {
    r4State = !r4State;
    digitalWrite(RELAY4_PIN, r4State?HIGH:LOW);
    updateButtons();
    return;
  }
  if (key=='C') {
    r5State = !r5State;
    digitalWrite(RELAY5_PIN, r5State?HIGH:LOW);
    updateButtons();
    return;
  }
  if (key=='D') {
    stopAll();
    tft.fillScreen(ST77XX_RED);
    tft.setCursor(60,90);
    tft.setTextColor(ST77XX_WHITE);
    tft.setTextSize(3);
    tft.print("STOPPED");
    delay(1000);
    readyToStart = false;
    drawInputScreen();
    return;
  }

  if (!readyToStart) {
    // Numeric input / confirm
    if (isdigit(key)) {
      volumes[inputPump] = volumes[inputPump]*10 + (key - '0');
      // Update display
      tft.fillRect(50,80,200,30,ST77XX_BLACK);
      tft.setCursor(60,100);
      tft.setTextColor(ST77XX_GREEN);
      tft.setTextSize(3);
      tft.print(volumes[inputPump]);
      tft.print(" ml");
    }
    else if (key=='#') {
      durations[inputPump] = volumes[inputPump]*600UL;
      inputPump++;
      if (inputPump<3) drawInputScreen();
      else {
        readyToStart = true;
        drawMainMenu();
      }
    }
  } else if (key=='A') {
    runSimultaneous();
    readyToStart = false;
    drawInputScreen();
  }
}

void drawInputScreen() {
  tft.fillScreen(ST77XX_BLACK);
  tft.setTextSize(2);
  tft.setTextColor(ST77XX_WHITE);
  tft.setCursor(30,30);
  tft.print("Enter Pump ");
  tft.print(inputPump+1);
  tft.print(" Volume");
  tft.drawRect(40,70,240,50,ST77XX_WHITE);
  tft.setCursor(60,100);
  tft.setTextColor(ST77XX_YELLOW);
  tft.print("0 ml");
}

void drawMainMenu() {
  tft.fillScreen(ST77XX_BLUE);
  tft.setTextSize(2);
  tft.setTextColor(ST77XX_WHITE);
  tft.setCursor(90,10);
  tft.print("Control Panel");
  // Draw buttons
  Btn b[4] = {btnStart, btnR4, btnR5, btnStop};
  for (int i=0;i<4;i++) {
    tft.fillRect(b[i].x,b[i].y,BTN_W,BTN_H,ST77XX_WHITE);
    tft.setTextColor(ST77XX_BLACK);
    tft.setCursor(b[i].x+5,b[i].y+8);
    tft.print(b[i].label);
  }
  updateButtons();
}

void updateButtons() {
  // Highlight R4/R5 state
  // R4
  tft.fillRect(btnR4.x,btnR4.y,BTN_W,BTN_H, r4State?ST77XX_GREEN:ST77XX_WHITE);
  tft.setTextColor(r4State?ST77XX_BLACK:ST77XX_BLACK);
  tft.setCursor(btnR4.x+15,btnR4.y+8);
  tft.print("R4");
  // R5
  tft.fillRect(btnR5.x,btnR5.y,BTN_W,BTN_H, r5State?ST77XX_GREEN:ST77XX_WHITE);
  tft.setTextColor(r5State?ST77XX_BLACK:ST77XX_BLACK);
  tft.setCursor(btnR5.x+15,btnR5.y+8);
  tft.print("R5");
}

void runSimultaneous() {
  // Show status UI
  tft.fillScreen(ST77XX_BLACK);
  tft.setTextSize(2);
  for (int i=0;i<3;i++) {
    tft.setCursor(20,20 + i*40);
    tft.setTextColor(ST77XX_CYAN);
    tft.print("P"); tft.print(i+1);
    tft.print(" Rem: ");
  }
  unsigned long startT = millis();
  unsigned long maxDur = max(durations[0], max(durations[1], durations[2]));
  while (millis()-startT < maxDur) {
    unsigned long elapsed = millis()-startT;
    for (int i=0;i<3;i++) {
      int remMl = (durations[i]>elapsed)?(durations[i]-elapsed)/600UL:0;
      tft.fillRect(100,20 + i*40,100,30,ST77XX_BLACK);
      tft.setCursor(100,20 + i*40);
      tft.setTextColor(ST77XX_YELLOW);
      tft.print(remMl); tft.print(" ml");
    }
    // allow B/C/D/A keys inside
    char k = keypad.getKey();
    if (k=='B'||k=='C'||k=='D') {
      // fallback to main loop for handling
      break;
    }
  }
  // Stop pumps
  digitalWrite(RELAY1_PIN,LOW);
  digitalWrite(RELAY2_PIN,LOW);
  digitalWrite(RELAY3_PIN,LOW);
  // Show done
  tft.fillScreen(ST77XX_GREEN);
  tft.setCursor(70,90);
  tft.setTextColor(ST77XX_BLACK);
  tft.setTextSize(3);
  tft.print("DONE");
  delay(1000);
}

void stopAll() {
  digitalWrite(RELAY1_PIN,LOW);
  digitalWrite(RELAY2_PIN,LOW);
  digitalWrite(RELAY3_PIN,LOW);
  digitalWrite(RELAY4_PIN,LOW);
  digitalWrite(RELAY5_PIN,LOW);
  r4State = r5State = false;
}