# eff_tools.py
from mcp.server.fastmcp import FastMCP

import io
import numpy as np
import pandas as pd
import requests
from scipy.interpolate import RegularGridInterpolator

mcp = FastMCP("MotorEfficiencyTools")


@mcp.tool()
def calc_motor_efficiency(
    cltc_url: str,
    map_url: str,
    gear_ratio: float = 13.67,
) -> dict:
    """
    输入CLTC工况Excel与电机效率Map Excel的URL，以及传动比gear_ratio，
    输出驱动效率与回馈效率（百分比）。

    Returns:
      {
        "drive_efficiency_percent": float,
        "regen_efficiency_percent": float
      }
    """
    # 1) 下载并读取 CLTC 工况文件
    resp_cltc = requests.get(cltc_url)
    resp_cltc.raise_for_status()
    df_cycle = pd.read_excel(io.BytesIO(resp_cltc.content))
    df_cycle.columns = [str(c).strip() for c in df_cycle.columns]

    # 2) 下载并读取 电机 Map 文件（第一行是转速列名，第一列是扭矩index）
    resp_map = requests.get(map_url)
    resp_map.raise_for_status()
    df_map = pd.read_excel(io.BytesIO(resp_map.content), index_col=0)

    # 3) Map 预处理
    map_speeds = pd.to_numeric(df_map.columns, errors="coerce").values
    map_torques = pd.to_numeric(df_map.index, errors="coerce").values
    map_values = np.nan_to_num(df_map.values, nan=0.01)

    # 若是百分比，转成0~1
    if np.nanmax(map_values) > 1.0:
        map_values = map_values / 100.0

    # RegularGridInterpolator 要求坐标递增
    if map_torques[0] > map_torques[-1]:
        map_torques = map_torques[::-1]
        map_values = map_values[::-1, :]
    if map_speeds[0] > map_speeds[-1]:
        map_speeds = map_speeds[::-1]
        map_values = map_values[:, ::-1]

    interp_func = RegularGridInterpolator(
        (map_torques, map_speeds),
        map_values,
        bounds_error=False,
        fill_value=0.01,
    )

    # 4) 工况转换（轮端 -> 电机端）
    if "轮端转速" in df_cycle.columns and "轮端扭矩" in df_cycle.columns:
        wheel_speed = df_cycle["轮端转速"]
        wheel_torque = df_cycle["轮端扭矩"]
    else:
        # 兜底：假设第3列转速，第4列扭矩
        wheel_speed = df_cycle.iloc[:, 2]
        wheel_torque = df_cycle.iloc[:, 3]

    motor_speeds = wheel_speed * float(gear_ratio)
    motor_torques = wheel_torque / float(gear_ratio)

    # 5) 能量积分
    drive_mech_kwh = drive_elec_kwh = 0.0
    regen_mech_kwh = regen_elec_kwh = 0.0
    dt = 1.0 / 3600.0  # 1s -> hour

    query_points = np.column_stack((motor_torques, np.abs(motor_speeds)))
    efficiencies = interp_func(query_points)

    # 功率(kW)：T(Nm)*n(rpm)/9550
    powers_kw = (motor_torques * motor_speeds) / 9550.0

    for p_kw, eff in zip(powers_kw, efficiencies):
        if abs(p_kw) < 1e-6:
            continue

        eff = max(0.01, float(eff))

        if p_kw > 0:  # 驱动
            e_mech = float(p_kw) * dt
            e_elec = e_mech / eff
            drive_mech_kwh += e_mech
            drive_elec_kwh += e_elec
        else:  # 回馈
            e_mech = abs(float(p_kw)) * dt
            e_elec = e_mech * eff
            regen_mech_kwh += e_mech
            regen_elec_kwh += e_elec

    drive_eff_pct = (drive_mech_kwh / drive_elec_kwh * 100.0) if drive_elec_kwh > 0 else 0.0
    regen_eff_pct = (regen_elec_kwh / regen_mech_kwh * 100.0) if regen_mech_kwh > 0 else 0.0

    return {
        "drive_efficiency_percent": round(drive_eff_pct, 2),
        "regen_efficiency_percent": round(regen_eff_pct, 2),
    }


if __name__ == "__main__":
    mcp.run(transport="stdio")