"""
Python 包装层

提供圆柱/圆锥滚子轴承切片应力计算的高级 Python 接口。
封装底层 Fortran 计算核心（通过 f2py）。
"""

from typing import Optional, Sequence, Callable, Tuple
import numpy as np
from numpy.typing import NDArray

from .models import (
    RollerParams,
    RacewayParams,
    CrownType,
    RacewayType,
    SliceResult,
)

# 尝试导入 f2py 编译的模块
_FORTRAN_AVAILABLE = False
_taper_calculate: Optional[Callable[..., Tuple[
    NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], float, float, float, int
]]] = None

try:
    from . import _taper_slice
    _taper_calculate = _taper_slice.taper_calculate
    _FORTRAN_AVAILABLE = True
except ImportError:
    import warnings
    warnings.warn(
        "Fortran 计算核心未编译，请运行 'python build.py' 编译模块。"
        "当前切片计算功能不可用。"
    )


def slice_stress(
    roller: RollerParams,
    raceway: RacewayParams,
    load: float,
    n_slice: int = 30,
) -> SliceResult:
    """计算滚子-滚道接触的切片应力分布
    
    使用影响系数法（Influence Coefficient Method）计算滚子与滚道
    接触时各切片的应力分布，考虑凸度修形和滚子倾斜的影响。
    
    Args:
        roller: 滚子几何参数
        raceway: 滚道参数
        load: 滚子载荷 (N)
        n_slice: 切片数量 (1-50)，默认 30
    
    Returns:
        SliceResult: 包含各切片应力、接触半宽度等计算结果
    
    Raises:
        RuntimeError: Fortran 模块未编译时抛出
        ValueError: 参数无效时抛出
    
    Example:
        >>> from slice_wrapper import slice_stress, RollerParams, RacewayParams
        >>> from slice_wrapper import CrownType, RacewayType
        >>> 
        >>> # 定义滚子参数
        >>> roller = RollerParams(
        ...     d1=10.0, d2=10.0, length=9.6,
        ...     crown_type=CrownType.LOGARITHMIC,
        ...     curve_q=14100.0
        ... )
        >>> 
        >>> # 定义外圈滚道参数（注意直径为负值）
        >>> raceway = RacewayParams(
        ...     diameter=-150.0,
        ...     raceway_type=RacewayType.OUTER
        ... )
        >>> 
        >>> # 计算切片应力
        >>> result = slice_stress(roller, raceway, load=1340.0, n_slice=30)
        >>> print(f"最大应力: {result.max_stress:.0f} MPa")
        >>> print(f"应力均匀度: {result.stress_uniformity:.2%}")
    
    Note:
        - 切片数 n_slice 最大支持 50
        - 对数凸度 (crown_type=2) 在 load == curve_q 时可实现均匀应力分布
        - 外圈滚道的 diameter 应使用负值（Hertz 接触符号约定）
    
    Theory:
        该函数使用影响系数法求解滚子-滚道接触问题：
        1. 将滚子沿轴向分成 n_slice 个切片
        2. 建立切片间的弹性耦合关系（影响系数矩阵）
        3. 使用高斯消元法求解线性方程组
        4. 迭代收敛得到各切片的接触应力分布
    """
    if not _FORTRAN_AVAILABLE or _taper_calculate is None:
        raise RuntimeError(
            "Fortran 计算核心未编译。请先运行 'python build.py' 编译模块。"
        )
    
    # 参数验证
    if load <= 0:
        raise ValueError(f"载荷必须为正数，当前值: {load}")
    if not 1 <= n_slice <= 50:
        raise ValueError(f"切片数必须在 1-50 之间，当前值: {n_slice}")
    
    # 确定设计载荷
    curve_q = roller.curve_q if roller.curve_q > 0 else load
    
    # 调用 Fortran 计算（此时 _taper_calculate 已确保不为 None）
    p, a, q_slice, dc, q1, qm, ierr = _taper_calculate(
        d1=roller.d1,
        d2=roller.d2,
        rl=roller.length,
        alfa=roller.alfa,
        q=load,
        tilt=roller.tilt,
        m1=int(roller.crown_type),
        n_slice=n_slice,
        m2=int(raceway.raceway_type),
        dr1=raceway.diameter,
        fai=raceway.fai,
        curveq=curve_q,
    )
    
    # 构建结果
    return SliceResult(
        n_slice=n_slice,
        stress=np.array(p, dtype=np.float64),
        half_width=np.array(a, dtype=np.float64),
        slice_force=np.array(q_slice, dtype=np.float64),
        deflection=float(dc),
        equilibrium_load=float(q1),
        equilibrium_moment=float(qm),
        converged=(ierr == 0),
        error_code=int(ierr),
    )


def batch_slice_stress(
    roller: RollerParams,
    raceway: RacewayParams,
    loads: Sequence[float],
    n_slice: int = 30,
) -> list[SliceResult]:
    """批量计算多个载荷下的切片应力分布
    
    对多个载荷值进行切片应力计算，适用于载荷分布已知的情况。
    
    Args:
        roller: 滚子几何参数
        raceway: 滚道参数
        loads: 载荷序列 (N)，支持 list[int], list[float], NDArray 等
        n_slice: 切片数量，默认 30
    
    Returns:
        list[SliceResult]: 各载荷对应的计算结果列表
    
    Example:
        >>> # 假设已知各滚子的载荷分布
        >>> loads = [1340.0, 1128.3, 689.2, 324.8]
        >>> results = batch_slice_stress(roller, raceway, loads)
        >>> for i, result in enumerate(results):
        ...     print(f"载荷 {loads[i]:.1f}N: 最大应力 {result.max_stress:.0f} MPa")
    """
    results = []
    for load in loads:
        if load > 0:
            result = slice_stress(roller, raceway, load, n_slice)
            results.append(result)
        else:
            # 无载荷，返回零结果
            results.append(SliceResult(
                n_slice=n_slice,
                stress=np.zeros(n_slice),
                half_width=np.zeros(n_slice),
                slice_force=np.zeros(n_slice),
                deflection=0.0,
                equilibrium_load=0.0,
                equilibrium_moment=0.0,
                converged=True,
                error_code=0,
            ))
    return results


def is_fortran_available() -> bool:
    """检查 Fortran 计算核心是否可用
    
    Returns:
        bool: True 表示 Fortran 模块已编译可用
    """
    return _FORTRAN_AVAILABLE
