"""
==========================================================
Train Automation Visualization using Arduino and Python

Arduino Code: https://stefan.bracher.info/microcontrollers.php
==========================================================

This script reads serial data from an Arduino to track the movement of two model trains
on a four-segment loop. It visualizes the train positions and dynamically updates
the powered track segment. The trains move in 6-second intervals along the tracks
when their corresponding segment is powered.

Features:
- Auto-detects Arduino serial port (or defaults to a specified port)
- Reads serial data to determine which track segment is powered
- Updates the visualization in real-time using Matplotlib
- Moves two trains along their respective tracks
- Changes track color to indicate power status

Requirements:
- Python 3.x
- Matplotlib for visualization
- PySerial for serial communication with Arduino
===========================================================
"""

import serial
import serial.tools.list_ports
import time
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.animation import FuncAnimation

def find_arduino_port():
    """Auto-detects the Arduino serial port."""
    ports = list(serial.tools.list_ports.comports())
    print("\nSearching for Arduino...")

    for port in ports:
        print(f"Found: {port.device} ({port.description})")
        if "Arduino" in port.description or "CH340" in port.description or "USB Serial" in port.description:
            print(f"Arduino detected at {port.device}\n")
            return port.device

    print("No Arduino auto-detected.")
    return None

SERIAL_PORT = find_arduino_port() or "/dev/cu.usbmodem14101"
print(f"Attempting to connect to {SERIAL_PORT}...")

try:
    ser = serial.Serial(SERIAL_PORT, 9600, timeout=1)
    time.sleep(2)
    print("Connection successful!")
except Exception as e:
    print(f"Error: Could not open serial port {SERIAL_PORT}\n{e}")
    exit()

# Track segment positions
track_positions = {
    "1": (-1, 1), "2": (1, 1), "3": (1, -1), "4": (-1, -1)
}

# Mapping track power to segments
track_segments = {"A": "1", "C": "2", "B": "3", "D": "4"}
segment_connections = {"1": "2", "2": "3", "3": "4", "4": "1"}

# Initialize figure
fig, ax = plt.subplots(figsize=(5, 5))
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title("Two Trains Loop Visualization")

# Draw track layout
track_patches = {}
for seg, (start, end) in zip(track_segments.values(), [
        ((-1, 1), (1, 1)), ((1, 1), (1, -1)), ((1, -1), (-1, -1)), ((-1, -1), (-1, 1))]):
    line, = ax.plot([start[0], end[0]], [start[1], end[1]], 'gray', lw=5)
    track_patches[seg] = line

# Train markers
train1_marker = plt.Circle(track_positions["1"], 0.2, color="red")
train2_marker = plt.Circle(track_positions["3"], 0.2, color="blue")
ax.add_patch(train1_marker)
ax.add_patch(train2_marker)

# Train states
train_positions = {"train1": "1", "train2": "3"}
train_progress = {"train1": 0, "train2": 0}
train_moving = {"train1": False, "train2": False}
active_segment = None
animation_started = False
TRAIN_SPEED = 0.0333  # Adjusted for 6-second movement

def update(frame):
    """Updates train positions and track colors based on serial data."""
    global train_positions, train_progress, train_moving, active_segment, animation_started

    if ser.in_waiting > 0:
        line = ser.readline().decode("utf-8").strip()
        print("Serial:", line)

        for track, segment in track_segments.items():
            if f"Powering Track {track}" in line:
                active_segment = segment
                if track == "A":
                    animation_started = True
                for train in train_positions:
                    if train_positions[train] == active_segment and not train_moving[train]:
                        train_progress[train] = 0
                        train_moving[train] = True
                break

    if not animation_started:
        return []

    for seg, line in track_patches.items():
        line.set_color("blue" if seg == active_segment else "gray")

    for train in train_positions:
        if train_moving[train]:
            current_seg = train_positions[train]
            next_seg = segment_connections[current_seg]
            start = track_positions[current_seg]
            end = track_positions[next_seg]
            train_x = start[0] + (end[0] - start[0]) * train_progress[train]
            train_y = start[1] + (end[1] - start[1]) * train_progress[train]

            if train == "train1":
                train1_marker.set_center((train_x, train_y))
            else:
                train2_marker.set_center((train_x, train_y))

            train_progress[train] += TRAIN_SPEED
            if train_progress[train] >= 1:
                train_progress[train] = 1
                train_moving[train] = False
                train_positions[train] = next_seg

    return train1_marker, train2_marker, *track_patches.values()

ani = FuncAnimation(fig, update, interval=200, cache_frame_data=False)
plt.show()
ser.close()