"""Data classes and enums for the Robstride motor parameters/types"""
from functools import cached_property
from feathersdk.utils.feathertypes import Optional, Union, TYPE_CHECKING, Literal, Iterable, TypeVar, Generic
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

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."""
    Undervoltage = 1
    Overcurrent = 1 << 1
    Overtemp = 1 << 2
    MagneticEncoding = 1 << 3
    GridlockOverload = 1 << 4
    Uncalibrated = 1 << 5

    @classmethod
    def parse_errors(cls, extra_byte: int) -> list['RobstrideMotorError']:
        """Parse the errors from the extra byte of the can_id of a feedback message."""
        return [error for error in cls if error.value & extra_byte]


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

    @classmethod
    def parse_mode(cls, extra_byte: int) -> 'RobstrideMotorMode':
        """Parse the mode from the extra byte of the can_id of a feedback message."""
        return cls((extra_byte & 0b1100_0000) >> 6)


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: 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
    
    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"]


_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}")


@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=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=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=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=10.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=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=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=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=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=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=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=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=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=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=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, default=40.0, range=_vrange(float('-inf'), float('inf')),
    ))
    "Position loop Kp"
    spd_kp: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x701F, mode="rw", dtype=S_float32le, default=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=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=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, default=20.0, range=_vrange(float('-inf'), float('inf')),
    ))
    "Velocity mode acceleration (rad/s^2)"
    vel_max: RobstrideParam[float] = field(default_factory=lambda: RobstrideParam(
        id=0x7024, mode="rw", dtype=S_float32le, default=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=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=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=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=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=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=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
