"""
数据类型定义

定义圆柱/圆锥滚子轴承切片计算所需的输入参数和输出结果的数据类。
所有类都使用 dataclass 装饰器，支持完整的类型提示。
"""

from dataclasses import dataclass
from enum import IntEnum
from typing import Optional
import numpy as np
from numpy.typing import NDArray


class CrownType(IntEnum):
    """凸度类型枚举
    
    用于指定滚子的凸度修形类型。
    
    Attributes:
        STRAIGHT: 直线（无修形）
        ARC: 圆弧凸度
        LOGARITHMIC: 对数凸度（推荐，可实现最优应力分布）
    """
    STRAIGHT = 0
    """直线（无修形），边缘应力高"""
    
    ARC = 1
    """圆弧凸度，需要指定圆弧半径"""
    
    LOGARITHMIC = 2
    """对数凸度（推荐），理论上可实现均匀应力分布"""


class RacewayType(IntEnum):
    """滚道类型枚举
    
    用于指定接触的滚道类型。
    
    Attributes:
        PLANE: 平面接触
        INNER: 内圈滚道（凸面）
        OUTER: 外圈滚道（凹面）
    """
    PLANE = 0
    """平面接触"""
    
    INNER = 1
    """内圈滚道（凸面）"""
    
    OUTER = 2
    """外圈滚道（凹面），注意 diameter 使用负值"""


@dataclass
class RollerParams:
    """滚子几何参数
    
    定义圆柱/圆锥滚子的几何尺寸和修形参数。
    
    Attributes:
        d1: 滚子小端直径 (mm)
        d2: 滚子大端直径 (mm)，圆柱滚子 d1=d2
        length: 滚子有效长度 (mm)
        alfa: 滚子半圆锥角 (度)，圆柱滚子 alfa=0
        tilt: 滚子倾斜角度 (度)，0 表示无倾斜
        crown_type: 凸度类型，见 CrownType 枚举
        curve_q: 设计载荷 (N)，用于计算对数凸度修形量
    
    Example:
        >>> # 圆柱滚子
        >>> roller = RollerParams(
        ...     d1=10.0, d2=10.0, length=9.6,
        ...     alfa=0.0, tilt=0.0,
        ...     crown_type=CrownType.LOGARITHMIC,
        ...     curve_q=14100.0
        ... )
        
        >>> # 圆锥滚子
        >>> taper_roller = RollerParams(
        ...     d1=8.0, d2=10.0, length=12.0,
        ...     alfa=10.0, tilt=0.0,
        ...     crown_type=CrownType.LOGARITHMIC,
        ...     curve_q=20000.0
        ... )
    
    Note:
        - 对数凸度修形在 load == curve_q 时可实现均匀应力分布
        - tilt > 0 会导致一端应力集中，另一端可能脱离接触
    """
    d1: float
    """滚子小端直径 (mm)"""
    
    d2: float
    """滚子大端直径 (mm)，圆柱滚子 d1=d2"""
    
    length: float
    """滚子有效长度 (mm)"""
    
    alfa: float = 0.0
    """滚子半圆锥角 (度)，圆柱滚子 alfa=0"""
    
    tilt: float = 0.0
    """滚子倾斜角度 (度)"""
    
    crown_type: CrownType = CrownType.LOGARITHMIC
    """凸度类型"""
    
    curve_q: float = 0.0
    """设计载荷 (N)，用于对数凸度计算，0 表示使用实际载荷"""
    
    def __post_init__(self):
        """参数验证"""
        if self.d1 <= 0:
            raise ValueError(f"滚子小端直径必须为正数，当前值: {self.d1}")
        if self.d2 <= 0:
            raise ValueError(f"滚子大端直径必须为正数，当前值: {self.d2}")
        if self.length <= 0:
            raise ValueError(f"滚子有效长度必须为正数，当前值: {self.length}")
        if self.d1 > self.d2:
            import warnings
            warnings.warn(f"滚子小端直径 d1={self.d1} 大于大端直径 d2={self.d2}，请确认")
        
        # 转换枚举类型
        if isinstance(self.crown_type, int):
            self.crown_type = CrownType(self.crown_type)
    
    @property
    def is_cylindrical(self) -> bool:
        """是否为圆柱滚子"""
        return abs(self.d1 - self.d2) < 1e-6 and abs(self.alfa) < 1e-6
    
    @property
    def diameter(self) -> float:
        """滚子直径（圆柱滚子）或平均直径（圆锥滚子）"""
        return (self.d1 + self.d2) / 2.0


@dataclass
class RacewayParams:
    """滚道参数
    
    定义滚道的几何特征。
    
    Attributes:
        diameter: 滚道直径 (mm)
            - 内圈滚道：正值
            - 外圈滚道：负值（Hertz 接触理论符号约定）
        raceway_type: 滚道类型，见 RacewayType 枚举
        fai: 滚道半圆锥角 (度)，圆柱滚子轴承为 0
    
    Example:
        >>> # 内圈滚道
        >>> inner = RacewayParams(
        ...     diameter=110.0,
        ...     raceway_type=RacewayType.INNER,
        ...     fai=0.0
        ... )
        
        >>> # 外圈滚道（注意负号）
        >>> outer = RacewayParams(
        ...     diameter=-150.0,
        ...     raceway_type=RacewayType.OUTER,
        ...     fai=0.0
        ... )
    
    Note:
        外圈滚道 diameter 使用负值是 Hertz 接触理论的符号约定：
        - 正曲率半径 → 凸面（滚子、内圈）
        - 负曲率半径 → 凹面（外圈）
    """
    diameter: float
    """滚道直径 (mm)，外圈使用负值"""
    
    raceway_type: RacewayType = RacewayType.INNER
    """滚道类型"""
    
    fai: float = 0.0
    """滚道半圆锥角 (度)"""
    
    def __post_init__(self):
        """参数验证"""
        # 转换枚举类型
        if isinstance(self.raceway_type, int):
            self.raceway_type = RacewayType(self.raceway_type)
        
        # 检查符号约定
        if self.raceway_type == RacewayType.OUTER and self.diameter > 0:
            import warnings
            warnings.warn(
                f"外圈滚道 diameter 通常使用负值（当前值: {self.diameter}），"
                "请确认是否正确"
            )


@dataclass
class SliceResult:
    """切片计算结果
    
    包含滚子-滚道接触的切片应力分布计算结果。
    
    Attributes:
        n_slice: 切片数量
        stress: 各切片接触中心应力 (MPa)
        half_width: 各切片接触半宽度 (mm)
        slice_force: 各切片接触力 (N)
        deflection: 滚子变形量 (mm)
        equilibrium_load: 平衡载荷 (N)
        equilibrium_moment: 平衡力矩 (N·mm)
        converged: 是否收敛
        error_code: 错误码 (0=成功)
    
    Example:
        >>> result = slice_stress(roller, raceway, load=1340.0, n_slice=30)
        >>> print(f"最大应力: {result.max_stress:.0f} MPa")
        >>> print(f"应力分布: {result.stress}")
        >>> print(f"各切片力: {result.slice_force}")
    """
    n_slice: int
    """切片数量"""
    
    stress: NDArray[np.float64]
    """各切片接触中心应力 (MPa)"""
    
    half_width: NDArray[np.float64]
    """各切片接触半宽度 (mm)"""
    
    slice_force: NDArray[np.float64]
    """各切片接触力 (N)"""
    
    deflection: float
    """滚子变形量 (mm)"""
    
    equilibrium_load: float
    """平衡载荷 (N)，应接近输入载荷"""
    
    equilibrium_moment: float
    """平衡力矩 (N·mm)"""
    
    converged: bool = True
    """是否收敛"""
    
    error_code: int = 0
    """错误码: 0=成功, 1=方程求解失败, 2=切片数无效, 3=载荷无效"""
    
    @property
    def max_stress(self) -> float:
        """最大接触应力 (MPa)"""
        return float(np.max(self.stress)) if len(self.stress) > 0 else 0.0
    
    @property
    def min_stress(self) -> float:
        """最小接触应力 (MPa)，不包括 0（非接触区）"""
        nonzero = self.stress[self.stress > 0]
        return float(np.min(nonzero)) if len(nonzero) > 0 else 0.0
    
    @property
    def mean_stress(self) -> float:
        """平均接触应力 (MPa)，仅计算接触区"""
        nonzero = self.stress[self.stress > 0]
        return float(np.mean(nonzero)) if len(nonzero) > 0 else 0.0
    
    @property
    def contact_slices(self) -> int:
        """接触切片数量（应力 > 0）"""
        return int(np.sum(self.stress > 0))
    
    @property
    def stress_uniformity(self) -> float:
        """应力均匀度 (0~1)，1 表示完全均匀
        
        定义为: 1 - (最大应力 - 最小应力) / 平均应力
        """
        if self.mean_stress == 0:
            return 0.0
        variation = (self.max_stress - self.min_stress) / self.mean_stress
        return max(0.0, 1.0 - variation)
    
    def get_slice_positions(self, roller_length: float) -> NDArray[np.float64]:
        """获取各切片的中心位置
        
        Args:
            roller_length: 滚子有效长度 (mm)
        
        Returns:
            各切片中心位置，从 0 到 roller_length
        """
        h = roller_length / self.n_slice / 2  # 半切片宽度
        positions = np.array([h * (2 * i + 1) for i in range(self.n_slice)])
        return positions
    
    def get_slice_width(self, roller_length: float) -> float:
        """获取切片宽度
        
        Args:
            roller_length: 滚子有效长度 (mm)
        
        Returns:
            单个切片的宽度 (mm)
        """
        return roller_length / self.n_slice
    
    def get_slice_spacing(self, roller_length: float) -> float:
        """获取相邻切片中心之间的间距
        
        注意：切片间距 = 切片宽度 = roller_length / n_slice
        
        Args:
            roller_length: 滚子有效长度 (mm)
        
        Returns:
            相邻切片中心之间的间距 (mm)
        """
        return roller_length / self.n_slice
    
    def print_slice_info(self, roller_length: float) -> None:
        """打印切片划分信息
        
        Args:
            roller_length: 滚子有效长度 (mm)
        """
        width = self.get_slice_width(roller_length)
        positions = self.get_slice_positions(roller_length)
        print(f"切片划分信息:")
        print(f"  滚子有效长度: {roller_length} mm")
        print(f"  切片数量: {self.n_slice}")
        print(f"  切片宽度: {width:.6f} mm")
        print(f"  切片宽度 × 切片数 = {width * self.n_slice:.6f} mm")
        print(f"  第一片中心位置: {positions[0]:.6f} mm")
        print(f"  最后片中心位置: {positions[-1]:.6f} mm")
