// ---------------------------------------------------------------------------
//  Fencing ohm tester - Raspberry Pi Pico / Pico 2 (Arduino-Pico core)
//
//  No modes: plug something in and it works out what it is (a cord in both
//  sockets, a weapon on a cord, or the probe on a lame) and answers with a
//  red / yellow / green light and a sound, so nobody has to read numbers
//  with their hands full. The screen adds the detail: a big word and bars.
//
//  Hardware plumbing only: the four CD74HC4051 muxes, the ADS1115, the OLED,
//  the LEDs, the piezo, the button and the fast break sampler. Every number,
//  verdict, light and sound comes from ohm_logic.cpp, which is unit tested on
//  a PC (test_ohm_logic.cpp). Pin map and circuit: ../DESIGN.md.
//
//  Board: "Raspberry Pi Pico" or "Raspberry Pi Pico 2". Library: U8g2.
//  Prebuilt files in ../firmware/: hold BOOTSEL, plug in, drag the .uf2 on.
// ---------------------------------------------------------------------------
#include <Wire.h>
#include <U8g2lib.h>
#include <pico/mutex.h>
#include <hardware/gpio.h>
#include <initializer_list>
#include "ohm_logic.h"

// ------------------------------------------------------------------ pins
static const uint8_t PIN_SDA = 4, PIN_SCL = 5;
static const uint8_t PIN_X[3] = {10, 11, 12};   // U1 drive + U3 sense+ address
static const uint8_t PIN_Y[3] = {13, 14, 15};   // U2 return + U4 sense- address
static const uint8_t PIN_DRIVE_EN = 16;         // U1 /E: LOW = test current on
static const uint8_t PIN_LED_RED = 17;
static const uint8_t PIN_LED_YEL = 18;
static const uint8_t PIN_LED_GRN = 19;
static const uint8_t PIN_BTN = 20;
static const uint8_t PIN_BUZZER = 21;           // passive piezo, driven with tone()
static const uint8_t PIN_REF_ADC = 26;          // R_REF top node, dropout detector
static const uint8_t PIN_VSYS_ADC = 29;         // VSYS / 3
static const uint8_t PIN_VBUS = 24;             // high when USB is plugged in

static const float R_REF = 10.0f;               // the 0.1 % reference resistor

// Pico ADC counts at the R_REF node: about 62 when a wire conducts (5 mA x
// 10 ohm = 50 mV), 0 when it is broken. Half way is the break threshold.
static const int FLEX_THRESHOLD = 25;
// In a cord session the sampler watches for breaks this long between full
// ohm passes.
static const uint32_t FLEX_WINDOW_MS = 700;

U8G2_SSD1306_128X64_NONAME_F_HW_I2C oled(U8G2_R0, U8X8_PIN_NONE);

// --------------------------------------------------------------- ADS1115
static const uint8_t ADS_ADDR = 0x48;
static const uint8_t MUX_01 = 0, MUX_23 = 3;          // AIN0-AIN1, AIN2-AIN3
static const uint8_t PGA_4096 = 1, PGA_256 = 5;
static uint8_t adsRate = 6;                            // 475 SPS on a real ADS1115
static bool adsPresent = false, adsIs12bit = false;

static void adsWrite(uint8_t reg, uint16_t v) {
    Wire.beginTransmission(ADS_ADDR);
    Wire.write(reg); Wire.write(v >> 8); Wire.write(v & 0xFF);
    Wire.endTransmission();
}

static uint16_t adsRead(uint8_t reg) {
    Wire.beginTransmission(ADS_ADDR);
    Wire.write(reg);
    Wire.endTransmission();
    Wire.requestFrom(ADS_ADDR, (uint8_t)2);
    uint16_t hi = Wire.read(), lo = Wire.read();
    return (hi << 8) | lo;
}

// One single-shot conversion, comparator off. Returns signed counts on the
// 16-bit scale (an ADS1015 left-justifies its 12 bits, so the scale matches).
static int32_t adsConvert(uint8_t mux, uint8_t pga) {
    uint16_t cfg = 0x8000 | (mux << 12) | (pga << 9) | 0x0100 | (adsRate << 5) | 0x0003;
    adsWrite(1, cfg);
    uint32_t t0 = micros();
    delayMicroseconds(300);
    while (!(adsRead(1) & 0x8000) && micros() - t0 < 40000) {}
    return (int16_t)adsRead(0);
}

// A real ADS1115 at rate code 4 takes 7.8 ms; an ADS1015 (sold as a 1115 on
// cheap modules) takes 0.6 ms at the same code. Timing tells them apart.
static void adsDetect() {
    Wire.beginTransmission(ADS_ADDR);
    adsPresent = Wire.endTransmission() == 0;
    if (!adsPresent) return;
    adsRate = 4;
    uint32_t t0 = micros();
    adsConvert(MUX_23, PGA_256);
    adsIs12bit = micros() - t0 < 3000;
    adsRate = adsIs12bit ? 4 : 6;   // ADS1015 code 4 = 1600 SPS; ADS1115 code 6 = 475 SPS
}

// ------------------------------------------------------------------ muxes
static void selectX(uint8_t ch) { for (int i = 0; i < 3; ++i) digitalWrite(PIN_X[i], (ch >> i) & 1); }
static void selectY(uint8_t ch) { for (int i = 0; i < 3; ++i) digitalWrite(PIN_Y[i], (ch >> i) & 1); }
static void drive(bool on) { digitalWrite(PIN_DRIVE_EN, on ? LOW : HIGH); }

// The Meter the logic asks for: Kelvin, ratiometric, auto-ranged.
struct HardwareMeter : Meter {
    float ohms(Term x, Term y) override {
        if (!adsPresent) return R_OPEN;
        selectX(x); selectY(y);
        drive(true);
        delayMicroseconds(200);
        RawReading r{};
        uint8_t pga = PGA_256;
        r.xyFsr = 0.256f;
        r.xy = adsConvert(MUX_01, pga);
        if (r.xy > 32000 || r.xy < -32000) {          // over about 40 ohm: wide range
            pga = PGA_4096; r.xyFsr = 4.096f;
            r.xy = adsConvert(MUX_01, pga);
        }
        r.refFsr = 0.256f;
        r.ref = adsConvert(MUX_23, PGA_256);
        float R = ratioToOhms(r, R_REF);
        if (R <= CLOSED_MAX) {
            // Closed path: X now sits at a defined voltage with the current
            // off, so the zero-current readings are the ADC's own offset.
            drive(false);
            delayMicroseconds(200);
            r.xyOff  = adsConvert(MUX_01, pga);
            r.refOff = adsConvert(MUX_23, PGA_256);
            R = ratioToOhms(r, R_REF);
        }
        drive(false);
        return R;
    }
} meter;

// ------------------------------------------------------------------ state
static AutoTester at;
static CordMemory cordMem;
static bool soundOn = true;

// Flex: core 1 owns the muxes while flexRun is set; core 0 only reads the
// tracker under the mutex.
static FlexTracker flex;
static volatile bool flexRun = false, flexIdle = true;
static Term flexFar[3] = {T_J2A, T_J2B, T_J2C};
static bool flexWasArmed = false;
auto_init_mutex(flexLock);

static void stopFlex() {
    flexRun = false;
    while (!flexIdle) {}
    drive(false);
}

// ------------------------------------------------------------- sound
// Non-blocking: a short list of (frequency, ms) steps played from loop().
struct Note { uint16_t hz, ms; };
static Note tune[4];
static uint8_t tuneLen = 0, tuneAt = 0;
static uint32_t noteEnd = 0;

static void play(Sound s) {
    if (!soundOn || s == SND_NONE) return;
    static const Note TICK[]  = {{3200, 18}};
    static const Note BREAK[] = {{2000, 30}};
    static const Note PASS[]  = {{2700, 70}, {0, 40}, {3600, 110}};
    static const Note FAIL[]  = {{900, 160}, {0, 70}, {900, 260}};
    const Note* src; uint8_t n;
    switch (s) {
        case SND_TICK:  src = TICK;  n = 1; break;
        case SND_BREAK: src = BREAK; n = 1; break;
        case SND_PASS:  src = PASS;  n = 3; break;
        default:        src = FAIL;  n = 3; break;
    }
    for (uint8_t i = 0; i < n; ++i) tune[i] = src[i];
    tuneLen = n; tuneAt = 0; noteEnd = 0;
}

static void serviceSound() {
    if (tuneAt >= tuneLen) return;
    uint32_t now = millis();
    if (noteEnd && now < noteEnd) return;
    if (noteEnd) ++tuneAt;
    if (tuneAt >= tuneLen) { noTone(PIN_BUZZER); noteEnd = 0; return; }
    const Note& nt = tune[tuneAt];
    if (nt.hz) tone(PIN_BUZZER, nt.hz); else noTone(PIN_BUZZER);
    noteEnd = now + nt.ms;
}

// --------------------------------------------------------------- lights
static void showLights() {
    bool on = !at.blink || (millis() / 250) % 2 == 0;   // 2 Hz blink = do the next step
    digitalWrite(PIN_LED_RED, at.light == L_RED && on);
    digitalWrite(PIN_LED_YEL, at.light == L_YELLOW && on);
    digitalWrite(PIN_LED_GRN, at.light == L_GREEN && on);
}

// --------------------------------------------------------------- display
static int batteryPct = 100;

static void readBattery() {   // core 0 only, never while core 1 owns the ADC
    float vsys = analogRead(PIN_VSYS_ADC) * 3.3f * 3.0f / 4095.0f + 0.3f;   // + diode drop
    int pct = (int)((vsys - 3.3f) / (4.15f - 3.3f) * 100.0f);
    batteryPct = pct < 0 ? 0 : pct > 100 ? 100 : pct;
}

static void textAt(int x, int y, const char* s) { oled.drawStr(x, y, s); }

static void header(const char* title) {
    oled.setFont(u8g2_font_7x13B_tr);
    textAt(0, 11, title);
    oled.setFont(u8g2_font_5x7_tr);
    char b[12];
    if (digitalRead(PIN_VBUS)) snprintf(b, sizeof b, "%sUSB", soundOn ? "" : "mute ");
    else                       snprintf(b, sizeof b, "%s%d%%", soundOn ? "" : "mute ", batteryPct);
    textAt(128 - oled.getStrWidth(b), 9, b);
    oled.drawHLine(0, 13, 128);
}

static void bigWord(const char* w) {
    oled.setFont(u8g2_font_logisoso16_tr);
    textAt(0, 34, w);
}

// A horizontal bar: value against a full scale, with tick marks at the
// limits, so "how close to failing" reads at a glance.
static void bar(int x, int y, int w, float ohms, float scale, float tick1, float tick2 = 0) {
    oled.drawFrame(x, y, w, 7);
    if (ohms < LEAK_MAX) {
        int fill = ohms >= scale ? w - 2 : (int)(ohms / scale * (w - 2));
        if (fill > 0) oled.drawBox(x + 1, y + 1, fill, 5);
    }
    for (float t : {tick1, tick2}) {
        if (t <= 0) continue;
        int tx = x + 1 + (int)(t / scale * (w - 2));
        oled.drawVLine(tx, y - 1, 9);
    }
}

static void ohmsText(float r, char* o, int n) { formatOhms(r, o, n); }

static void drawReady() {
    header("READY");
    oled.setFont(u8g2_font_6x10_tr);
    textAt(0, 26, "Cord: both sockets");
    textAt(0, 38, "Weapon: on a cord in");
    textAt(0, 48, "  J1, press the tip");
    textAt(0, 60, "Lame: touch the probe");
}

static void drawCord() {
    header(at.cord.type == CORD_CLIP ? "CORD (clip)" : "BODY CORD");
    // The big word says the single most important thing.
    char w[24] = "";
    static const char* P[3] = {"A-B", "B-C", "A-C"};
    int shortAt = -1, leakAt = -1;
    for (int i = 0; i < 3; ++i) {
        if (at.shortVerdict(i) == V_SHORT && shortAt < 0) shortAt = i;
        if (at.shortVerdict(i) == V_LEAKY && leakAt < 0) leakAt = i;
    }
    if (at.breaks)              snprintf(w, sizeof w, "BREAKS %lu", (unsigned long)at.breaks);
    else if (shortAt >= 0)      snprintf(w, sizeof w, "SHORT %s", P[shortAt]);
    else if (at.light == L_RED) snprintf(w, sizeof w, "FAIL");
    else if (leakAt >= 0)       snprintf(w, sizeof w, "LEAKY %s", P[leakAt]);
    else if (at.light == L_GREEN) snprintf(w, sizeof w, "PASS");
    else                        snprintf(w, sizeof w, "...");
    bigWord(w);
    static const char L[3] = {'A', 'B', 'C'};
    oled.setFont(u8g2_font_5x7_tr);
    for (int i = 0; i < 3; ++i) {
        int y = 39 + i * 9;
        char lab[2] = {L[i], 0};
        textAt(0, y + 6, lab);
        const WireResult& wr = at.cord.wire[i];
        float shown = wr.ohms > at.worstWire[i] ? wr.ohms : at.worstWire[i];
        char o[12];
        if (wr.crossed) {
            static const char* T[] = {"", "", "", "to J2-A", "to J2-B", "to J2-C", "to CLIP"};
            snprintf(o, sizeof o, "%s", T[wr.farEnd]);
            bar(8, y, 76, R_OPEN, 2.0f, LIM_CORD);
        } else {
            bar(8, y, 76, shown, 2.0f, LIM_CORD);
            ohmsText(shown, o, sizeof o);
        }
        textAt(88, y + 6, o);
    }
}

static void drawFoilSabre() {
    const FoilSabreTest& f = at.fs;
    header(f.seenOpen ? "FOIL" : "FOIL/SABRE");
    const char* w = at.light == L_RED ? "FAIL" : f.partial ? "DIRTY TIP" : at.blink ? "PRESS TIP" :
                    at.light == L_GREEN ? "PASS" : "...";
    bigWord(w);
    char o[12], line[32];
    ohmsText(f.worstRest, o, sizeof o);
    bar(0, 39, 84, f.worstRest, 4.0f, LIM_SABRE, LIM_FOIL);
    oled.setFont(u8g2_font_5x7_tr);
    textAt(88, 45, o);
    if (f.seenOpen)        snprintf(line, sizeof line, "Tip opens ok  x%u", (unsigned)f.presses);
    else if (f.seenProbe)  { ohmsText(f.worstProbe, o, sizeof o); snprintf(line, sizeof line, "Sabre blade %s", o); }
    else if (f.worstRest <= LIM_SABRE) snprintf(line, sizeof line, "Sabre ok. Foil: press tip");
    else                   snprintf(line, sizeof line, "Foil ok if tip opens");
    textAt(0, 55, line);
    textAt(0, 63, cordMem.valid ? "cord subtracted" : "includes the cord");
}

static void drawEpee() {
    const EpeeTest& e = at.ep;
    header("EPEE");
    Verdict g = e.guardVerdict();
    const char* w = g == V_SHORT ? "SHORT" : at.light == L_RED ? "FAIL" : at.blink ? "LET GO" :
                    g == V_LEAKY ? "LEAKY" : at.light == L_GREEN ? "PASS" : "...";
    bigWord(w);
    char o[12], line[32];
    ohmsText(e.worstPressed, o, sizeof o);
    bar(0, 39, 84, e.seenPressed ? e.worstPressed : R_OPEN, 4.0f, LIM_EPEE);
    oled.setFont(u8g2_font_5x7_tr);
    textAt(88, 45, e.seenPressed ? o : "-");
    snprintf(line, sizeof line, "Guard %s   presses %u", g == V_PASS ? "ok" : verdictWord(g), (unsigned)e.presses);
    textAt(0, 55, line);
    textAt(0, 63, cordMem.valid ? "cord subtracted" : "includes the cord");
}

static void drawProbe() {
    const LameTest& p = at.probe;
    header("LAME / MASK");
    char o[12], line[32];
    if (p.liveOhms < SHORT_MAX) ohmsText(p.liveOhms, o, sizeof o); else snprintf(o, sizeof o, "--");
    bigWord(o);
    bar(0, 39, 126, p.liveOhms < SHORT_MAX ? p.liveOhms : R_OPEN, 10.0f, LIM_WIRE, LIM_LAME);
    if (p.seenContact) {   // worst-spot marker under the bar
        int wx = 1 + (int)((p.worst > 10.0f ? 10.0f : p.worst) / 10.0f * 124);
        oled.drawTriangle(wx - 3, 50, wx + 3, 50, wx, 47);
    }
    oled.setFont(u8g2_font_5x7_tr);
    if (p.seenContact) { ohmsText(p.worst, o, sizeof o); snprintf(line, sizeof line, "Worst %s  tap=new", o); }
    else snprintf(line, sizeof line, "Slide the probe slowly");
    textAt(0, 57, line);
    textAt(0, 64, "ok: 1R wire, 5R lame/mask");
}

static void drawOdd() {
    header("CHECK");
    bigWord("FAIL");
    oled.setFont(u8g2_font_6x10_tr);
    textAt(0, 48, "A-C closed: a tip");
    textAt(0, 60, "wire on the guard?");
}

static void draw() {
    oled.clearBuffer();
    switch (at.shown) {
        case S_CORD:      drawCord(); break;
        case S_FOILSABRE: drawFoilSabre(); break;
        case S_EPEE:      drawEpee(); break;
        case S_PROBE:     drawProbe(); break;
        case S_ODD:       drawOdd(); break;
        default:          drawReady(); break;
    }
    oled.sendBuffer();
}

// ----------------------------------------------------------------- button
// Tap = start over. Hold one second = sound on / off.
static void pollButton() {
    static bool down = false, longFired = false;
    static uint32_t since = 0;
    bool pressed = digitalRead(PIN_BTN) == LOW;
    uint32_t now = millis();
    if (pressed && !down) { down = true; longFired = false; since = now; }
    if (pressed && down && !longFired && now - since > 1000) {
        longFired = true;
        soundOn = !soundOn;
        if (soundOn) play(SND_TICK);
    }
    if (!pressed && down) {
        down = false;
        if (!longFired && now - since > 30) {
            stopFlex();
            at.reset();
            flexWasArmed = false;
            play(SND_TICK);
        }
    }
}

// Everything that must keep running while a measurement or flex window is
// in progress.
static void housekeeping() {
    pollButton();
    serviceSound();
    showLights();
}

// --------------------------------------------------------------- one pass
static void measurePass() {
    uint32_t now = millis();
    if (at.sit == S_CORD) {
        at.updateCord(testCord(meter), cordMem, now);
    } else {
        at.update(scanAll(meter), cordMem, now);
        if (at.sit == S_CORD) at.updateCord(testCord(meter), cordMem, millis());
    }
    play(at.sound);
}

// In a cord session, between full ohm passes, hand the muxes to core 1 and
// count breaks while the user flexes the cord.
static void flexWindow() {
    if (!flexWasArmed) {
        mutex_enter_blocking(&flexLock); flex.reset(); mutex_exit(&flexLock);
        flexWasArmed = true;
    }
    for (int i = 0; i < 3; ++i) flexFar[i] = at.cord.wire[i].farEnd;
    flexRun = true;
    uint32_t end = millis() + FLEX_WINDOW_MS, lastDraw = 0;
    while (millis() < end && flexRun) {
        mutex_enter_blocking(&flexLock);
        uint32_t total = flex.totalBreaks();
        mutex_exit(&flexLock);
        at.noteBreaks(total);
        play(at.sound);
        housekeeping();
        if (millis() - lastDraw > 150) { draw(); lastDraw = millis(); }
        delay(2);
    }
    stopFlex();
}

// ------------------------------------------------------------------ setup
void setup() {
    for (int i = 0; i < 3; ++i) { pinMode(PIN_X[i], OUTPUT); pinMode(PIN_Y[i], OUTPUT); }
    pinMode(PIN_DRIVE_EN, OUTPUT);
    drive(false);
    for (uint8_t p : {PIN_LED_RED, PIN_LED_YEL, PIN_LED_GRN}) {
        pinMode(p, OUTPUT);
        gpio_set_drive_strength(p, GPIO_DRIVE_STRENGTH_12MA);   // ~8 mA through 150 ohm
    }
    pinMode(PIN_BUZZER, OUTPUT);
    pinMode(PIN_BTN, INPUT_PULLUP);
    pinMode(PIN_VBUS, INPUT);
    analogReadResolution(12);

    Wire.setSDA(PIN_SDA);
    Wire.setSCL(PIN_SCL);
    Wire.begin();
    Wire.setClock(400000);
    oled.setBusClock(400000);
    oled.begin();

    // Lamp test: each light and the piezo once, so a dead LED is obvious.
    for (uint8_t p : {PIN_LED_RED, PIN_LED_YEL, PIN_LED_GRN}) { digitalWrite(p, HIGH); delay(180); digitalWrite(p, LOW); }
    tone(PIN_BUZZER, 3200, 60);

    adsDetect();
    readBattery();
    oled.clearBuffer();
    oled.setFont(u8g2_font_7x13B_tr);
    textAt(0, 12, "OHM TESTER");
    oled.setFont(u8g2_font_6x10_tr);
    if (!adsPresent) {
        textAt(0, 30, "ADC NOT FOUND");
        textAt(0, 42, "Check ADS1115 SDA/SCL");
        textAt(0, 54, "and ADDR to GND.");
        digitalWrite(PIN_LED_RED, HIGH);
    } else {
        textAt(0, 28, adsIs12bit ? "ADC: 12-bit (ADS1015)" : "ADC: ADS1115 16-bit");
        textAt(0, 42, "Tap: start over");
        textAt(0, 54, "Hold: sound on/off");
    }
    oled.sendBuffer();
    delay(adsPresent ? 1800 : 5000);
}

void loop() {
    housekeeping();
    measurePass();
    housekeeping();
    static uint32_t lastBatt = 0;
    if (millis() - lastBatt > 2000) { readBattery(); lastBatt = millis(); }
    draw();
    if (at.flexArmed(millis())) flexWindow();
    else flexWasArmed = flexWasArmed && at.sit == S_CORD;
}

// ----------------------------------------------------- core 1: flex sampler
// Round-robin over the three wires: switch the muxes, let them settle, read
// the R_REF node. About 25 us per wire, so each wire is looked at every
// ~0.1 ms and any break of 0.2 ms or longer is counted.
void setup1() {}

void loop1() {
    if (!flexRun) { flexIdle = true; return; }
    flexIdle = false;
    static const Term J1[3] = {T_J1A, T_J1B, T_J1C};
    drive(true);
    while (flexRun) {
        for (int ln = 0; ln < 3; ++ln) {
            selectX(J1[ln]); selectY(flexFar[ln]);
            delayMicroseconds(15);
            bool on = analogRead(PIN_REF_ADC) > FLEX_THRESHOLD;
            mutex_enter_blocking(&flexLock);
            flex.sample(ln, on, micros());
            mutex_exit(&flexLock);
        }
    }
    drive(false);
    flexIdle = true;
}