import time
from pymodbus.client import ModbusSerialClient

# --- CONFIGURATION ---
PORT = "COM6"
BAUDRATE = 9600
SLAVE_ID = 1

# Keep this conservative until you have confirmed the setup is mechanically safe.
MAX_TEST_RPM = 500.0
RPM_SCALE = 10  # Drive value = RPM * 10

# If your drive manual gives you an actual speed feedback register, put it here.
# Leave as None to verify only that the commanded target speed registers match.
REG_ACTUAL_SPEED_FEEDBACK = None

# --- T3D TARGET REGISTERS ---
REG_CONTROL_MODE = 4       # P-004: Control Mode (1 = Speed Mode)
REG_TORQUE_LIMIT = 145     # P-145: Internal Torque Limit
REG_SPEED_1 = 137          # P-137: Internal Speed Register 1
REG_SPEED_2 = 138          # P-138: Internal Speed Register 2
REG_SERVO_ENABLE = 98      # P-098: Forced Internal Enable
REG_RUN_COMMAND = 96       # P-096: Internal Run Authorization
REG_VIRTUAL_DI = 110       # P-110: Virtual DI Binary Mask


def rpm_to_drive_value(rpm):
    return int(round(rpm * RPM_SCALE))


def drive_value_to_rpm(value):
    if value >= 32768:
        value -= 65536
    return value / RPM_SCALE


def ensure_ok(response, action):
    if response is None:
        raise RuntimeError(f"{action} failed: no response from drive")
    if hasattr(response, "isError") and response.isError():
        raise RuntimeError(f"{action} failed: {response}")
    return response


def write_register(client, register, value, label):
    response = client.write_register(register, value, device_id=SLAVE_ID)
    ensure_ok(response, label)


def read_register(client, register, label):
    response = client.read_holding_registers(register, count=1, device_id=SLAVE_ID)
    ensure_ok(response, label)
    return response.registers[0]


def stop_drive(client):
    print("Stopping drive...")
    for register, value, label in (
        (REG_RUN_COMMAND, 0, "Stop run command"),
        (REG_VIRTUAL_DI, 0, "Clear virtual DI mask"),
        (REG_SERVO_ENABLE, 0, "Disable servo"),
    ):
        try:
            write_register(client, register, value, label)
            time.sleep(0.05)
        except Exception as exc:
            print(f"Warning: {label} failed: {exc}")


def prepare_drive(client):
    print("Preparing drive...")
    write_register(client, REG_CONTROL_MODE, 1, "Set speed mode")
    time.sleep(0.1)

    write_register(client, REG_TORQUE_LIMIT, 150, "Set torque limit")
    time.sleep(0.1)

    # Mask 15 turns on DI1, DI2, DI3, and DI4 simultaneously.
    write_register(client, REG_VIRTUAL_DI, 15, "Set virtual DI mask")
    time.sleep(0.1)

    write_register(client, REG_SERVO_ENABLE, 1, "Enable servo")
    time.sleep(0.4)


def command_speed(client, rpm):
    scaled_speed = rpm_to_drive_value(rpm)
    print(f"\nCommanding {rpm:.1f} RPM...")

    write_register(client, REG_SPEED_1, scaled_speed, "Write speed register 1")
    write_register(client, REG_SPEED_2, scaled_speed, "Write speed register 2")
    time.sleep(0.1)

    speed_1 = drive_value_to_rpm(read_register(client, REG_SPEED_1, "Read speed register 1"))
    speed_2 = drive_value_to_rpm(read_register(client, REG_SPEED_2, "Read speed register 2"))

    print(f"Target readback: P-137={speed_1:.1f} RPM, P-138={speed_2:.1f} RPM")

    if abs(speed_1 - rpm) > 0.1 or abs(speed_2 - rpm) > 0.1:
        print("Warning: target speed readback does not match the requested speed.")


def monitor_speed(client, rpm, seconds):
    print(f"Running for {seconds:.1f} seconds. Press Ctrl+C to stop.")
    write_register(client, REG_RUN_COMMAND, 1, "Start run command")

    started_at = time.monotonic()
    while time.monotonic() - started_at < seconds:
        if REG_ACTUAL_SPEED_FEEDBACK is None:
            print(f"Commanded target: {rpm:.1f} RPM")
        else:
            actual_value = read_register(
                client,
                REG_ACTUAL_SPEED_FEEDBACK,
                "Read actual speed feedback",
            )
            actual_rpm = drive_value_to_rpm(actual_value)
            error = actual_rpm - rpm
            print(f"Target: {rpm:.1f} RPM | Actual: {actual_rpm:.1f} RPM | Error: {error:+.1f} RPM")
        time.sleep(1)

    write_register(client, REG_RUN_COMMAND, 0, "Stop run command")
    print("Run complete.")


def get_speed_from_user():
    raw = input("\nEnter target RPM, 's' to stop, or 'q' to quit: ").strip().lower()
    if raw in {"q", "quit", "exit"}:
        return "quit"
    if raw in {"s", "stop"}:
        return "stop"

    try:
        rpm = float(raw)
    except ValueError:
        print("Please enter a number, 's', or 'q'.")
        return None

    if abs(rpm) > MAX_TEST_RPM:
        print(f"Refusing {rpm:.1f} RPM because MAX_TEST_RPM is {MAX_TEST_RPM:.1f}.")
        return None

    return rpm


def main():
    client = ModbusSerialClient(
        port=PORT,
        baudrate=BAUDRATE,
        parity="N",
        stopbits=1,
        bytesize=8,
        timeout=1,
    )

    if not client.connect():
        print(f"Failed to connect to serial port {PORT}.")
        return

    print(f"Serial port {PORT} opened successfully.")

    try:
        prepare_drive(client)

        while True:
            speed = get_speed_from_user()
            if speed is None:
                continue
            if speed == "quit":
                break
            if speed == "stop":
                stop_drive(client)
                prepare_drive(client)
                continue

            command_speed(client, speed)

            run_seconds_raw = input("Run this speed for how many seconds? [5]: ").strip()
            run_seconds = float(run_seconds_raw) if run_seconds_raw else 5.0
            monitor_speed(client, speed, run_seconds)

    except KeyboardInterrupt:
        print("\nStopping sequence initiated by user.")
    except Exception as exc:
        print(f"\nError: {exc}")
    finally:
        stop_drive(client)
        client.close()
        print("Safe.")


if __name__ == "__main__":
    main()