"""
PROJECT SLOW-SWEEP — Closed-loop encoder-controlled rotation.

Runs the turntable in continuous mode at min speed, watches the encoder,
and sends STOP at exactly the target angle. Then waits out the remainder
of the step interval. Achieves ultra-slow rotation (20-30 min/rev) with
encoder-verified precision.

Physics:
  Min speed = 77s/rev = 2400 counts/rev -> 32ms/count
  IR stop latency ~200ms -> ~6 count overshoot at min speed
  We target (desired_counts - OVERSHOOT_COUNTS) to land on the mark.

Usage:
  python3 slow_sweep.py [--minutes-per-rev 25] [--direction ccw]
                        [--step-deg 6.0] [--revolutions N]
"""

import serial, time, threading, subprocess, argparse
import numpy as np

WROOM_PORT       = "/dev/ttyACM0"
ATOM_PORT        = "/dev/ttyACM1"
MIC_DEVICE       = "hw:3,0"
MIC_RATE         = 44100
LOG              = "/home/person/Desktop/Turntable_Hacking/slow_sweep.log"

COUNTS_PER_DEG   = 2400 / 360.0        # 6.667 counts/deg
OVERSHOOT_COUNTS = 6                    # ~200ms IR latency at min speed

# ── Logging ───────────────────────────────────────────────────────────────────

def log(msg):
    ts = time.strftime("%H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)
    with open(LOG, "a") as f:
        f.write(line + "\n")

# ── Beep listener ─────────────────────────────────────────────────────────────

class BeepListener:
    CHUNK_MS    = 20
    MIN_BEEP_MS = 40
    MIN_GAP_MS  = 120
    REPORT_MS   = 400

    def __init__(self):
        self._lock   = threading.Lock()
        self._event  = threading.Event()
        self._counts = []
        self.threshold = 50
        self._stop  = threading.Event()

    def calibrate(self):
        proc = subprocess.Popen(
            ['arecord', '-D', MIC_DEVICE, '-f', 'S16_LE',
             '-r', str(MIC_RATE), '-c', '1', '-d', '2'],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
        data = proc.stdout.read()
        proc.kill()
        if data:
            s = np.frombuffer(data, dtype=np.int16).astype(np.float32)
            ambient = float(np.sqrt(np.mean(s**2)))
        else:
            ambient = 0.0
        self.threshold = max(ambient * 1.3, 50)
        log(f"MIC calibrated: ambient={ambient:.1f}  threshold={self.threshold:.1f}")

    def start(self):
        self._stop.clear()
        threading.Thread(target=self._run, daemon=True).start()

    def stop(self):
        self._stop.set()

    def wait_beep(self, timeout=3.0):
        self._event.clear()
        if self._event.wait(timeout):
            with self._lock:
                return self._counts.pop(0) if self._counts else 1
        return 0

    def _run(self):
        chunk_n    = int(MIC_RATE * self.CHUNK_MS / 1000)
        chunk_bytes = chunk_n * 2
        proc = subprocess.Popen(
            ['arecord', '-D', MIC_DEVICE, '-f', 'S16_LE',
             '-r', str(MIC_RATE), '-c', '1'],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)

        beep_count = 0
        in_beep    = False
        beep_ms    = 0
        gap_ms     = 0

        while not self._stop.is_set():
            chunk = proc.stdout.read(chunk_bytes)
            if not chunk:
                break
            rms  = float(np.sqrt(np.mean(
                np.frombuffer(chunk, dtype=np.int16).astype(np.float32)**2)))
            loud = rms > self.threshold

            if loud:
                if not in_beep:
                    in_beep = True; beep_ms = 0; gap_ms = 0
                else:
                    beep_ms += self.CHUNK_MS; gap_ms = 0
            else:
                if in_beep:
                    gap_ms += self.CHUNK_MS
                    if gap_ms >= self.MIN_GAP_MS:
                        if beep_ms >= self.MIN_BEEP_MS:
                            beep_count += 1
                        in_beep = False; beep_ms = 0
                else:
                    if beep_count > 0:
                        gap_ms += self.CHUNK_MS
                        if gap_ms >= self.REPORT_MS:
                            with self._lock:
                                self._counts.append(beep_count)
                            self._event.set()
                            beep_count = 0; gap_ms = 0
        proc.kill()

# ── Serial helpers ────────────────────────────────────────────────────────────

def send_ir(atom, cmd, delay=0.15):
    atom.write(f"{cmd}\n".encode())
    time.sleep(delay)

def read_enc_alive(wroom, timeout=2.0):
    """Read next ALIVE heartbeat and return encoder count."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        line = wroom.readline().decode(errors="replace").strip()
        if "[ALIVE]" in line and "enc=" in line:
            try:
                return int(line.split("enc=")[1].split()[0])
            except:
                pass
    return None

def read_enc_live(wroom):
    """Read next [ENC] line (fires every encoder count change during motion)."""
    while True:
        line = wroom.readline().decode(errors="replace").strip()
        if "[ENC]" in line and "counts=" in line:
            try:
                return int(line.split("counts=")[1].split()[0])
            except:
                pass

def wait_stopped(wroom, settle_s=1.5, timeout=10.0):
    """Wait until encoder is stable for settle_s seconds."""
    last_enc  = None
    last_move = time.time()
    deadline  = time.time() + timeout
    while time.time() < deadline:
        line = wroom.readline().decode(errors="replace").strip()
        if "[ALIVE]" not in line or "enc=" not in line:
            continue
        try:
            enc = int(line.split("enc=")[1].split()[0])
        except:
            continue
        if last_enc is not None and enc != last_enc:
            last_move = time.time()
        last_enc = enc
        if time.time() - last_move >= settle_s:
            break
    return last_enc

# ── Closed-loop step ──────────────────────────────────────────────────────────

def closed_loop_step(wroom, atom, mic, target_deg, direction_sign, interval_s):
    """
    Run continuous rotation, stop when encoder hits target_deg.
    direction_sign: +1 for CCW (encoder increases), -1 for CW (encoder decreases).
    Returns (actual_deg, overshoot_deg).
    """
    target_counts = int(target_deg * COUNTS_PER_DEG)
    stop_at       = target_counts - OVERSHOOT_COUNTS  # fire stop early

    cycle_start = time.time()

    # Ensure table is fully stopped before starting
    wait_stopped(wroom, settle_s=1.0, timeout=5.0)

    # Flush stale serial data, note start position
    wroom.reset_input_buffer()
    enc_start = read_enc_alive(wroom, timeout=2.0)
    if enc_start is None:
        log("WARNING: could not read start encoder position")
        enc_start = 0

    # Start continuous rotation — retry until 1-beep confirms table started
    started = False
    for attempt in range(3):
        send_ir(atom, "continuous", delay=0.1)
        b = mic.wait_beep(timeout=1.5)
        if b == 1:
            started = True
            break
        # No beep — table may already be in continuous mode, try startstop to reset then retry
        send_ir(atom, "startstop", delay=0.1)
        mic.wait_beep(timeout=1.0)
        time.sleep(0.3)
    if not started:
        return 0.0, 0.0, enc_start  # skip this step, report zero move

    # Watch encoder until stop_at counts reached (with timeout guard)
    deadline = time.time() + 60.0
    while time.time() < deadline:
        enc = read_enc_live(wroom)
        delta = (enc - enc_start) * direction_sign
        if delta >= stop_at:
            send_ir(atom, "startstop", delay=0.0)
            break
    else:
        send_ir(atom, "startstop", delay=0.0)  # safety stop on timeout

    # Let it fully stop
    enc_final = wait_stopped(wroom, settle_s=1.5, timeout=10.0) or enc_start
    actual_counts = abs(enc_final - enc_start)
    actual_deg    = actual_counts / COUNTS_PER_DEG
    overshoot_deg = actual_deg - target_deg

    # Wait out the rest of the interval
    elapsed        = time.time() - cycle_start
    wait_remaining = max(0.0, interval_s - elapsed)
    if wait_remaining > 0:
        time.sleep(wait_remaining)

    return actual_deg, overshoot_deg, enc_final

# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="SLOW-SWEEP closed-loop pulse")
    parser.add_argument("--minutes-per-rev", type=float, default=25.0)
    parser.add_argument("--direction", choices=["cw", "ccw"], default="ccw")
    parser.add_argument("--step-deg", type=float, default=6.0,
                        help="Degrees per step (default 6.0)")
    parser.add_argument("--revolutions", type=float, default=0,
                        help="Stop after N revolutions (0 = run forever)")
    args = parser.parse_args()

    steps_per_rev  = 360.0 / args.step_deg
    interval_s     = (args.minutes_per_rev * 60.0) / steps_per_rev
    direction_sign = +1 if args.direction == "ccw" else -1
    target_counts  = int(args.step_deg * COUNTS_PER_DEG)

    log("=" * 55)
    log("SLOW-SWEEP (closed-loop encoder control)")
    log(f"  direction      : {args.direction}")
    log(f"  step size      : {args.step_deg}° ({target_counts} counts)")
    log(f"  overshoot comp : {OVERSHOOT_COUNTS} counts")
    log(f"  steps/rev      : {steps_per_rev:.1f}")
    log(f"  interval       : {interval_s:.1f}s/step")
    log(f"  speed target   : {args.minutes_per_rev:.1f} min/rev")
    if args.revolutions:
        log(f"  target         : {args.revolutions} rev ({int(args.revolutions*steps_per_rev)} steps)")
    log("=" * 55)

    wroom = serial.Serial(WROOM_PORT, 115200, timeout=1)
    atom  = serial.Serial(ATOM_PORT,  115200, timeout=1)
    time.sleep(1)
    wroom.reset_input_buffer()

    mic = BeepListener()
    mic.calibrate()
    mic.start()

    enc_home = read_enc_alive(wroom) or 0
    log(f"Home encoder: {enc_home}")

    # Set min speed (25x slowdown -> 77s/rev = slowest)
    log("Setting min speed (25x slowdown)...")
    for _ in range(25):
        send_ir(atom, "0x43", delay=0.05)
    time.sleep(0.5)

    # Set direction
    dir_cmd = args.direction
    send_ir(atom, dir_cmd)
    b = mic.wait_beep(timeout=2.0)
    log(f"Direction set to {args.direction} (beep={b})")
    time.sleep(0.5)

    max_steps  = int(args.revolutions * steps_per_rev) if args.revolutions else 0
    step_count = 0
    total_deg  = 0.0
    enc_pos    = enc_home

    log("Starting sweep — Ctrl+C to stop")
    log(f"{'step':>5}  {'moved':>7}  {'overshoot':>9}  {'total':>10}  {'pos':>8}")

    try:
        while True:
            actual_deg, overshoot_deg, enc_pos = closed_loop_step(
                wroom, atom, mic, args.step_deg, direction_sign, interval_s)

            step_count += 1
            total_deg  += actual_deg
            total_rev   = total_deg / 360.0
            offset_deg  = (enc_pos - enc_home) * 0.15

            log(f"{step_count:>5}  {actual_deg:>6.1f}°  {overshoot_deg:>+8.1f}°  "
                f"{total_deg:>7.0f}° ({total_rev:.2f}rev)  {offset_deg:>+7.1f}°")

            if actual_deg > args.step_deg * 2:
                log(f"WARNING: step {step_count} overran {actual_deg:.1f}° (target {args.step_deg}°) — IR stop not received?")

            if max_steps and step_count >= max_steps:
                log(f"Target reached: {args.revolutions} rev in {step_count} steps")
                break

    except KeyboardInterrupt:
        log("")
        log(f"Stopped after {step_count} steps")
        log(f"Total: {total_deg:.1f}° = {total_deg/360:.2f} rev")

    finally:
        mic.stop()
        atom.close()
        wroom.close()

if __name__ == "__main__":
    main()