dune-weaver/modules/core/pattern_manager.py

import os
import threading
import time
import random
import logging
from datetime import datetime
from tqdm import tqdm
from modules.connection import connection_manager
from modules.core.state import state
from math import pi
from modules.led.led_controller import effect_playing, effect_idle

# Configure logging
logger = logging.getLogger(__name__)

# Global state
THETA_RHO_DIR = './patterns'
os.makedirs(THETA_RHO_DIR, exist_ok=True)

def list_theta_rho_files():
    files = []
    for root, _, filenames in os.walk(THETA_RHO_DIR):
        for file in filenames:
            relative_path = os.path.relpath(os.path.join(root, file), THETA_RHO_DIR)
            files.append(relative_path)
    logger.debug(f"Found {len(files)} theta-rho files")
    return files

def parse_theta_rho_file(file_path):
    """Parse a theta-rho file and return a list of (theta, rho) pairs."""
    coordinates = []
    try:
        logger.debug(f"Parsing theta-rho file: {file_path}")
        with open(file_path, 'r') as file:
            for line in file:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                try:
                    theta, rho = map(float, line.split())
                    coordinates.append((theta, rho))
                except ValueError:
                    logger.warning(f"Skipping invalid line: {line}")
                    continue
    except Exception as e:
        logger.error(f"Error reading file: {e}")
        return coordinates

    # Normalization Step
    if coordinates:
        first_theta = coordinates[0][0]
        normalized = [(theta - first_theta, rho) for theta, rho in coordinates]
        coordinates = normalized
        logger.debug(f"Parsed {len(coordinates)} coordinates from {file_path}")

    return coordinates

def get_clear_pattern_file(clear_pattern_mode, path=None):
    """Return a .thr file path based on pattern_name and table type."""
    if not clear_pattern_mode or clear_pattern_mode == 'none':
        return
    
    # Define patterns for each table type
    clear_patterns = {
        'dune_weaver': {
            'clear_from_out': './patterns/clear_from_out.thr',
            'clear_from_in': './patterns/clear_from_in.thr',
            'clear_sideway': './patterns/clear_sideway.thr'
        },
        'dune_weaver_mini': {
            'clear_from_out': './patterns/clear_from_out_mini.thr',
            'clear_from_in': './patterns/clear_from_in_mini.thr',
            'clear_sideway': './patterns/clear_sideway_mini.thr'
        },
        'dune_weaver_pro': {
            'clear_from_out': './patterns/clear_from_out_pro.thr',
            'clear_from_in': './patterns/clear_from_in_pro.thr',
            'clear_sideway': './patterns/clear_sideway_pro.thr'
        }
    }
    
    # Get patterns for current table type, fallback to standard patterns if type not found
    table_patterns = clear_patterns.get(state.table_type, clear_patterns['dune_weaver'])
    
    logger.debug(f"Clear pattern mode: {clear_pattern_mode} for table type: {state.table_type}")
    
    if clear_pattern_mode == "random":
        return random.choice(list(table_patterns.values()))

    if clear_pattern_mode == 'adaptive':
        if not path:
            logger.warning("No path provided for adaptive clear pattern")
            return random.choice(list(table_patterns.values()))
            
        coordinates = parse_theta_rho_file(path)
        if not coordinates:
            logger.warning("No valid coordinates found in file for adaptive clear pattern")
            return random.choice(list(table_patterns.values()))
            
        first_rho = coordinates[0][1]
        if first_rho < 0.5:
            return table_patterns['clear_from_out']
        else:
            return random.choice([table_patterns['clear_from_in'], table_patterns['clear_sideway']])
    else:
        if clear_pattern_mode not in table_patterns:
            return False
        return table_patterns[clear_pattern_mode]
            
            
def move_polar(theta, rho):
    """
    This functions take in a pair of theta rho coordinate, compute the distance to travel based on current theta, rho,
    and translate the motion to gcode jog command and sent to grbl. 
    
    Since having similar steps_per_mm will make x and y axis moves at around the same speed, we have to scale the 
    x_steps_per_mm and y_steps_per_mm so that they are roughly the same. Here's the range of motion:
    
    X axis (angular): 50mm = 1 revolution
    Y axis (radial): 0 => 20mm = theta 0 (center) => 1 (perimeter)
    
    Args:
        theta (_type_): _description_
        rho (_type_): _description_
    """
    # Adding soft limit to reduce hardware sound
    soft_limit_inner = 0.01
    if rho < soft_limit_inner:
        rho = soft_limit_inner
    
    soft_limit_outter = 0.015
    if rho > (1-soft_limit_outter):
        rho = (1-soft_limit_outter)
    
    if state.gear_ratio == 6.25:
        x_scaling_factor = 2
        y_scaling_factor = 3.7
    else:
        x_scaling_factor = 2
        y_scaling_factor = 5
    
    delta_theta = theta - state.current_theta
    delta_rho = rho - state.current_rho
    x_increment = delta_theta * 100 / (2 * pi * x_scaling_factor)  # Added -1 to reverse direction
    y_increment = delta_rho * 100 / y_scaling_factor
    
    x_total_steps = state.x_steps_per_mm * (100/x_scaling_factor)
    y_total_steps = state.y_steps_per_mm * (100/y_scaling_factor)
        
    offset = x_increment * (x_total_steps * x_scaling_factor / (state.gear_ratio * y_total_steps * y_scaling_factor))

    if state.gear_ratio == 6.25:
        y_increment -= offset
    else:
        y_increment += offset
    
    new_x_abs = state.machine_x + x_increment
    new_y_abs = state.machine_y + y_increment
    
    # dynamic_speed = compute_dynamic_speed(rho, max_speed=state.speed)
    
    connection_manager.send_grbl_coordinates(round(new_x_abs, 3), round(new_y_abs,3), state.speed)
    state.current_theta = theta
    state.current_rho = rho
    state.machine_x = new_x_abs
    state.machine_y = new_y_abs
    
def pause_execution():
    logger.info("Pausing pattern execution")
    with state.pause_condition:
        state.pause_requested = True
    return True

def resume_execution():
    logger.info("Resuming pattern execution")
    with state.pause_condition:
        state.pause_requested = False
        state.pause_condition.notify_all()
    return True
    
def reset_theta():
    logger.info('Resetting Theta')
    state.current_theta = 0
    connection_manager.update_machine_position()

def set_speed(new_speed):
    state.speed = new_speed
    logger.info(f'Set new state.speed {new_speed}')

def run_theta_rho_file(file_path):
    """Run a theta-rho file by sending data in optimized batches with tqdm ETA tracking."""
    
    # Check if connection is still valid, if not, restart
    # if not connection_manager.get_status_response() and isinstance(state.conn, connection_manager.WebSocketConnection):
    #     logger.info('Cannot get status response, restarting connection')
    #     connection_manager.restart_connection(home=False)
    # if (state.conn.is_connected() if state.conn else False):
    #     logger.error('Connection not established')
    #     return
    # if not file_path:
    #     return
    
    coordinates = parse_theta_rho_file(file_path)
    total_coordinates = len(coordinates)

    if total_coordinates < 2:
        logger.warning("Not enough coordinates for interpolation")
        state.current_playing_file = None
        state.execution_progress = None
        return

    state.execution_progress = (0, total_coordinates, None)
    
    # stop actions without resetting the playlist
    stop_actions(clear_playlist=False)

    state.current_playing_file = file_path
    state.execution_progress = (0, 0, None)
    state.stop_requested = False
    logger.info(f"Starting pattern execution: {file_path}")
    logger.info(f"t: {state.current_theta}, r: {state.current_rho}")
    reset_theta()
    
    if state.led_controller:
            effect_playing(state.led_controller)
            
    
    with tqdm(
        total=total_coordinates,
        unit="coords",
        desc=f"Executing Pattern {file_path}",
        dynamic_ncols=True,
        disable=False,  # Force enable the progress bar
        mininterval=1.0  # Optional: reduce update frequency to prevent flooding
    ) as pbar:
        for i, coordinate in enumerate(coordinates):
            theta, rho = coordinate
            if state.stop_requested:
                logger.info("Execution stopped by user")
                if state.led_controller:
                    effect_idle(state.led_controller)
                break
            
            if state.skip_requested:
                logger.info("Skipping pattern...")
                connection_manager.check_idle()
                if state.led_controller:
                    effect_idle(state.led_controller)
                break

            # Wait for resume if paused
            if state.pause_requested:
                logger.info("Execution paused...")
                if state.led_controller:
                    effect_idle(state.led_controller)
                pause_event.wait()
                logger.info("Execution resumed...")
                if state.led_controller:
                    effect_playing(state.led_controller)

            move_polar(theta, rho)
            
            if i != 0:
                pbar.update(1)
                estimated_remaining_time = (total_coordinates - i) / pbar.format_dict['rate'] if pbar.format_dict['rate'] and total_coordinates else 0
                elapsed_time = pbar.format_dict['elapsed']
                state.execution_progress = (i, total_coordinates, estimated_remaining_time, elapsed_time)

    connection_manager.check_idle()

    state.current_playing_file = None
    state.execution_progress = None
    logger.info("Pattern execution completed")

def run_theta_rho_files(file_paths, pause_time=0, clear_pattern=None, run_mode="single", shuffle=False):
    """Run multiple .thr files in sequence with options."""
    state.stop_requested = False
    
    # Set initial playlist state
    state.playlist_mode = run_mode
    state.current_playlist_index = 0

    try:
        while True:
            # Construct the complete pattern sequence
            pattern_sequence = []
            for path in file_paths:
                # Add clear pattern if specified
                if clear_pattern and clear_pattern != 'none':
                    clear_file_path = get_clear_pattern_file(clear_pattern, path)
                    if clear_file_path:
                        pattern_sequence.append(clear_file_path)
                
                # Add main pattern
                pattern_sequence.append(path)

            # Shuffle if requested
            if shuffle:
                # Get pairs of patterns (clear + main) to keep them together
                pairs = [pattern_sequence[i:i+2] for i in range(0, len(pattern_sequence), 2)]
                random.shuffle(pairs)
                # Flatten the pairs back into a single list
                pattern_sequence = [pattern for pair in pairs for pattern in pair]
                logger.info("Playlist shuffled")

            # Set the playlist to the first pattern
            state.current_playlist = pattern_sequence
            # Execute the pattern sequence
            for idx, file_path in enumerate(pattern_sequence):
                state.current_playlist_index = idx
                if state.stop_requested:
                    logger.info("Execution stopped")
                    return

                # Update state for main patterns only
                logger.info(f"Running pattern {file_path}")
                
                # Execute the pattern
                run_theta_rho_file(file_path)

                # Handle pause between patterns
                if idx < len(pattern_sequence) - 1 and not state.stop_requested and pause_time > 0 and not state.skip_requested:
                    logger.info(f"Pausing for {pause_time} seconds")
                    pause_start = time.time()
                    while time.time() - pause_start < pause_time:
                        if state.skip_requested:
                            logger.info("Pause interrupted by stop/skip request")
                            break
                        time.sleep(1)
                    
                state.skip_requested = False

            if run_mode == "indefinite":
                logger.info("Playlist completed. Restarting as per 'indefinite' run mode")
                if pause_time > 0:
                    logger.debug(f"Pausing for {pause_time} seconds before restarting")
                    time.sleep(pause_time)
                continue
            else:
                logger.info("Playlist completed")
                break

    finally:
        # Clean up progress update task
        # if progress_update_task:
        #     progress_update_task.cancel()
        #     try:
        #         await progress_update_task
        #     except asyncio.CancelledError:
        #         pass
        #     progress_update_task = None
            
        # Clear all state variables
        state.current_playing_file = None
        state.execution_progress = None
        state.current_playlist = None
        state.current_playlist_index = None
        state.playlist_mode = None
        
        if state.led_controller:
            effect_idle(state.led_controller)
        
        logger.info("All requested patterns completed (or stopped) and state cleared")

def stop_actions(clear_playlist = True):
    """Stop all current actions."""
    with state.pause_condition:
        state.pause_requested = False
        state.stop_requested = True
        state.current_playing_file = None
        state.execution_progress = None
        state.is_clearing = False
        if clear_playlist:
            # Clear playlist state
            state.current_playlist = None
            state.current_playlist_index = None
            state.playlist_mode = None
        state.pause_condition.notify_all()
        connection_manager.update_machine_position()

def get_status():
    """Get the current execution status."""
    # Update state.is_clearing based on current file
    if state.current_playing_file in CLEAR_PATTERNS.values():
        state.is_clearing = True
    else:
        state.is_clearing = False

    return {
        "ser_port": state.port,
        "stop_requested": state.stop_requested,
        "pause_requested": state.pause_requested,
        "current_playing_file": state.current_playing_file,
        "execution_progress": state.execution_progress,
        "current_playing_index": state.current_playlist_index,
        "current_playlist": state.current_playlist,
        "is_clearing": state.is_clearing,
        "current_speed": state.speed
    }