import random
from typing import Dict, Optional

class IoTSimulator:
    """
    Simulates a robotic arm with four sensors:
    - temperature
    - vibration
    - motor current
    - position error
    Supports fault injection for testing anomaly detection.
    """

    def __init__(self, seed: int = 42):
        """
        Initialize the simulator with baseline values and fault mode.
        Args:
            seed: Random seed for reproducibility.
        """
        random.seed(seed)
        self.step = 0
        self.base_temp = 35.0
        self.base_vibration = 0.1
        self.base_current = 2.0
        self.base_position_error = 0.01
        self.fault_mode: Optional[str] = None  # 'overheat', 'vibration', 'stall', 'drift'

    def set_fault(self, fault_type: Optional[str]):
        """
        Set the fault mode.
        Args:
            fault_type: One of 'overheat', 'vibration', 'stall', 'drift', or None.
        """
        self.fault_mode = fault_type

    def read(self) -> Dict[str, float]:
        """
        Simulate a sensor reading, possibly with injected fault.
        Returns:
            Dictionary with sensor keys: temperature, vibration, motor_current, position_error.
        """
        self.step += 1
        # Normal variation
        temp = self.base_temp + random.gauss(0, 0.5)
        vib = self.base_vibration + random.gauss(0, 0.02)
        current = self.base_current + random.gauss(0, 0.1)
        pos_err = self.base_position_error + random.gauss(0, 0.005)

        # Inject fault
        if self.fault_mode == 'overheat':
            temp += 5 + 0.1 * self.step
        elif self.fault_mode == 'vibration':
            vib += 0.5 + 0.02 * self.step
        elif self.fault_mode == 'stall':
            current += 1.0 + 0.1 * self.step
        elif self.fault_mode == 'drift':
            pos_err += 0.02 + 0.002 * self.step

        return {
            'temperature': max(0, temp),
            'vibration': max(0, vib),
            'motor_current': max(0, current),
            'position_error': max(0, pos_err)
        }