"""
endurance_sweep.py — 30 min/rev endurance test, 5 passes, burst-pause + accel/decel.

Half-step mode: 14400 steps/rev, 0.025°/step
Thermal: 6° burst with trapezoidal ramp, release coils 28s, repeat.

Timing:
  60 bursts/rev × (1.84s burst + 28.16s cooldown) = exactly 30 min/rev
  5 passes (CW CCW CW CCW CW) × 30 min = 2.5 hours

Run:  python3 endurance_sweep.py
Stop: Ctrl-C (releases coils cleanly)
"""

import serial, time, sys

PORT          = '/dev/ttyACM0'
BAUD          = 115200
STEPS_PER_REV = 14400
BURST_DEG     = 6.0
BURST_STEPS   = round(BURST_DEG * STEPS_PER_REV / 360)   # 240
BURSTS_PER_REV = 360 // int(BURST_DEG)                    # 60
TARGET_REV_SEC = 30 * 60                                   # 1800s

# Trapezoidal accel/decel: (n_steps, delay_ms) segments
# 40 accel + 160 cruise + 40 decel = 240 steps
ACCEL = [(10, 20), (10, 15), (10, 10), (10, 7)]
CRUISE = (160, 5)
DECEL  = list(reversed(ACCEL))

assert sum(n for n,_ in ACCEL) + CRUISE[0] + sum(n for n,_ in DECEL) == BURST_STEPS

BURST_MS   = (sum(n*d for n,d in ACCEL)
              + CRUISE[0] * CRUISE[1]
              + sum(n*d for n,d in DECEL))                 # 1840 ms
COOLDOWN_SEC = TARGET_REV_SEC / BURSTS_PER_REV - BURST_MS / 1000.0  # 28.16 s

# 5 passes alternating direction: CW CCW CW CCW CW
PASS_DIRS = [('CW', 1), ('CCW', -1), ('CW', 1), ('CCW', -1), ('CW', 1)]
TOTAL_PASSES = len(PASS_DIRS)
ETA_MIN = TOTAL_PASSES * TARGET_REV_SEC / 60.0


# ---------------------------------------------------------------------------

def connect():
    s = serial.Serial(PORT, BAUD, timeout=2)
    s.setDTR(False); s.setRTS(False)
    time.sleep(1)
    s.reset_input_buffer()
    return s

def send_step(s, n, delay_ms):
    s.write(f'STEP {n} {delay_ms}\n'.encode())
    deadline = time.time() + abs(n) * delay_ms / 1000.0 + 5
    while time.time() < deadline:
        line = s.readline().decode(errors='replace').strip()
        if '[STEP]' in line:
            return line
    return None

def read_enc(s):
    s.write(b'ENC\n')
    deadline = time.time() + 3
    while time.time() < deadline:
        line = s.readline().decode(errors='replace').strip()
        if '[ENC]' in line:
            counts = int(line.split('counts=')[1].split()[0])
            deg    = float(line.split('deg=')[1])
            return counts, deg
    return None, None

def release(s):
    s.write(b'RELEASE\n')
    time.sleep(0.1)

def fmt(sec):
    m, s = divmod(int(sec), 60)
    h, m = divmod(m, 60)
    return f'{h:02d}:{m:02d}:{s:02d}'

def do_burst(s, direction):
    """Execute one 6° burst with trapezoidal accel/decel, then release coils."""
    profile = ACCEL + [CRUISE] + DECEL
    for n_steps, delay_ms in profile:
        send_step(s, n_steps * direction, delay_ms)
    release(s)


def main():
    print('=' * 62)
    print('  ENDURANCE SWEEP — 30 MIN/REV  HALF-STEP + BURST-PAUSE')
    print('=' * 62)
    print(f'  Steps/rev    : {STEPS_PER_REV}')
    print(f'  Burst size   : {BURST_DEG}° ({BURST_STEPS} steps, trapezoidal ramp)')
    print(f'  Ramp profile : 10@20ms → 10@15 → 10@10 → 10@7 → 160@5 → reverse')
    print(f'  Burst time   : {BURST_MS}ms')
    print(f'  Cooldown     : {COOLDOWN_SEC:.2f}s coils released')
    print(f'  Bursts/rev   : {BURSTS_PER_REV}  →  {TARGET_REV_SEC/60:.0f} min/rev')
    print(f'  Passes       : {TOTAL_PASSES}  (CW CCW CW CCW CW)')
    print(f'  ETA total    : ~{ETA_MIN:.0f} min  ({ETA_MIN/60:.2f} hr)')
    print('=' * 62)

    s = connect()
    s.write(b'ZERO\n')
    time.sleep(0.3)
    print('  [ZERO] encoder zeroed — starting\n')

    wall_start = time.time()
    pass_log   = []

    try:
        for pass_num, (leg, direction) in enumerate(PASS_DIRS, start=1):
            pass_start = time.time()
            print(f'  Pass {pass_num}/{TOTAL_PASSES} — {leg}  (started {time.strftime("%H:%M:%S")})')

            for burst_num in range(1, BURSTS_PER_REV + 1):
                burst_start = time.time()

                do_burst(s, direction)

                # Cooldown with live countdown display
                pct     = burst_num / BURSTS_PER_REV * 100
                elapsed = time.time() - pass_start
                remain  = (BURSTS_PER_REV - burst_num) * (BURST_MS/1000.0 + COOLDOWN_SEC)
                wall_elapsed = time.time() - wall_start
                wall_remain  = (TOTAL_PASSES * TARGET_REV_SEC) - wall_elapsed

                sys.stdout.write(
                    f'\r    burst {burst_num:2d}/{BURSTS_PER_REV}'
                    f'  {pct:5.1f}%'
                    f'  pass={fmt(elapsed)}'
                    f'  pass_rem~{fmt(remain)}'
                    f'  total_rem~{fmt(wall_remain)}'
                    f'  [cool {COOLDOWN_SEC:.0f}s]   '
                )
                sys.stdout.flush()

                # Wait out the remainder of the 30s window
                elapsed_burst = time.time() - burst_start
                sleep_time = max(0, (BURST_MS/1000.0 + COOLDOWN_SEC) - elapsed_burst)
                time.sleep(sleep_time)

            counts, deg = read_enc(s)
            pass_time = time.time() - pass_start
            pass_log.append({
                'pass': pass_num, 'leg': leg,
                'enc_counts': counts, 'enc_deg': deg, 'time_s': pass_time,
            })
            print(f'\n    Pass {pass_num} done — '
                  f'enc={counts} ({deg}°)  time={fmt(pass_time)}\n')

    except KeyboardInterrupt:
        print('\n\n  [INTERRUPTED — releasing coils]')

    release(s)
    s.close()

    total_elapsed = time.time() - wall_start
    print('=' * 62)
    print('  RUN SUMMARY')
    print('=' * 62)
    for p in pass_log:
        print(f"  Pass {p['pass']:2d}  {p['leg']:3s}  "
              f"enc={str(p['enc_counts']):>7}  ({str(p['enc_deg']):>8}°)  "
              f"t={fmt(p['time_s'])}")
    print('=' * 62)
    print(f'  Total time : {fmt(total_elapsed)}')
    print('  Coils released. Done.')


if __name__ == '__main__':
    main()