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

#define SCREEN_WIDTH 128     // OLED display width, in pixels
#define SCREEN_HEIGHT 64     // OLED display height, in pixels#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C  // See datasheet for Address
#define OLED_RESET -1        // Arduino reset
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

Adafruit_INA219 ina219;  // 0x40

void setup(void) {
  Serial.begin(115200);
  while (!Serial) {
    delay(1);  // Required for some development boards
  }

  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;)
      ;  // Don't proceed, loop forever
  }
  delay(1000);

  pinMode(A2, INPUT_PULLUP);

  display.clearDisplay();
  display.setTextSize(3);               // Normal 1:1 pixel scale
  display.setTextColor(SSD1306_WHITE);  // Draw white text

  // Initialize the INA219.
  // By default the initialization will use the largest range (32V, 2A).  However
  // you can call a setCalibration function to change this range (see comments).
  if (!ina219.begin()) {
    Serial.println("Failed to find INA219 chip");
    while (1) { delay(10); }
  }
  // To use a slightly lower 32V, 1A range (higher precision on amps):
  //ina219.setCalibration_32V_1A();
  // Or to use a lower 16V, 400mA range (higher precision on volts and amps):
  //ina219.setCalibration_16V_400mA();

  Serial.println("Measuring voltage and current with INA219 ...");
}

void loop(void) {
  float shuntvoltage = 0;
  float busvoltage = 0;
  float current_mA = 0;
  float loadvoltage = 0;
  float power_mW = 0;

  shuntvoltage = ina219.getShuntVoltage_mV();
  busvoltage = ina219.getBusVoltage_V();
  current_mA = ina219.getCurrent_mA();
  power_mW = ina219.getPower_mW();
  loadvoltage = busvoltage + (shuntvoltage / 1000);

  if (current_mA < 0) current_mA = 0;

  if (digitalRead(A2) == HIGH) {

    Serial.println("HI");

    display.clearDisplay();
    display.setCursor(0, 4);  // Row 1 Column 5
    if (current_mA >= 1000) {
      display.print(F("A "));
      display.println(current_mA / 1000, 3);
    } else {
      display.print(F("mA "));
      if (current_mA < 10.0) {
        display.println(current_mA, 2);
      } else if (current_mA < 100.0) {
        display.println(current_mA, 1);
      } else {
        display.println(current_mA, 0);
      }
    }
    display.setCursor(0, 36);
    display.print(F("V "));
    display.println(busvoltage);
    display.display();
  } else {

    Serial.println("LO");

    display.clearDisplay();
    display.setCursor(0, 4);  // Row 1 Column 5
    if (power_mW >= 1000) {
      display.print(F("W "));
      display.println(power_mW / 1000, 3);
    } else {
      display.print(F("mW "));
      if (power_mW < 10.0) {
        display.println(power_mW, 2);
      } else if (power_mW < 100.0) {
        display.println(power_mW, 1);
      } else {
        display.println(power_mW, 0);
      }
    }
    display.setCursor(0, 36);
    display.print(F("L "));
    display.println(loadvoltage);
    display.display();
  }

  Serial.print("Bus Voltage:   ");
  Serial.print(busvoltage);
  Serial.println(" V");
  Serial.print("Shunt Voltage: ");
  Serial.print(shuntvoltage);
  Serial.println(" mV");
  Serial.print("Load Voltage:  ");
  Serial.print(loadvoltage);
  Serial.println(" V");
  Serial.print("Current:       ");
  Serial.print(current_mA);
  Serial.println(" mA");
  Serial.print("Power:         ");
  Serial.print(power_mW);
  Serial.println(" mW");
  Serial.println("");

  delay(500);
}