
import numpy as np
import json
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import Dict
from bstart_trb.models import SliceResult
import pandas as pd
from bstart_trb import (
    slice_stress,
    batch_slice_stress,
    RollerParams,
    RacewayParams,
    CrownType,
    RacewayType,
    is_fortran_available,
)

# 检查 Fortran 模块是否可用
print(f"Fortran 计算核心: {'已加载' if is_fortran_available() else '未编译'}")


df = pd.read_excel('demo.xlsx')


print(df.head())


# 按滚子 j 聚合数据
df_agg = df.groupby('j').agg({
    'psi_j,[rad]': 'mean',
    'q_jk,[N]': 'sum',
    'k': 'count'  # 切片数
}).reset_index()

df_agg.rename(columns={'k': 'n_slice'}, inplace=True)

print("聚合结果:")
print(df_agg)


# 过滤掉 q_jk,[N] 为 0 的行
df_agg = df_agg[df_agg['q_jk,[N]'] != 0]

print("\n过滤后结果:")
print(df_agg)



def calculate_roller_stress(
    j: int,
    tilt_rad: float,
    load: float,
    n_slice: int,
    # 固定参数
    d1: float = 9,
    d2: float = 9,
    length: float = 13.6,
    alfa: float = 0.0,
    crown_type: CrownType = CrownType.STRAIGHT,
    curve_q: float = 500 * 1000,
    raceway_diameter: float = -52.5,
    raceway_type: RacewayType = RacewayType.OUTER,
    fai: float = 0.0,
) -> Dict:
    """
    计算单个滚子的切片应力分布
    
    Parameters
    ----------
    j : int
        滚子编号
    tilt_rad : float
        倾斜角 (弧度)
    load : float
        载荷 (N)
    n_slice : int
        切片数
    
    Returns
    -------
    Dict
        包含计算结果的字典
    """
    # 将弧度转换为角度
    tilt_deg = tilt_rad * 180 / np.pi
    
    # 定义滚子参数
    roller = RollerParams(
        d1=d1,
        d2=d2,
        length=length,
        alfa=alfa,
        tilt=tilt_deg,
        crown_type=crown_type,
        curve_q=curve_q,
    )
    
    # 定义滚道参数
    raceway = RacewayParams(
        diameter=raceway_diameter,
        raceway_type=raceway_type,
        fai=fai,
    )
    
    # 计算切片应力
    result = slice_stress(roller, raceway, load=load, n_slice=n_slice)
    
    # 获取切片位置
    positions = result.get_slice_positions(roller.length)
    
    return {
        'j': j,
        'roller': roller,
        'raceway': raceway,
        'result': result,
        'positions': positions,
        'load': load,
        'n_slice': n_slice,
        'tilt_deg': tilt_deg,
    }



# 批量计算所有滚子
results = []
for _, row in df_agg.iterrows():
    calc_result = calculate_roller_stress(
        j=int(row['j']),
        tilt_rad=row['psi_j,[rad]']/2,
        load=row['q_jk,[N]'],
        n_slice=int(row['n_slice']),
    )
    results.append(calc_result)
    
    # 打印计算结果
    r = calc_result['result']
    print(f"\n滚子 {calc_result['j']}:")
    print(f"  倾斜角: {calc_result['tilt_deg']:.6f}°")
    print(f"  载荷: {calc_result['load']:.2f} N")
    print(f"  切片数: {calc_result['n_slice']}")
    print(f"  收敛: {r.converged}")
    print(f"  最大应力: {r.max_stress:.0f} MPa")
    print(f"  平均应力: {r.mean_stress:.0f} MPa")

print(f"\n共计算 {len(results)} 个滚子")


# 导出 JSON 结果
json_data = []
for calc in results:
    r = calc['result']
    json_data.append({
        'j': calc['j'],
        'tilt_deg': calc['tilt_deg'],
        'load': calc['load'],
        'n_slice': calc['n_slice'],
        'max_stress': r.max_stress,
        'mean_stress': r.mean_stress,
        'stress': r.stress.tolist() if hasattr(r.stress, 'tolist') else list(r.stress),
        'half_width': r.half_width.tolist() if hasattr(r.half_width, 'tolist') else list(r.half_width),
        'slice_force': r.slice_force.tolist() if hasattr(r.slice_force, 'tolist') else list(r.slice_force),
    })

with open('slice_results.json', 'w', encoding='utf-8') as f:
    json.dump(json_data, f, indent=2, ensure_ascii=False)

print("\n计算结果已导出为 slice_results.json")



# 绘制复合图
import plotly.express as px

# 使用不同颜色区分不同滚子
colors = px.colors.qualitative.Set1

fig = make_subplots(
    rows=1, cols=2,
    subplot_titles=("Slice Stress Distribution", "Slice Contact Half-Width Distribution")
)

for i, calc in enumerate(results):
    j = calc['j']
    positions = calc['positions']
    result = calc['result']
    
    # Stress distribution
    fig.add_trace(
        go.Scatter(
            x=positions, y=result.stress,
            mode='lines+markers',
            name=f'Roller {j}',
            marker=dict(size=6),
            hovertemplate=f'Roller {j}<br>Position: %{{x:.4f}} mm<br>Stress: %{{y:.0f}} MPa<extra></extra>',
            legendgroup=f'roller_{j}',
        ),
        row=1, col=1
    )
    
    # Contact half-width distribution
    fig.add_trace(
        go.Scatter(
            x=positions, y=result.half_width,
            mode='lines+markers',
            name=f'Roller {j}',
            marker=dict(size=6),
            hovertemplate=f'Roller {j}<br>Position: %{{x:.4f}} mm<br>Half-width: %{{y:.4f}} mm<extra></extra>',
            legendgroup=f'roller_{j}',
            showlegend=False,
        ),
        row=1, col=2
    )

fig.update_xaxes(title_text="Position along Roller Length (mm)", row=1, col=1)
fig.update_xaxes(title_text="Position along Roller Length (mm)", row=1, col=2)
fig.update_yaxes(title_text="Contact Stress (MPa)", row=1, col=1)
fig.update_yaxes(title_text="Contact Half-Width (mm)", row=1, col=2)

fig.update_layout(
    title_text=f"Roller Slice Stress Distribution Comparison ({len(results)} Rollers)",
    height=500,
    showlegend=True,
    legend=dict(
        orientation="h",
        yanchor="bottom",
        y=1.02,
        xanchor="center",
        x=0.5
    )
)

# Show chart
fig.show()

# Export first chart
fig.write_html("stress_distribution.html")
print("\nStress distribution chart exported to stress_distribution.html")



# Plot slice force comparison: Original data vs Calculated results
# Filter rollers with load
roller_ids = df_agg['j'].tolist()
df_filtered = df[df['j'].isin(roller_ids)]

fig2 = go.Figure()

for i, j in enumerate(roller_ids):
    # Get original data
    df_roller = df_filtered[df_filtered['j'] == j]
    
    # Original data q_jk (solid line)
    fig2.add_trace(
        go.Scatter(
            x=df_roller['k'],
            y=df_roller['q_jk,[N]'],
            mode='lines+markers',
            name=f'Roller {j} (Original)',
            marker=dict(size=6),
            hovertemplate=f'Roller {j} (Original)<br>Slice k: %{{x}}<br>q_jk: %{{y:.2f}} N<extra></extra>',
        )
    )
    
    # Get calculated slice_force
    calc_result = next((r for r in results if r['j'] == j), None)
    if calc_result:
        result = calc_result['result']
        # Slice index starts from 1, matching original data
        k_values = list(range(1, len(result.slice_force) + 1))
        
        # Calculated slice_force (dashed line)
        fig2.add_trace(
            go.Scatter(
                x=k_values,
                y=result.slice_force,
                mode='lines+markers',
                name=f'Roller {j} (Calculated)',
                marker=dict(size=6, symbol='diamond'),
                hovertemplate=f'Roller {j} (Calculated)<br>Slice k: %{{x}}<br>slice_force: %{{y:.2f}} N<extra></extra>',
            )
        )

fig2.update_xaxes(title_text="Slice Index k")
fig2.update_yaxes(title_text="Slice Load (N)")

fig2.update_layout(
    title_text="Slice Force Comparison: Original Data vs Calculated Results",
    height=500,
    showlegend=True,
    legend=dict(
        orientation="h",
        yanchor="bottom",
        y=1.02,
        xanchor="center",
        x=0.5
    )
)

fig2.show()

# Export second chart
fig2.write_html("slice_force_comparison.html")
print("\nSlice force comparison chart exported to slice_force_comparison.html")


