scanengine-3/scanning/sc3_scan_model.py

"""
Scan Model Class.
Holds the scan configuration, and computes the
required start and end points for various specified angles.
Python implementation of SC3ScanModel.cs
"""

import math
from decimal import Decimal, InvalidOperation
from typing import List, Optional
import csv


class SC3ScanModel:
    """
    Scan Model for generating scan paths and rotated scans.
    Manages scan configuration and computes scan coordinates at various angles.
    """

    def __init__(self):
        # Private variables
        self._x_origin: Decimal = Decimal('0.0')
        self._y_origin: Decimal = Decimal('0.0')
        self._x_delta: Decimal = Decimal('0.0')
        self._y_delta: Decimal = Decimal('0.0')
        self._row_spacing: Decimal = Decimal('0.0')
        self._laser_frequency: Decimal = Decimal('2000.0')  # Default: 2000 Hz
        self._scan_velocity: Decimal = Decimal('100.0')  # Default: 100 mm/s
        self._scan_acceleration: Decimal = Decimal('0.0')
        self._scan_angles: int = 0
        self._points_required: int = 0
        self._rows_required: int = 0
        self._points_per_line: int = 0

        # Optical axis centerline in stage coordinates
        # Stage: MLS203-1
        self._optical_x_origin: Decimal = Decimal('55.0')
        self._optical_y_origin: Decimal = Decimal('37.5')

        # Data storage
        self._scan_coordinates: List[List[Decimal]] = []
        self._scan_velocities: List[List[Decimal]] = []
        self._scan_accelerations: List[List[Decimal]] = []
        self._rotated_coordinates: List[List[List[Decimal]]] = []

        # Constants
        self._deg2rad: float = math.pi / 180.0

    # Properties
    @property
    def x_origin(self) -> Decimal:
        """X coordinate of scan origin (mm)"""
        return self._x_origin

    @x_origin.setter
    def x_origin(self, value: Decimal):
        try:
            self._x_origin = Decimal(str(value))
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid x_origin value: {value}") from e

    @property
    def y_origin(self) -> Decimal:
        """Y coordinate of scan origin (mm)"""
        return self._y_origin

    @y_origin.setter
    def y_origin(self, value: Decimal):
        try:
            self._y_origin = Decimal(str(value))
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid y_origin value: {value}") from e

    @property
    def x_delta(self) -> Decimal:
        """Total X distance to scan (mm)"""
        return self._x_delta

    @x_delta.setter
    def x_delta(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val < 0:
                raise ValueError("x_delta must be non-negative")
            self._x_delta = val
            self.calculate_points_per_line()
            self.calculate_points_required()
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid x_delta value: {value}") from e

    @property
    def y_delta(self) -> Decimal:
        """Total Y distance to scan (mm)"""
        return self._y_delta

    @y_delta.setter
    def y_delta(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val < 0:
                raise ValueError("y_delta must be non-negative")
            self._y_delta = val
            self.calculate_rows_required()
            self.calculate_points_required()
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid y_delta value: {value}") from e

    @property
    def row_spacing(self) -> Decimal:
        """Spacing between scan rows (mm)"""
        return self._row_spacing

    @row_spacing.setter
    def row_spacing(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val < 0:
                raise ValueError("row_spacing must be non-negative")
            self._row_spacing = val
            self.calculate_rows_required()
            self.calculate_points_required()
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid row_spacing value: {value}") from e

    @property
    def laser_frequency(self) -> Decimal:
        """Laser pulse frequency (Hz)"""
        return self._laser_frequency

    @laser_frequency.setter
    def laser_frequency(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val <= 0:
                raise ValueError("laser_frequency must be positive")
            self._laser_frequency = val
            self.calculate_points_per_line()
            self.calculate_points_required()
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid laser_frequency value: {value}") from e

    @property
    def scan_velocity(self) -> Decimal:
        """Scan velocity (mm/s)"""
        return self._scan_velocity

    @scan_velocity.setter
    def scan_velocity(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val <= 0:
                raise ValueError("scan_velocity must be positive")
            self._scan_velocity = val
            self.calculate_points_per_line()
            self.calculate_points_required()
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid scan_velocity value: {value}") from e

    @property
    def scan_acceleration(self) -> Decimal:
        """Scan acceleration (mm/s²)"""
        return self._scan_acceleration

    @scan_acceleration.setter
    def scan_acceleration(self, value: Decimal):
        try:
            val = Decimal(str(value))
            if val < 0:
                raise ValueError("scan_acceleration must be non-negative")
            self._scan_acceleration = val
        except (ValueError, InvalidOperation) as e:
            raise ValueError(f"Invalid scan_acceleration value: {value}") from e

    @property
    def scan_angles(self) -> int:
        """Number of scan angles to compute"""
        return self._scan_angles

    @scan_angles.setter
    def scan_angles(self, value: int):
        if value < 0:
            raise ValueError("scan_angles must be non-negative")
        self._scan_angles = value
        # Only compute rotated scans if we have base scan coordinates
        if self._scan_coordinates:
            self.compute_rotated_scans()

    @property
    def points_required(self) -> int:
        """Total number of points required for a single angle scan (computed)"""
        return self._points_required

    @property
    def rows_required(self) -> int:
        """Number of rows required for the scan (computed)"""
        return self._rows_required

    @property
    def points_per_line(self) -> int:
        """Number of points per scan line (computed)"""
        return self._points_per_line

    @property
    def scan_coordinates(self) -> List[List[Decimal]]:
        """List of scan coordinates [x_start, y_start, x_end, y_end] in mm"""
        return self._scan_coordinates

    @property
    def scan_velocities(self) -> List[List[Decimal]]:
        """List of velocity vectors [vx, vy] in mm/s for each angle"""
        return self._scan_velocities

    @property
    def scan_accelerations(self) -> List[List[Decimal]]:
        """List of acceleration vectors [ax, ay] in mm/s² for each angle"""
        return self._scan_accelerations

    @property
    def rotated_coordinates(self) -> List[List[List[Decimal]]]:
        """List of rotated scan coordinates for each angle in mm"""
        return self._rotated_coordinates

    @property
    def optical_x_origin(self) -> Decimal:
        """X coordinate of optical axis origin in mm (read-only, MLS203-1 stage)"""
        return self._optical_x_origin

    @property
    def optical_y_origin(self) -> Decimal:
        """Y coordinate of optical axis origin in mm (read-only, MLS203-1 stage)"""
        return self._optical_y_origin

    # Calculation Methods
    def calculate_points_per_line(self):
        """
        Calculates the number of data points per scan line based on
        x_delta, scan_velocity, and laser_frequency.

        Formula: points = (distance / velocity) * frequency
        """
        if self._scan_velocity != 0:
            self._points_per_line = int(
                (self._x_delta / self._scan_velocity) * self._laser_frequency
            )

    def calculate_points_required(self):
        """
        Calculates the total number of data points for a complete single-angle scan.
        Also triggers computation of the zero-angle scan coordinates.

        Formula: total_points = points_per_line * rows_required
        """
        if self._rows_required != 0:
            self._points_required = self._points_per_line * self._rows_required
            self.compute_zero_scan()

    def calculate_rows_required(self):
        """
        Calculates the number of scan rows needed based on y_delta and row_spacing.

        Formula: rows = ceil(y_delta / row_spacing)
        """
        if self._row_spacing == 0:
            self._rows_required = 0
        else:
            self._rows_required = int(math.ceil(self._y_delta / self._row_spacing))

    def compute_zero_scan(self):
        """
        Computes the zero-angle (reference) scan coordinates.
        Each coordinate is [x_start, y_start, x_end, y_end].
        """
        # Ignore the zero-row case
        if self._rows_required == 0:
            return

        y_offset = Decimal('0.0')
        self._scan_coordinates.clear()

        for row in range(self._rows_required + 1):
            # Calculate x/y origin/delta for each needed row
            y_offset = Decimal(row) * self._row_spacing

            coords = [
                self._x_origin,                    # x_start
                self._y_origin + y_offset,         # y_start
                self._x_origin + self._x_delta,    # x_end
                self._y_origin + y_offset          # y_end (same as y_start for horizontal scan)
            ]
            self._scan_coordinates.append(coords)

    def _rotate_point(self, x: Decimal, y: Decimal, cosine: Decimal, sine: Decimal) -> tuple:
        """
        Rotate a point around the optical axis origin.

        Args:
            x, y: Point coordinates to rotate
            cosine, sine: Precomputed cos and sin of rotation angle

        Returns:
            Tuple of (rotated_x, rotated_y)
        """
        # Rotation transformation:
        # Xr = (X - Xo)*cos(a) + (Y - Yo)*sin(a) + Xo
        # Yr = -(X - Xo)*sin(a) + (Y - Yo)*cos(a) + Yo
        x_rot = ((x - self._optical_x_origin) * cosine +
                 (y - self._optical_y_origin) * sine +
                 self._optical_x_origin)
        y_rot = (-(x - self._optical_x_origin) * sine +
                 (y - self._optical_y_origin) * cosine +
                 self._optical_y_origin)
        return (x_rot, y_rot)

    def _compute_rotated_aoi_bbox(self, angle_rad: float) -> tuple:
        """
        Compute the bounding box of the rotated area of interest.

        Rotates the four corners of the AoI rectangle and finds the
        min/max extents to create a bounding box.

        Args:
            angle_rad: Rotation angle in radians

        Returns:
            Tuple of (min_x, min_y, max_x, max_y) as Decimals
        """
        cosine = Decimal(str(math.cos(angle_rad)))
        sine = Decimal(str(math.sin(angle_rad)))

        # Define the four corners of the AoI rectangle
        corners = [
            (self._x_origin, self._y_origin),
            (self._x_origin + self._x_delta, self._y_origin),
            (self._x_origin + self._x_delta, self._y_origin + self._y_delta),
            (self._x_origin, self._y_origin + self._y_delta)
        ]

        # Rotate all corners
        rotated_corners = []
        for x, y in corners:
            x_rot, y_rot = self._rotate_point(x, y, cosine, sine)
            rotated_corners.append((x_rot, y_rot))

        # Find bounding box extents
        x_coords = [corner[0] for corner in rotated_corners]
        y_coords = [corner[1] for corner in rotated_corners]

        return (min(x_coords), min(y_coords), max(x_coords), max(y_coords))

    def compute_rotated_scans(self):
        """
        Computes rotated scan coordinates by rotating the area of interest (AoI)
        and generating horizontal (+x direction) scans through the bounding box
        of the rotated AoI.

        The rotation covers 0 to 180 degrees with spacing determined by scan_angles.
        For each angle:
        1. Rotate the AoI rectangle around the optical axis origin
        2. Compute the bounding box of the rotated rectangle
        3. Generate horizontal scan lines through the bounding box
        """
        if self._rows_required == 0:
            return

        # Compute spacing between scans
        if self._scan_angles == 0:
            angle_spacing = 180
        else:
            angle_spacing = 180 / self._scan_angles

        self._rotated_coordinates.clear()

        # Iterate through each required angle from 0 to 180 degrees
        i = 0
        while i < 180:
            current_angle_radians = i * (math.pi / 180.0)

            # Get bounding box of rotated AoI
            min_x, min_y, max_x, max_y = self._compute_rotated_aoi_bbox(current_angle_radians)

            # Calculate the y extent of the bounding box
            y_extent = max_y - min_y

            # Determine number of rows needed for this bounding box
            if self._row_spacing == 0:
                num_rows = 0
            else:
                num_rows = int(math.ceil(y_extent / self._row_spacing))

            temp_list = []

            # Generate horizontal scan lines through the bounding box
            for row in range(num_rows + 1):
                y_offset = Decimal(row) * self._row_spacing
                y_pos = min_y + y_offset

                # Create horizontal scan line at this y position
                scan_line = [
                    min_x,      # x_start
                    y_pos,      # y_start
                    max_x,      # x_end
                    y_pos       # y_end (same as y_start for horizontal scan)
                ]
                temp_list.append(scan_line)

            self._rotated_coordinates.append(temp_list)
            i += int(round(angle_spacing))

    def compute_kinematics(self, offset: int = 0):
        """
        Computes velocity and acceleration component vectors for each scan angle.

        For each angle, decomposes the scalar velocity and acceleration into
        X and Y components based on the scan direction angle.

        Args:
            offset: Angle offset in degrees (default: 0)
        """
        if self._scan_angles == 0:
            scan_increment = 180
        else:
            scan_increment = 180 // self._scan_angles

        self._scan_velocities.clear()
        self._scan_accelerations.clear()

        for i in range(self._scan_angles):
            deg_angle = scan_increment * i + offset
            angle_rad = deg_angle * self._deg2rad

            # Compute velocity components: V = V_mag * [cos(θ), sin(θ)]
            velocities = [
                Decimal(str(math.cos(angle_rad))) * self._scan_velocity,
                Decimal(str(math.sin(angle_rad))) * self._scan_velocity
            ]

            # Compute acceleration components: A = A_mag * [cos(θ), sin(θ)]
            accels = [
                Decimal(str(math.cos(angle_rad))) * self._scan_acceleration,
                Decimal(str(math.sin(angle_rad))) * self._scan_acceleration
            ]

            self._scan_velocities.append(velocities)
            self._scan_accelerations.append(accels)

    # Export Methods
    def export_zero_scan_csv(self, filename: Optional[str] = None) -> str:
        """
        Export the zero-angle scan coordinates to a CSV file.

        Args:
            filename: Output filename. If None, generates from row count.

        Returns:
            The filename that was written

        Raises:
            ValueError: If no scan coordinates have been computed
            IOError: If file cannot be written
        """
        if not self._scan_coordinates:
            raise ValueError("No scan coordinates available. Configure scan parameters first.")

        if filename is None:
            filename = f"scantest-{self._rows_required}rows.csv"

        try:
            with open(filename, 'w', newline='') as f:
                writer = csv.writer(f)
                for coords in self._scan_coordinates:
                    writer.writerow([str(c) for c in coords])
        except IOError as e:
            raise IOError(f"Failed to write file {filename}: {e}") from e

        return filename

    def export_rotated_scan_csv(self, angle_index: int, filename: Optional[str] = None) -> str:
        """
        Export a specific rotated scan to CSV.

        Args:
            angle_index: Index of the angle to export (0-based)
            filename: Output filename. If None, generates from angle and row count.

        Returns:
            The filename that was written

        Raises:
            ValueError: If angle_index is invalid or no rotated coordinates exist
            IOError: If file cannot be written
        """
        if not self._rotated_coordinates:
            raise ValueError("No rotated coordinates available. Set scan_angles first.")

        if angle_index < 0 or angle_index >= len(self._rotated_coordinates):
            raise ValueError(
                f"Invalid angle_index {angle_index}. Must be 0-{len(self._rotated_coordinates)-1}"
            )

        if filename is None:
            angle_deg = angle_index * (180 // self._scan_angles if self._scan_angles > 0 else 180)
            filename = f"scantest-{angle_deg:03d}deg-{self._rows_required}rows.csv"

        try:
            with open(filename, 'w', newline='') as f:
                writer = csv.writer(f)
                for coords in self._rotated_coordinates[angle_index]:
                    writer.writerow([str(c) for c in coords])
        except IOError as e:
            raise IOError(f"Failed to write file {filename}: {e}") from e

        return filename

    def export_all_rotated_scans_csv(self, output_dir: str = ".") -> List[str]:
        """
        Export all rotated scans to separate CSV files.

        Args:
            output_dir: Directory to write files to (default: current directory)

        Returns:
            List of filenames that were written

        Raises:
            ValueError: If no rotated coordinates exist
            IOError: If files cannot be written
        """
        if not self._rotated_coordinates:
            raise ValueError("No rotated coordinates available. Set scan_angles first.")

        import os
        filenames = []

        for angle_index in range(len(self._rotated_coordinates)):
            angle_deg = angle_index * (180 // self._scan_angles if self._scan_angles > 0 else 180)
            filename = f"scantest-{angle_deg:03d}deg-{self._rows_required}rows.csv"
            filepath = os.path.join(output_dir, filename)

            try:
                with open(filepath, 'w', newline='') as f:
                    writer = csv.writer(f)
                    for coords in self._rotated_coordinates[angle_index]:
                        writer.writerow([str(c) for c in coords])
                filenames.append(filepath)
            except IOError as e:
                raise IOError(f"Failed to write file {filepath}: {e}") from e

        return filenames

    def export_kinematics_csv(self, filename: Optional[str] = None) -> str:
        """
        Export velocity and acceleration data to CSV.
        Format: vx, vy, ax, ay for each angle.

        Args:
            filename: Output filename. If None, generates from row count.

        Returns:
            The filename that was written

        Raises:
            ValueError: If no kinematics data has been computed
            IOError: If file cannot be written
        """
        if not self._scan_velocities or not self._scan_accelerations:
            raise ValueError("No kinematics data available. Call compute_kinematics() first.")

        if filename is None:
            filename = f"kinematics-{self._rows_required}rows.csv"

        try:
            with open(filename, 'w', newline='') as f:
                writer = csv.writer(f)
                # Optional: write header
                writer.writerow(['vx', 'vy', 'ax', 'ay'])
                for i in range(len(self._scan_velocities)):
                    row = [
                        str(self._scan_velocities[i][0]),
                        str(self._scan_velocities[i][1]),
                        str(self._scan_accelerations[i][0]),
                        str(self._scan_accelerations[i][1])
                    ]
                    writer.writerow(row)
        except IOError as e:
            raise IOError(f"Failed to write file {filename}: {e}") from e

        return filename

    # Utility Methods
    def get_angle_list(self) -> List[int]:
        """
        Get the list of scan angles in degrees.

        Returns:
            List of angles in degrees for the configured scan_angles
        """
        if self._scan_angles == 0:
            return []

        scan_increment = 180 // self._scan_angles
        return [scan_increment * i for i in range(self._scan_angles)]

    def get_scan_info(self) -> dict:
        """
        Get a dictionary with current scan configuration and computed values.

        Returns:
            Dictionary containing scan parameters and computed values
        """
        return {
            'x_origin': float(self._x_origin),
            'y_origin': float(self._y_origin),
            'x_delta': float(self._x_delta),
            'y_delta': float(self._y_delta),
            'row_spacing': float(self._row_spacing),
            'laser_frequency': float(self._laser_frequency),
            'scan_velocity': float(self._scan_velocity),
            'scan_acceleration': float(self._scan_acceleration),
            'scan_angles': self._scan_angles,
            'points_per_line': self._points_per_line,
            'rows_required': self._rows_required,
            'points_required': self._points_required,
            'optical_x_origin': float(self._optical_x_origin),
            'optical_y_origin': float(self._optical_y_origin),
            'num_scan_coordinates': len(self._scan_coordinates),
            'num_rotated_angles': len(self._rotated_coordinates),
            'angle_list': self.get_angle_list(),
        }

    def __repr__(self) -> str:
        """String representation of the scan model."""
        return (
            f"SC3ScanModel("
            f"origin=({self._x_origin},{self._y_origin}), "
            f"delta=({self._x_delta},{self._y_delta}), "
            f"rows={self._rows_required}, "
            f"angles={self._scan_angles})"
        )