"""Data classes and enums for the Robstride motor parameters/types"""
from functools import cached_property
from feathersdk.utils.feathertypes import *
from dataclasses import dataclass, field
from feathersdk.comms.system import *
from enum import IntEnum, Enum
import math
from feathersdk.utils.common import make_exception_pickleable
from feathersdk.utils.logger import warning

if TYPE_CHECKING:
    from .robstride_motor import RobstrideMotor


class RobstrideVersion(Enum):
    RS00 = 0x00
    RS01 = 0x01
    RS02 = 0x02
    RS03 = 0x03
    RS04 = 0x04
    RS05 = 0x05
    RS06 = 0x06


RobstrideVersionType = Union[int, str, RobstrideVersion]


def get_robstride_version(ver: RobstrideVersionType) -> RobstrideVersion:
    """Get the RobstrideVersion enum from the input. Can pass int, str, or RobstrideVersion.
    
    EG: 2, '4', RobstrideVersion.RS02, 'RS04'
    """
    if isinstance(ver, RobstrideVersion):
        return ver
    if isinstance(ver, int):
        return RobstrideVersion(ver)
    if isinstance(ver, str):
        try:
            int_ver = int(ver)
        except ValueError:
            int_ver = None
        if int_ver is not None:
            return RobstrideVersion(int_ver)
        else:
            ver_str = ver.lower().replace("robstrideversion.", "")
            for v in RobstrideVersion:
                if v.name.lower() == ver_str:
                    return v
            raise ValueError(f"Invalid version string: \"{ver}\"")
    raise TypeError(f"Invalid version type: {type(ver)}")


class RobstrideRunMode(IntEnum):
    Operation = 0
    Position = 1
    Speed = 2
    Current = 3


class RobstrideMotorError(Enum):
    """Errors that can be returned in the CAN ID of a feedback message.
    
    Positive values are for motor feedback errors, negative values are for fault feedback errors. All errors 
    are powers of 2. So abs(value) is the error bit position.
    """
    UnderVoltage = 0
    OverCurrent = 1
    OverTemperature = 2
    MagneticEncoding = 3
    GridlockOverload = 4

    # Not in the CAN ID, but in the fault feedback message
    DriverChipFault = 5
    OverVoltage = 6
    EncoderNotCalibrated = 7

    @staticmethod
    def _motor_error_bits() -> Dict[int, 'RobstrideMotorError']:
        return {
            1 << 0: RobstrideMotorError.UnderVoltage, 
            1 << 1: RobstrideMotorError.OverCurrent, 
            1 << 2: RobstrideMotorError.OverTemperature, 
            1 << 3: RobstrideMotorError.MagneticEncoding, 
            1 << 4: RobstrideMotorError.GridlockOverload
        }

    @staticmethod
    def _fault_error_bits() -> Dict[int, 'RobstrideMotorError']:
        return {
            1 << 0: RobstrideMotorError.OverTemperature,
            1 << 1: RobstrideMotorError.DriverChipFault,
            1 << 2: RobstrideMotorError.UnderVoltage,
            1 << 3: RobstrideMotorError.OverVoltage,
            1 << 7: RobstrideMotorError.EncoderNotCalibrated,
            1 << 14: RobstrideMotorError.GridlockOverload,
        }

    @staticmethod
    def _parse_bits(value: int, bits: Dict[int, 'RobstrideMotorError'], name: str) -> Set['RobstrideMotorError']:
        ret = set(error for bit, error in bits.items() if bit & value)

        if value < 0 or value & ~sum(bits.keys()) != 0:
            warning(f"{name} value {value} has unknown error bits, or is negative")
        
        return ret

    @staticmethod
    def parse_motor_feedback_errors(extra_byte: int) -> Set['RobstrideMotorError']:
        """Parse the errors from the extra byte of the can_id of a feedback message."""
        return RobstrideMotorError._parse_bits(extra_byte, RobstrideMotorError._motor_error_bits(), "Motor feedback")
    
    @staticmethod
    def parse_fault_feedback_errors(fault_bytes: int) -> Set['RobstrideMotorError']:
        """Parse the fault feedback errors from the fault feedback message."""
        return RobstrideMotorError._parse_bits(fault_bytes, RobstrideMotorError._fault_error_bits(), "Fault feedback")
    
    @staticmethod
    def errors_to_int(errors: Iterable['RobstrideMotorError']) -> int:
        """Convert an iterable of errors to an integer with the bits set for the errors."""
        return sum(1 << error.value for error in errors)


class RobstrideMotorMode(Enum):
    """Motor modes that can be returned in the CAN ID of a feedback message."""
    Reset = 0
    Calibration = 1
    Run = 2
    Unknown = 3

    @staticmethod
    def parse_mode(extra_byte: int) -> 'RobstrideMotorMode':
        """Parse the mode from the extra byte of the can_id of a feedback message."""
        if (ret := RobstrideMotorMode((extra_byte & 0b1100_0000) >> 6)) is RobstrideMotorMode.Unknown:
            warning(f"Unknown Robstride motor mode parsed: {extra_byte}")
        return ret


class RobstrideMotorMsg(Enum):
    GetDeviceId = 0x00
    OperationControl = 0x01
    MotorFeedback = 0x02
    Enable = 0x03
    Disable = 0x04
    ZeroPos = 0x06
    SetID = 0x07
    BeginFlash = 0x0b
    FlashFileInfo = 0x0c
    FlashData = 0x0d
    EndFlash = 0x0e
    ReadParam = 0x11
    WriteParam = 0x12
    FaultFeedback = 0x15
    SaveMotorData = 0x16
    SetMotorBaudRate = 0x17
    SetMotorActiveReporting = 0x18
    SetMotorProtocol = 0x19


class RobstrideFeedbackMessageBounds:
    """Bounds for the feedback message."""
    ANGLE_MAP = {v: FloatRange(-4.0 * math.pi, 4.0 * math.pi) for v in RobstrideVersion}
    VELOCITY_MAP = {
        RobstrideVersion.RS00: FloatRange(-33.0, 33.0),
        RobstrideVersion.RS01: FloatRange(-44.0, 44.0),
        RobstrideVersion.RS02: FloatRange(-44.0, 44.0),
        RobstrideVersion.RS03: FloatRange(-20.0, 20.0),
        RobstrideVersion.RS04: FloatRange(-15.0, 15.0),
        RobstrideVersion.RS05: FloatRange(-50.0, 50.0),
        RobstrideVersion.RS06: FloatRange(-50.0, 50.0),
    }
    TORQUE_MAP = {
        RobstrideVersion.RS00: FloatRange(-14.0, 14.0),
        RobstrideVersion.RS01: FloatRange(-17.0, 17.0),
        RobstrideVersion.RS02: FloatRange(-17.0, 17.0),
        RobstrideVersion.RS03: FloatRange(-60.0, 60.0),
        RobstrideVersion.RS04: FloatRange(-120.0, 120.0),
        RobstrideVersion.RS05: FloatRange(-5.5, 5.5),
        RobstrideVersion.RS06: FloatRange(-36.0, 36.0),
    }
    TEMP_DIVISOR = 10  # Same for all versions


class RobstrideOperationCommand:
    """Command builder for Robstride motor operation control messages.
    
    Args:
        motor: The RobstrideMotor instance to send the command to.
        t_torque: Target torque in Nm. Must be within version-specific torque bounds.
        t_angle: Target angle in radians. Must be within [-4π, 4π].
        t_vel: Target velocity in rad/s. Must be within version-specific velocity bounds.
        kp: Proportional gain. Must be within version-specific KP bounds.
        kd: Derivative gain. Must be within version-specific KD bounds.
    """

    KP_BOUNDS_MAP = {
        RobstrideVersion.RS00: FloatRange(0.0, 500.0),
        RobstrideVersion.RS01: FloatRange(0.0, 500.0),
        RobstrideVersion.RS02: FloatRange(0.0, 500.0),
        RobstrideVersion.RS03: FloatRange(0.0, 5000.0),
        RobstrideVersion.RS04: FloatRange(0.0, 5000.0),
        RobstrideVersion.RS05: FloatRange(0.0, 500.0),
        RobstrideVersion.RS06: FloatRange(0.0, 5000.0),
    }

    KD_BOUNDS_MAP = {
        RobstrideVersion.RS00: FloatRange(0.0, 5.0),
        RobstrideVersion.RS01: FloatRange(0.0, 5.0),
        RobstrideVersion.RS02: FloatRange(0.0, 5.0),
        RobstrideVersion.RS03: FloatRange(0.0, 100.0),
        RobstrideVersion.RS04: FloatRange(0.0, 100.0),
        RobstrideVersion.RS05: FloatRange(0.0, 5.0),
        RobstrideVersion.RS06: FloatRange(0.0, 100.0),
    }

    def __init__(self, motor: 'RobstrideMotor', t_torque: float, t_angle: float, t_vel: float, kp: float, kd: float):
        self.motor: 'RobstrideMotor' = motor
        self.target_torque = t_torque
        self.target_angle = t_angle
        self.target_velocity = t_vel
        self.kp = kp
        self.kd = kd
    
    def _validate(self):
        for name, value, bounds in [
            ("target_torque", self.target_torque, RobstrideFeedbackMessageBounds.TORQUE_MAP[self.motor.version]), 
            ("target_angle", self.target_angle, RobstrideFeedbackMessageBounds.ANGLE_MAP[self.motor.version]), 
            ("target_velocity", self.target_velocity, RobstrideFeedbackMessageBounds.VELOCITY_MAP[self.motor.version]), 
            ("kp", self.kp, RobstrideOperationCommand.KP_BOUNDS_MAP[self.motor.version]), 
            ("kd", self.kd, RobstrideOperationCommand.KD_BOUNDS_MAP[self.motor.version])
        ]:
            if value < bounds[0] or value > bounds[1]:
                raise ValueError(f"{name} value ({value}) is out of bounds [{bounds[0]}, {bounds[1]}]")

    def build_command(self) -> Tuple[int, bytes]:
        """Build the command bytes for the operation control message."""
        self._validate()
        
        version = self.motor.version
        t_bytes = S_uint16be.map_and_pack(self.target_torque, *RobstrideFeedbackMessageBounds.TORQUE_MAP[version])
        a_bytes = S_uint16be.map_and_pack(self.target_angle, *RobstrideFeedbackMessageBounds.ANGLE_MAP[version])
        v_bytes = S_uint16be.map_and_pack(self.target_velocity, *RobstrideFeedbackMessageBounds.VELOCITY_MAP[version])
        kp_bytes = S_uint16be.map_and_pack(self.kp, *RobstrideOperationCommand.KP_BOUNDS_MAP[version])
        kd_bytes = S_uint16be.map_and_pack(self.kd, *RobstrideOperationCommand.KD_BOUNDS_MAP[version])
        
        data = a_bytes + v_bytes + kp_bytes + kd_bytes
        can_id = self.motor._msg_can_id(RobstrideMotorMsg.OperationControl, self.motor.motor_id)
        return ((can_id & 0xFF00_00FF) | (t_bytes[0] << 16) | (t_bytes[1] << 8), data)


# Type for the union of all possible parameter value types
RobstrideParamValue = Union[float, int, RobstrideRunMode]
T = TypeVar("T", bound=RobstrideParamValue)


@dataclass
class RobstrideParam(Generic[T]):
    id: int
    mode: Literal["r", "w", "rw"]
    dtype: StructType
    default: Dict[RobstrideVersion, T]
    range: Dict[RobstrideVersion, FloatRange]
    name: str = field(default="")
    value: T = field(default=None)

    def __post_init__(self):
        for version in RobstrideVersion:
            if version not in self.range:
                raise KeyError(f"Version {version} not found in range for parameter {self.name}")
            if self.range[version].min > self.range[version].max:
                raise ValueError(f"Invalid range {self.range[version]} for parameter {self.name}")
        assert self.default is not None, f"Default value is required"
        self.value = self.default[RobstrideVersion.RS00]
    
    def is_within_range(self, value: float, version: RobstrideVersionType) -> bool:
        """Checks if value is within the acceptible range for the parameter for the given robstride motor version"""
        version = get_robstride_version(version)
        return self.range[version].min <= value <= self.range[version].max
    
    @property
    def readable(self) -> bool:
        return self.mode in ["r", "rw"]
    
    @property
    def writable(self) -> bool:
        return self.mode in ["w", "rw"]
    
    def unpack_and_clamp(self, data: bytes, version: RobstrideVersionType) -> T:
        """Unpack the data and clamp it to the range for the parameter."""
        return self.default[version].__class__(self.range[version].clamp(self.dtype.unpack(data)))


_CURRENT_LIMIT_MAP = {
    RobstrideVersion.RS00: 16.0,
    RobstrideVersion.RS01: 23.0,
    RobstrideVersion.RS02: 16.0,
    RobstrideVersion.RS03: 43.0,
    RobstrideVersion.RS04: 90.0,
    RobstrideVersion.RS05: 11.0,
    RobstrideVersion.RS06: 57.0,
}

_SPEED_LIMIT_MAP = {
    RobstrideVersion.RS00: 33.0,
    RobstrideVersion.RS01: 44.0,
    RobstrideVersion.RS02: 33.0,
    RobstrideVersion.RS03: 20.0,
    RobstrideVersion.RS04: 20.0,
    RobstrideVersion.RS05: 50.0,
    RobstrideVersion.RS06: 50.0,
}


# Used to mark a (-inf, inf) range as one that is unknown, as opposed to a testing infinity range
_NO_KNOWN_RANGE = {v: FloatRange(float('-inf'), float('inf')) for v in RobstrideVersion}

def _vrange(
    _min: Optional[Union[float, Dict[RobstrideVersion, float]]] = None,
    _max: Optional[Union[float, Dict[RobstrideVersion, float]]] = None
) -> Dict[RobstrideVersion, FloatRange]:
    """Helper function to create a range for a parameter.
    Pass floats for that range for all versions, or a single dictionary of version -> value for the range. If that
    dictionary is passed in _min, will use range (-d[v], d[v]), otherwise if passed in _max, will use range (0, d[v]).
    """
    if isinstance(_min, (float, int)):
        return {v: FloatRange(float(_min), float(_max)) for v in RobstrideVersion}
    elif _min is not None:
        assert isinstance(_min, dict) and _max is None, f"Bad args: {_min}, {_max}"
        return {v: FloatRange(float(-_min[v]), float(_min[v])) for v in RobstrideVersion}
    elif _max is not None:
        assert isinstance(_max, dict) and _min is None, f"Bad args: {_min}, {_max}"
        return {v: FloatRange(0.0, float(_max[v])) for v in RobstrideVersion}
    else:
        raise TypeError(f"Bad args: {_min}, {_max}")


def _vdef(default: T) -> Dict[RobstrideVersion, T]:
    """Helper function to create a default value for a parameter for all versions"""
    return {v: default for v in RobstrideVersion}


@dataclass
class RobstrideParamList:
    """Contains all the parameters for all robstride versions"""
    
    run_mode: RobstrideParam[RobstrideRunMode] = field(default_factory=lambda: RobstrideParam(
        id=0x7005, mode="rw", dtype=S_uint8le, default=_vdef(RobstrideRunMode.Operation), range=_vrange(0, 3),
    ))
    "Run mode selector (0: operation, 1: position PP, 2: velocity, 3: current)"
    iq_ref: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7006, mode="rw", dtype=S_float32le, default=_vdef(0.0), range=_vrange(_CURRENT_LIMIT_MAP),
    ))
    "Current mode Iq command (A)"
    spd_ref: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x700A, mode="rw", dtype=S_float32le, default=_vdef(0.0), range=_vrange(_SPEED_LIMIT_MAP),
    ))
    "Velocity mode rotational speed command (rad/s)"
    limit_torque: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x700B, mode="rw", dtype=S_float32le, default=_vdef(5.0),
        range={
            RobstrideVersion.RS00: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS00].max),
            RobstrideVersion.RS01: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS01].max),
            RobstrideVersion.RS02: FloatRange(0, 14.0),  # Don't know why this is apparantly different
            RobstrideVersion.RS03: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS03].max),
            RobstrideVersion.RS04: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS04].max),
            RobstrideVersion.RS05: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS05].max),
            RobstrideVersion.RS06: FloatRange(0, RobstrideFeedbackMessageBounds.TORQUE_MAP[RobstrideVersion.RS06].max),
        }
    ))
    "Torque limit (Nm)"
    cur_kp: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7010, mode="rw", dtype=S_float32le, default=_vdef(0.17), range=_vrange(float('-inf'), float('inf')),
    ))
    "Current loop Kp"
    cur_ki: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7011, mode="rw", dtype=S_float32le, default=_vdef(0.012), range=_vrange(float('-inf'), float('inf')),
    ))
    "Current loop Ki"
    cur_filt_gain: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7014, mode="rw", dtype=S_float32le, default=_vdef(0.1), range=_vrange(0, 1.0),
    ))
    "Current loop filter gain"
    loc_ref: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7016, mode="rw", dtype=S_float32le, default=_vdef(0.0), range=_vrange(float('-inf'), float('inf')),
    ))
    "Position mode angle command (rad)"
    limit_spd: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7017, mode="rw", dtype=S_float32le, default=_vdef(0.0), range=_vrange(None, _SPEED_LIMIT_MAP),
    ))
    "CSP mode speed limit (rad/s)"
    limit_cur: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7018, mode="rw", dtype=S_float32le, default=_vdef(10.0), range=_vrange(None, _CURRENT_LIMIT_MAP),
    ))
    "Velocity/position mode current limit (A)"
    mech_pos: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7019, mode="r", dtype=S_float32le, default=_vdef(0.0), range=_NO_KNOWN_RANGE,
    ))
    "Mechanical angle of the loading coil (rad)"
    iqf: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701A, mode="r", dtype=S_float32le, default=_vdef(0.0), range=_vrange(_CURRENT_LIMIT_MAP),
    ))
    "Iq filter value (A)"
    mech_vel: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701B, mode="r", dtype=S_float32le, default=_vdef(0.0), range=_vrange(_SPEED_LIMIT_MAP),
    ))
    "Load speed (rad/s)"
    vbus: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701C, mode="r", dtype=S_float32le, default=_vdef(0.0), range=_NO_KNOWN_RANGE,
    ))
    "Bus voltage (V)"
    loc_kp: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701E, mode="rw", dtype=S_float32le, range=_vrange(float('-inf'), float('inf')), default={
            RobstrideVersion.RS00: 40.0,
            RobstrideVersion.RS01: 40.0,
            RobstrideVersion.RS02: 40.0,
            RobstrideVersion.RS03: 60.0,
            RobstrideVersion.RS04: 60.0,
            RobstrideVersion.RS05: 40.0,
            RobstrideVersion.RS06: 40.0,
        },
    ))
    "Position loop Kp"
    spd_kp: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701F, mode="rw", dtype=S_float32le, default=_vdef(6.0), range=_vrange(float('-inf'), float('inf')),
    ))
    "Velocity loop Kp"
    spd_ki: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7020, mode="rw", dtype=S_float32le, default=_vdef(0.02), range=_vrange(float('-inf'), float('inf')),
    ))
    "Velocity loop Ki"
    spd_filt_gain: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7021, mode="rw", dtype=S_float32le, default=_vdef(0.1), range=_vrange(float('-inf'), float('inf')),
    ))
    "Speed filter gain"
    acc_rad: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7022, mode="rw", dtype=S_float32le, range=_vrange(float('-inf'), float('inf')), default={
            RobstrideVersion.RS00: 20.0,
            RobstrideVersion.RS01: 20.0,
            RobstrideVersion.RS02: 20.0,
            RobstrideVersion.RS03: 20.0,
            RobstrideVersion.RS04: 15.0,  # Why is this different???
            RobstrideVersion.RS05: 20.0,
            RobstrideVersion.RS06: 20.0,
        },
    ))
    "Velocity mode acceleration (rad/s^2)"
    vel_max: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7024, mode="rw", dtype=S_float32le, default=_vdef(10.0), range=_NO_KNOWN_RANGE,
    ))
    "PP mode speed (rad/s)"
    acc_set: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7025, mode="rw", dtype=S_float32le, default=_vdef(10.0), range=_NO_KNOWN_RANGE,
    ))
    "PP mode acceleration (rad/s^2)"
    epscan_time: RobstrideParam[int] = field(default_factory=lambda: RobstrideParam(
        id=0x7026, mode="rw", dtype=S_uint16le, default=_vdef(1), range=_vrange(*S_uint16le.range),
    ))
    "Active report interval (1 -> 10ms, +1 adds 5ms)"
    can_timeout: RobstrideParam[int] = field(default_factory=lambda: RobstrideParam(
        id=0x7028, mode="rw", dtype=S_uint32le, default=_vdef(0), range=_vrange(*S_uint32le.range),
    ))
    "CAN timeout threshold (20000 = 1s)"
    zero_sta: RobstrideParam[int] = field(default_factory=lambda: RobstrideParam(
        id=0x7029, mode="rw", dtype=S_uint8le, default=_vdef(0), range=_vrange(*S_uint8le.range),
    ))
    "Zero flag (0: [0, 2π], 1: [-pi, pi])"
    damper: RobstrideParam[int] = field(default_factory=lambda: RobstrideParam(
        id=0x702A, mode="rw", dtype=S_uint8le, default=_vdef(0), range=_vrange(*S_uint8le.range),
    ))
    "Damping switch"
    add_offset: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x702B, mode="rw", dtype=S_float32le, default=_vdef(0.0), range=_vrange(float('-inf'), float('inf')),
    ))
    "Zero offset"

    @cached_property
    def as_dict(self) -> Dict[str, RobstrideParam]:
        """Returns dictionary of parameter name to parameter"""
        return {name: param for name, param in self.__dict__.items() if isinstance(param, RobstrideParam)}
    
    def __getitem__(self, key: str) -> RobstrideParam:
        return self.as_dict[key.lower()]
    
    def __setitem__(self, key: str, value: RobstrideParamValue) -> None:
        self.as_dict[key.lower()].value = value
    
    def __iter__(self) -> Iterable[RobstrideParam]:
        return iter(self.as_dict.values())
    
    def __post_init__(self) -> None:
        for name, param in self.as_dict.items():
            param.name = name


_ROBSTRIDE_PARAMS = RobstrideParamList()
_ROBSTRIDE_PARAMS_BY_ID = {param.id: param for param in _ROBSTRIDE_PARAMS.as_dict.values()}


def get_param_info(param_id: Union[int, str, RobstrideParam]) -> RobstrideParam:
    """Get the parameter info for the given parameter ID or name.
    
    Args:
        param_id: The parameter ID or name. Case insensitive.
    """
    try:
        if isinstance(param_id, RobstrideParam):
            return param_id
        if isinstance(param_id, str):
            return _ROBSTRIDE_PARAMS[param_id.lower()]
        return _ROBSTRIDE_PARAMS_BY_ID[param_id]
    except KeyError:
        raise UnknownRobstrideParamError(param_id=param_id)


@make_exception_pickleable
class UnknownRobstrideParamError(Exception):
    """Raised when an unknown Robstride parameter is encountered."""
    def __init__(self, *, param_id: Union[int, str]):
        super().__init__(f"Unknown Robstride parameter: {param_id}")
        self.param_id = param_id
