// ============================================================
//  HV Transformer Variac Stepper Controller — VERSION 24.5
// ============================================================

#include <Arduino.h>
#include <EEPROM.h>
#include <math.h>

// ── 1. PIN DEFINITIONS ──────────────────────────────────────
#define STEP_PIN           9
#define DIR_PIN            10
#define LIMIT_SWITCH_PIN   2

// ── 2. DIRECTION LOGIC ─────────────────────────────────────
#define DIR_UP   LOW
#define DIR_DOWN HIGH

// ── 3. USER CONFIG ─────────────────────────────────────────
float stepsPerKV        = 167.0f;
float defaultSpeedVPS   = 1.0f;
static const float SW0_VOLTAGE     = 1.2f;
static const float HOME_VOLTAGE    = 0.5f;
static const float MAX_SAFE_KV     = 36.0f;
static const float ACCEL_NORMAL_SPS2 = 1100.0f;
static const float SPEED_MIN_SPS   = 50.0f;
static const float MAX_SPEED_KVPS  = 2.0f;
static const int EEPROM_CAL_ADDR   = 8;

// ── 4. FIXED / DERIVED CONSTANTS ───────────────────────────
static const uint32_t TIMER_HZ      = 2000000UL;
static const uint16_t STEP_PULSE_TICKS = 16;
static const long MIN_BRAKE_STEPS   = 5;

// ── 5. STATE VARIABLES ─────────────────────────────────────
enum MotionState { MS_IDLE, MS_MOVING, MS_BRAKE_FOR_REVERSAL };
volatile MotionState motionState = MS_IDLE;

volatile long finalTarget = 0;
volatile long curPos      = 0;
volatile long targetPos   = 0;
volatile bool moving      = false;
volatile bool dirUp       = true;
volatile float curSpeed   = SPEED_MIN_SPS;
volatile float maxSpeed   = 2.0f;
volatile float accel      = ACCEL_NORMAL_SPS2;
volatile uint16_t nextOCR = 1000;

float lastPrintedKV    = -999.0f;
bool homingActive      = true;
bool profileAborted    = false;
volatile bool sw0Detected = false;

// ── 6. FORWARD DECLARATIONS ────────────────────────────────
long kvToSteps(float kv);
float stepsToKV(long steps);
long atomicReadLong(volatile long &v);
void atomicWriteLong(volatile long &v, long val);
void haltSafe();
void startMove(long target, float vps);
void findPhysicalHome();
bool moveToKV(float kv);
void handleMotionState();
void waitForMotionComplete();
bool limitIsActive();
void reportProgress();
void printStatus();
void printHelp();
void cmdSync();
bool runProfile(const char *spec);
void dispatchCommand(char *line);

// ── 7. INTERRUPT SERVICE ROUTINES ─────────────────────────
ISR(TIMER1_COMPA_vect) {
    if (!moving) return;
    if (dirUp && curPos >= kvToSteps(MAX_SAFE_KV)) { moving = false; return; }

    uint16_t thisA = OCR1A;
    PORTB |= _BV(PB1);
    OCR1B = thisA + STEP_PULSE_TICKS;
    TIMSK1 |= (1 << OCIE1B);

    curPos += dirUp ? 1 : -1;
    long dist = dirUp ? (targetPos - curPos) : (curPos - targetPos);

    if (!homingActive && !dirUp && curPos <= kvToSteps(HOME_VOLTAGE)) { moving = false; return; }
    if (dist <= 0) { moving = false; return; }

    float v = curSpeed;
    float dt = (float)nextOCR / (float)TIMER_HZ;
    float brakeDist = (v * v) / (2.0f * accel);
    if ((float)dist <= brakeDist) { v -= accel * dt; }
    else if (v < maxSpeed) { v += accel * dt; if (v > maxSpeed) v = maxSpeed; }
    if (v < SPEED_MIN_SPS) v = SPEED_MIN_SPS;
    curSpeed = v;
    nextOCR = (uint16_t)min(65535.0f, (float)TIMER_HZ / v);
    OCR1A = thisA + nextOCR;
}
ISR(TIMER1_COMPB_vect) { PORTB &= ~_BV(PB1); TIMSK1 &= ~(1 << OCIE1B); }

void haltSafe() {
    noInterrupts();
    TCCR1B = 0;
    TIMSK1 = 0;
    PORTB &= ~_BV(PB1);
    moving = false;
    targetPos = curPos;
    interrupts();
}

void startMove(long tgt, float vps) {
    if (!moving) curSpeed = SPEED_MIN_SPS;
    long pos = atomicReadLong(curPos);
    bool goUp = (tgt > pos);
    digitalWrite(DIR_PIN, goUp ? DIR_UP : DIR_DOWN);
    noInterrupts();
    targetPos = tgt;
    dirUp = goUp;
    float limited = min(vps, MAX_SPEED_KVPS);
    maxSpeed = max(1.0f, limited * stepsPerKV);
    accel = ACCEL_NORMAL_SPS2;
    moving = true;
    TCCR1A = 0;
    TCCR1B = (1 << WGM12) | (1 << CS11);
    TCNT1 = 0;
    OCR1A = 1000;
    nextOCR = 1000;
    TIMSK1 = (1 << OCIE1A);
    interrupts();
}

// ── 8. HOMING ROUTINE ──────────────────────────────────────
void findPhysicalHome() {
    Serial.println(F("\n--- [INIT] SEEKING SW0 ---"));
    homingActive = true;
    sw0Detected = false;

    if (!limitIsActive()) {
        Serial.println(F("Initial State: SW1 (High). Fast Seek to SW0..."));
        atomicWriteLong(curPos, kvToSteps(100.0f));
        startMove(kvToSteps(0.0f), 2.0f);
        delay(150);
        while (digitalRead(LIMIT_SWITCH_PIN) == HIGH) { delay(5); }
        haltSafe();
        delay(500);
        atomicWriteLong(curPos, kvToSteps(SW0_VOLTAGE));
        Serial.println(F("SW0 Found. Handing off to Precision Calibration..."));
    }

    if (limitIsActive()) {
        Serial.println(F("Executing Precision Back-off..."));
        startMove(kvToSteps(2.0f), 0.5f);
        delay(100);
        while (limitIsActive() && moving) { delay(5); }
        haltSafe();
        delay(200);
    }

    atomicWriteLong(curPos, kvToSteps(SW0_VOLTAGE));
    sw0Detected = false;
    Serial.println(F("Reference Set. Moving to 500V floor..."));
    startMove(kvToSteps(HOME_VOLTAGE), 0.5f);
    while (moving) { handleMotionState(); delay(10); }
    haltSafe();
    homingActive = false;
    sw0Detected = limitIsActive();
    reportProgress();
    Serial.println(F("--- SYSTEM READY ---"));
}

// ── 9. MOTION CONTROL ─────────────────────────────────────
bool moveToKV(float kv) {
    if (kv > MAX_SAFE_KV) kv = MAX_SAFE_KV;
    if (kv < HOME_VOLTAGE) kv = HOME_VOLTAGE;

    long desiredSteps = kvToSteps(kv);
    long posSnap = atomicReadLong(curPos);
    bool upNow; float vSnap;
    noInterrupts(); upNow = dirUp; vSnap = curSpeed; interrupts();

    lastPrintedKV = -999.0f;
    bool reversalNeeded = (moving && ((desiredSteps > posSnap) != upNow));
    if (reversalNeeded) {
        float vAbs = fabsf(vSnap);
        long brakeDist = (long)ceilf((vAbs * vAbs) / (2.0f * ACCEL_NORMAL_SPS2));
        if (brakeDist < MIN_BRAKE_STEPS) brakeDist = MIN_BRAKE_STEPS;
        long brakeTarget = upNow ? (posSnap + brakeDist) : (posSnap - brakeDist);
        finalTarget = desiredSteps;
        motionState = MS_BRAKE_FOR_REVERSAL;
        startMove(brakeTarget, defaultSpeedVPS);
        return true;
    }
    finalTarget = desiredSteps;
    motionState = MS_MOVING;
    startMove(desiredSteps, defaultSpeedVPS);
    return true;
}

void handleMotionState() {
    bool stillMoving = moving;
    if (motionState == MS_BRAKE_FOR_REVERSAL && !stillMoving) {
        motionState = MS_MOVING; startMove(finalTarget, defaultSpeedVPS); return;
    }
    if (motionState == MS_MOVING && !stillMoving) {
        motionState = MS_IDLE; reportProgress();
    }

    if (limitIsActive()) {
        if (homingActive) sw0Detected = true;
        else if (!sw0Detected && !dirUp) {
            haltSafe();
            atomicWriteLong(curPos, kvToSteps(SW0_VOLTAGE));
            sw0Detected = true;
            if (motionState == MS_MOVING) startMove(finalTarget, defaultSpeedVPS);
        }
    } else sw0Detected = false;
}

void waitForMotionComplete() {
    while (motionState != MS_IDLE) { handleMotionState(); reportProgress(); delay(10); }
}

inline bool limitIsActive() { return digitalRead(LIMIT_SWITCH_PIN) == LOW; }

void reportProgress() {
    long pos = atomicReadLong(curPos);
    float currentKV = stepsToKV(pos);
    float rounded = roundf(currentKV * 20.0f) / 20.0f;
    bool physicalPinPressed = limitIsActive();
    if (rounded != lastPrintedKV) {
        Serial.println(rounded, 2);
        Serial.println(physicalPinPressed ? F("SW0") : F("SW1"));
        lastPrintedKV = rounded;
    }
}

// ── 10. PROFILE EXECUTION ──────────────────────────────────
bool runProfile(const char *spec) {
    profileAborted = false;
    Serial.println(F("=== PROFILE START ==="));
    const char *p = spec;
    while (*p) {
        while (*p == ' ' || *p == '\t' || *p == ';') p++;
        if (*p == '\0') break;
        float targetKV = atof(p);
        const char *comma = strchr(p, ',');
        if (!comma) break;
        long holdSeconds = atol(comma + 1);

        Serial.print(F("Moving to: ")); Serial.print(targetKV); Serial.println(F("kV"));
        moveToKV(targetKV); waitForMotionComplete();
        if (profileAborted) break;

        Serial.print(F("Holding for ")); Serial.print(holdSeconds); Serial.println(F("s..."));
        unsigned long startHold = millis();
        long lastSecRem = -1;

        while ((millis() - startHold) < (unsigned long)holdSeconds * 1000UL) {
            handleMotionState();
            long secRem = holdSeconds - ((millis() - startHold) / 1000UL);
            if (secRem != lastSecRem) { Serial.print(secRem); Serial.print(F(" ")); lastSecRem = secRem; }
            if (Serial.available()) { if (Serial.peek() == '0') { Serial.read(); profileAborted = true; break; } }
            if (profileAborted) break;
            delay(10);
        }
        Serial.println();
        if (profileAborted) break;
        const char *semi = strchr(p, ';'); if (!semi) break; p = semi + 1;
    }
    if (profileAborted) Serial.println(F("PROFILE ABORTED"));
    else Serial.println(F("=== PROFILE COMPLETE ==="));
    moveToKV(HOME_VOLTAGE); waitForMotionComplete();
    return !profileAborted;
}

// ── 11. SYSTEM SHELL ───────────────────────────────────────
void cmdSync() {
    moveToKV(2.0f); waitForMotionComplete();
    Serial.println(F("Enter meter reading:"));
    while (!Serial.available());
    float measured = Serial.parseFloat();
    if (measured > 0.1f) {
        stepsPerKV = (float)atomicReadLong(curPos) / measured;
        EEPROM.put(EEPROM_CAL_ADDR, stepsPerKV);
        Serial.println(F("Sync saved."));
    }
    moveToKV(HOME_VOLTAGE); waitForMotionComplete();
}

void printStatus() {
    long pos = atomicReadLong(curPos);
    Serial.println(F("\n--- STATUS ---------------"));
    Serial.print(F("Voltage : ")); Serial.print(stepsToKV(pos), 2); Serial.println(F(" kV"));
    Serial.print(F("Limit   : ")); Serial.println(limitIsActive() ? F("CLOSED (SW0)") : F("OPEN (SW1)"));
    Serial.print(F("Speed   : ")); Serial.print(curSpeed / stepsPerKV, 2); Serial.println(F(" kV/s"));
    Serial.println(F("--------------------------"));
}

void printHelp() {
    Serial.println(F("\n--- COMMANDS -------------------"));
    Serial.println(F("  <kV>             move (0.5–36.0)"));
    Serial.println(F("  0                abort / home"));
    Serial.println(F("  speed <v>        set speed kV/s"));
    Serial.println(F("  sync             calibrate at 2kV"));
    Serial.println(F("  profile <k,s;…>  test sequence"));
    Serial.println(F("  status           show info"));
    Serial.println(F("--------------------------------"));
}

void dispatchCommand(char *line) {
    while (*line == ' ' || *line == '\t') line++;
    if (*line == '\0') return;
    if (strcmp(line, "0") == 0) { profileAborted = true; moveToKV(HOME_VOLTAGE); }
    else if (strcasecmp(line, "status") == 0) printStatus();
    else if (strcasecmp(line, "help") == 0 || line[0] == '?') printHelp();
    else if (strncasecmp(line, "speed ", 6) == 0) defaultSpeedVPS = atof(line + 6);
    else if (strncasecmp(line, "profile ", 8) == 0) runProfile(line + 8);
    else if (strcasecmp(line, "sync") == 0) cmdSync();
    else moveToKV(atof(line));
}

// ── 12. SETUP & LOOP ───────────────────────────────────────
void setup() {
    Serial.begin(115200);
    pinMode(STEP_PIN, OUTPUT);
    pinMode(DIR_PIN, OUTPUT);
    pinMode(LIMIT_SWITCH_PIN, INPUT_PULLUP);
    delay(50);

    float saved;
    EEPROM.get(EEPROM_CAL_ADDR, saved);
    if (!isnan(saved) && saved > 10.0f) stepsPerKV = saved;

    Serial.println(F("\n--- HV VARIAC CONTROLLER ---"));
    Serial.println(F("Ready to Home? (y/n)"));
    while (!Serial.available());
    if (tolower(Serial.read()) == 'y') findPhysicalHome();
    else {
        homingActive = false;
        atomicWriteLong(curPos, kvToSteps(HOME_VOLTAGE));
    }
}

void loop() {
    static char lineBuf[96]; static size_t lineLen = 0;
    while (Serial.available()) {
        char c = Serial.read(); Serial.write(c);
        if (c == '\n' || c == '\r') {
            lineBuf[lineLen] = '\0';
            if (lineLen > 0) dispatchCommand(lineBuf);
            lineLen = 0;
        } else if (lineLen < 95) lineBuf[lineLen++] = c;
    }
    handleMotionState();
    static unsigned long lastR = 0;
    if (millis() - lastR >= 20) { reportProgress(); lastR = millis(); }
}

// ── 13. UTILITY FUNCTIONS ──────────────────────────────────
long kvToSteps(float kv) { return (long)lroundf(kv * stepsPerKV); }
float stepsToKV(long steps) { return (float)steps / stepsPerKV; }
long atomicReadLong(volatile long &v) { noInterrupts(); long t = v; interrupts(); return t; }
void atomicWriteLong(volatile long &v, long val) { noInterrupts(); v = val; interrupts(); }