# The ChronoPulse: laptop side. Draws each screen and streams it to the board.
# Fadi G. Haddad, HTMAA 2026. Pair with chronopulse.ino running on the board.
# Built on Quentin Bolsee's computer_link example (qpad-xiao, MIT License):
# https://forge.cba.mit.edu/quentin.bolsee/qpad-xiao
# Needs: pip3 install numpy pyserial pillow
# Run:   python3 health_monitor.py

import argparse
import glob
import math
import os
import time

import numpy as np
import serial
from PIL import Image, ImageDraw, ImageFont

import libs.hdlc_encoding as hdlc_encoding

W, H = 128, 64

# must match the enums in the Arduino sketch
FINE, CAUTION, DANGER, OFF = 0, 1, 2, 3
WAVES, LOGO = 0, 1
LABELS = {FINE: "FINE", CAUTION: "CAUTION", DANGER: "DANGER"}

# how many milliseconds of heartbeat each screen pixel covers
MS_PER_PX = 20

# trace placement: baseline row and how tall the spikes are, in pixels
BASELINE_Y = 26
AMPLITUDE = 20

SCANLINES = True

TITLE = "THE CHRONOPULSE"        # title page: logo view with the lights off
MENU_TITLE = "INSTRUCTIONS"      # controls card: waves view with the lights off

# logo throb: its width as a fraction of the screen, at rest and at the peak of a beat
LOGO_REST_SCALE = 0.82
LOGO_PEAK_SCALE = 1.0

LOGO_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logo", "logo_full_bw.png")


# ---------------------------------------------------------------------------
# Heart-monitor waves. These MUST match the same functions in the Arduino
# sketch, so the trace and the lights stay in sync. t is in milliseconds.
# ---------------------------------------------------------------------------

def pulse(t, center, width):
    x = (t - center) / width
    return np.exp(-x * x)


def fine_ecg(now):
    t = now % 1000
    return (0.12 * pulse(t, 150, 35)
            - 0.15 * pulse(t, 265, 10)
            + 1.00 * pulse(t, 285, 12)
            - 0.35 * pulse(t, 305, 12)
            + 0.30 * pulse(t, 520, 60))


CAUTION_AT = [150, 480, 780, 1050, 1350, 1900, 2150]
CAUTION_H = [0.40, 0.50, 0.45, 1.00, 0.30, 0.70, 0.35]


def caution_ecg(now):
    t = now % 2400
    v = 0.08 * np.sin(2 * math.pi * t / 400.0)
    for at, h in zip(CAUTION_AT, CAUTION_H):
        v = v + h * pulse(t, at, 14) - 0.3 * h * pulse(t, at + 22, 12)
    return v


def danger_ecg(now):
    return (0.50 * np.sin(2 * math.pi * (now % 310) / 310.0)
            + 0.33 * np.sin(2 * math.pi * (now % 170) / 170.0 + 1.0)
            + 0.17 * np.sin(2 * math.pi * (now % 97) / 97.0 + 2.0))


ECG = {FINE: fine_ecg, CAUTION: caution_ecg, DANGER: danger_ecg}


# ---------------------------------------------------------------------------
# Drawing
# ---------------------------------------------------------------------------

# same as the lights: follow the upward parts of the trace; DANGER never settles
def beat_wave(health, now):
    if health == OFF:
        return 0.0
    v = float(ECG[health](np.float64(now)))
    if health == DANGER:
        return 0.5 + 0.5 * v
    return min(max(v, 0.0), 1.0)


# the logo, scaled to `scale` x screen width and centered
def draw_logo(source, scale):
    new_w = max(1, round(W * scale))
    new_h = max(1, round(source.height * new_w / source.width))
    logo = source.resize((new_w, new_h), Image.LANCZOS)
    logo = logo.point(lambda p: 255 if p > 90 else 0).convert("1")
    canvas = Image.new("1", (W, H), 0)
    canvas.paste(logo, ((W - new_w) // 2, (H - new_h) // 2))
    return canvas


def text_width(draw, text, font):
    left, _, right, _ = draw.textbbox((0, 0), text, font=font)
    return right - left


# a lit strip across the top with the text in dark pixels, like a menu banner
def draw_banner(draw, text, font):
    draw.rectangle((0, 0, W - 1, 11), fill=1)
    draw.text(((W - text_width(draw, text, font)) // 2, 0), text, font=font, fill=0)


# title page: the name in plain white text, the logo filling the space below it
def draw_title_page(source, font):
    img = Image.new("1", (W, H), 0)
    draw = ImageDraw.Draw(img)
    draw.text(((W - text_width(draw, TITLE, font)) // 2, 0), TITLE, font=font, fill=1)

    new_h = round(source.height * W / source.width)
    logo = source.resize((W, new_h), Image.LANCZOS)
    logo = logo.point(lambda p: 255 if p > 90 else 0).convert("1")
    top = 12 + (H - 12 - new_h) // 2
    img.paste(logo, (0, top))
    return img


# a small solid arrow for a direction pad, or the pad's number
def draw_symbol(draw, symbol, x, y, font):
    shapes = {
        "right": [(x, y + 2), (x, y + 10), (x + 5, y + 6)],
        "left":  [(x + 5, y + 2), (x + 5, y + 10), (x, y + 6)],
        "down":  [(x, y + 4), (x + 7, y + 4), (x + 3, y + 9)],
        "up":    [(x, y + 9), (x + 7, y + 9), (x + 3, y + 4)],
    }
    if symbol in shapes:
        draw.polygon(shapes[symbol], fill=1)
    else:
        draw.text((x, y), symbol, font=font, fill=1)


# title card shown when the lights are off: a menu banner plus the controls
def draw_menu(font):
    img = Image.new("1", (W, H), 0)
    draw = ImageDraw.Draw(img)

    draw_banner(draw, MENU_TITLE, font)

    # view pads first
    for x, (number, label) in zip((4, 68), (("1", "WAVES"), ("2", "LOGO"))):
        draw.text((x, 14), number, font=font, fill=1)
        draw.text((x + 10, 14), label, font=font, fill=1)

    draw_dpad(draw, font, cx=68, cy=45)
    return img


# a plus-shaped D-pad outline with an arrow in each arm, labeled like a controller diagram
def draw_dpad(draw, font, cx, cy, arm=5, reach=15):
    a, L = arm, reach
    outline = [(cx - a, cy - L), (cx + a, cy - L), (cx + a, cy - a), (cx + L, cy - a),
               (cx + L, cy + a), (cx + a, cy + a), (cx + a, cy + L), (cx - a, cy + L),
               (cx - a, cy + a), (cx - L, cy + a), (cx - L, cy - a), (cx - a, cy - a)]
    draw.polygon(outline, outline=1, fill=0)

    draw.polygon([(cx - 3, cy - 7), (cx + 3, cy - 7), (cx, cy - 12)], fill=1)   # up
    draw.polygon([(cx - 3, cy + 7), (cx + 3, cy + 7), (cx, cy + 12)], fill=1)   # down
    draw.polygon([(cx - 7, cy - 3), (cx - 7, cy + 3), (cx - 12, cy)], fill=1)   # left
    draw.polygon([(cx + 7, cy - 3), (cx + 7, cy + 3), (cx + 12, cy)], fill=1)   # right

    draw.text((cx + a + 4, cy - L - 2), "OFF", font=font, fill=1)
    draw.text((cx + a + 4, cy + L - 10), "DANGER", font=font, fill=1)
    draw.text((cx + L + 4, cy - 6), "FINE", font=font, fill=1)
    draw.text((cx - L - 4 - text_width(draw, "CAUTION", font), cy - 6), "CAUTION", font=font, fill=1)


def draw_waves(health, now, font):
    img = Image.new("1", (W, H), 0)
    draw = ImageDraw.Draw(img)

    if SCANLINES:
        for y in range(2, 50, 6):
            for x in range(0, W, 3):
                img.putpixel((x, y), 1)

    if health == OFF:
        return img

    # the right edge of the screen is "now"; older moments scroll off to the left
    xs = np.arange(W)
    times = now - (W - 1 - xs) * MS_PER_PX
    times = times.astype(np.float64)
    ys = BASELINE_Y - AMPLITUDE * ECG[health](times)
    ys = np.clip(ys, 0, 49)

    draw.line(list(zip(xs.tolist(), ys.tolist())), fill=1, width=1)
    draw.text((2, 52), LABELS[health], font=font, fill=1)
    return img


def img_to_bytes(img):
    # same packing as demo_sine_wave.py: 1024 bytes for the 128x64 screen
    arr = np.array(img, dtype=bool)
    data = b""
    for i in range(8):
        data += np.packbits(arr[i*8:(i+1)*8, :], axis=0, bitorder="little").tobytes()
    return data


# how long to wait for the board's reply before assuming it's gone, in seconds
REPLY_TIMEOUT_S = 2.0
RECONNECT_EVERY_S = 1.0


def parse_arguments():
    parser = argparse.ArgumentParser(description="The ChronoPulse health monitor")
    parser.add_argument("port", type=str, nargs="?", default=None,
                        help="Serial port (optional; found automatically if left out).")
    return parser.parse_args()


def find_port(requested):
    if requested:
        return requested
    ports = sorted(glob.glob("/dev/cu.usbmodem*"))
    return ports[0] if ports else None


# keeps trying until the board is plugged in and the port opens
def connect(requested):
    while True:
        port = find_port(requested)
        if port:
            try:
                ser = serial.Serial(port, 250000,
                                    timeout=REPLY_TIMEOUT_S,
                                    write_timeout=REPLY_TIMEOUT_S)
                print(f"Connected to {port}")
                return ser
            except (serial.SerialException, OSError) as e:
                print(f"Can't open {port} yet ({e}); retrying...")
        else:
            print("Board not found; waiting for it to be plugged in...")
        time.sleep(RECONNECT_EVERY_S)


def main():
    args = parse_arguments()
    ser = connect(args.port)

    font = ImageFont.load_default()
    logo_source = Image.open(LOGO_PATH).convert("L")

    # start on the title page until the board reports its own state
    health, view, now = OFF, LOGO, 0

    while True:
        if view == LOGO and health == OFF:
            img = draw_title_page(logo_source, font)
        elif view == LOGO:
            b = beat_wave(health, now)
            img = draw_logo(logo_source, LOGO_REST_SCALE + (LOGO_PEAK_SCALE - LOGO_REST_SCALE) * b)
        elif health == OFF:
            img = draw_menu(font)
        else:
            img = draw_waves(health, now, font)

        try:
            ser.write(hdlc_encoding.encode(img_to_bytes(img)))
            # the board replies with 1 status byte + 4 bytes of its clock
            reply = ser.read(5)
        except (serial.SerialException, OSError):
            reply = b""

        if len(reply) < 5:
            print("Lost the board; reconnecting...")
            try:
                ser.close()
            except Exception:
                pass
            time.sleep(RECONNECT_EVERY_S)
            ser = connect(args.port)
            continue

        status = reply[0]
        health = status & 0b11
        view = (status >> 2) & 1
        now = int.from_bytes(reply[1:5], "little")


if __name__ == "__main__":
    main()
