# -*- coding: utf-8 -*-
"""
MARS SSSR 轴向磁拉力 — 正式计算
=================================
第5轮探测确认: AFM 力数据在 3D lumped 力图, 命名沿用径向机惯例, 其 "Fr"
(法向力) 在 AFM 2.5D 展开模型中即轴向力:
  Fr_Rotor_OL_Lumped / Fr_Rotor_OC_Lumped (转子, 负载/空载)
  Fr_Stator_OL_Lumped / Fr_Stator_OC_Lumped (定子, 反作用)
节点: 转子 10 (36°步, 首尾重复共11点), 定子 12 (30°步, 共13点); 单位 N。
两个径向切片 (sec1 r=28.25mm, sec2 r=34.75mm) 分别读, 节点求和+切片求和
得净轴向力; 按时间步扫描得波形。

校核: (a) 定子合力 ≈ -转子合力; (b) Σ(Ft×r) ≈ 电磁转矩图;
      (c) 解析 F ≈ A/(2μ0)·mean(B²)。

用法: python axial_force_final.py [--quit]
"""
import json
import math
import os
import sys
import time

BASE = os.path.dirname(os.path.abspath(__file__))
MOT_SRC = os.path.join(BASE, "MARS-12S10P_SSSR_D76-C150_V5.0-0819.mot")
OUT_DIR = os.path.join(BASE, "output_motorcad")

MU0 = 4e-7 * math.pi
SEC_RADII_MM = [28.25, 34.75]      # AFM_SectionCentreRadius_Array
MAX_TSTEPS = 64
MAX_NODES = 40


def stats(ys):
    if not ys:
        return None
    return {"mean": sum(ys) / len(ys), "min": min(ys), "max": max(ys),
            "pk2pk": max(ys) - min(ys), "n": len(ys)}


def read_nodes(mc, graph, sec, tstep):
    """读某时间步的全部节点 (x=角度, y=力N); 首尾重复点保留由调用方处理。"""
    xs, ys = [], []
    for i in range(MAX_NODES):
        try:
            x, y = mc.get_magnetic_3d_graph_point(graph, sec, i, tstep)
        except Exception:
            break
        xs.append(x)
        ys.append(y)
    return xs, ys


def net_force_series(mc, graph):
    """净力时间序列: 对两切片、去重节点求和; 返回 (series, meta)。"""
    series = []
    meta = {"sections": {}}
    for tstep in range(MAX_TSTEPS):
        total = 0.0
        got = False
        for sec in (1, 2):
            xs, ys = read_nodes(mc, graph, sec, tstep)
            if not ys:
                continue
            got = True
            # 首尾重复 (0°与360°同一节点) 则去掉末点
            n_unique = len(ys) - 1 if (len(xs) > 1 and
                                       abs(xs[-1] - xs[0] - 360.0) < 1e-6) \
                else len(ys)
            total += sum(ys[:n_unique])
            if tstep == 0:
                meta["sections"][sec] = {"n_points": len(ys),
                                         "n_unique": n_unique, "x": xs}
        if not got:
            break
        series.append(total)
    return series, meta


def torque_from_ft(mc, graph):
    """t=0 时刻 Σ(Ft×r) 粗校核 (Nm)。"""
    tq = 0.0
    for sec, r_mm in zip((1, 2), SEC_RADII_MM):
        xs, ys = read_nodes(mc, graph, sec, 0)
        if not ys:
            return None
        n_unique = len(ys) - 1 if (len(xs) > 1 and
                                   abs(xs[-1] - xs[0] - 360.0) < 1e-6) \
            else len(ys)
        tq += sum(ys[:n_unique]) * (r_mm * 1e-3)
    return tq


def read_2d(mc, graph, maxpts=64):
    xs, ys = [], []
    for i in range(maxpts):
        try:
            x, y = mc.get_magnetic_graph_point(graph, i)
        except Exception:
            break
        xs.append(x)
        ys.append(y)
    return xs, ys


def main(argv):
    quit_after = "--quit" in argv
    os.makedirs(OUT_DIR, exist_ok=True)
    ts = time.strftime("%m%d_%H%M%S")

    from ansys.motorcad.core import MotorCAD, set_motorcad_exe
    # Motor-CAD 2026R1 新装机器可能未注册 MOTORCAD_ACTIVEX (pymotorcad 0.8.x
    # 仍依赖它); 此时显式定位 exe, 可用环境变量 MOTORCAD_EXE 覆盖路径。
    if not os.environ.get("MOTORCAD_ACTIVEX"):
        exe = os.environ.get(
            "MOTORCAD_EXE",
            r"D:\Program Files\ANSYS Inc\v261\motorcad\MotorCAD.exe")
        if os.path.isfile(exe):
            set_motorcad_exe(exe)
    print("启动 Motor-CAD (前台) ...")
    mc = MotorCAD()
    # /SCRIPTING 模式在部分机器上窗口创建但不显示 (任务栏有图标点不开),
    # 强制可见; 已可见时无副作用 (2026-08-26 同事复现机实测该问题)
    try:
        mc.set_visible(True)
    except Exception as e:
        print("  [提示] set_visible 失败 (不影响计算): %s" % e)
    results = {"when": ts, "source_mot": os.path.basename(MOT_SRC),
               "convention_note": ("AFM 2.5D 展开模型中 Fr(法向)=轴向力; "
                                   "OL=负载(RMS 21A), OC=空载开路")}
    try:
        mc.load_from_file(MOT_SRC)
        out_mot = os.path.join(OUT_DIR, "MARS_SSSR_axialF_%s.mot" % ts)
        mc.save_to_file(out_mot)
        results["work_mot"] = out_mot

        for key, names in [("RMSCurrent_A", ["RMSCurrent"]),
                           ("ShaftSpeed_rpm", ["ShaftSpeed"]),
                           ("Airgap_mm", ["Airgap"]),
                           ("Stator_Lam_Dia_mm", ["Stator_Lam_Dia"]),
                           ("Stator_Bore_mm", ["Stator_Bore"])]:
            try:
                results[key] = mc.get_variable(names[0])
            except Exception:
                results[key] = None
        for var in ["ElectromagneticForcesCalc_Load",
                    "ElectromagneticForcesCalc_OC"]:
            mc.set_variable(var, True)

        print("求解 (负载点 RMS %sA, OC+OL 力同算) ..." % results["RMSCurrent_A"])
        t0 = time.time()
        mc.do_magnetic_calculation()
        results["solve_seconds"] = time.time() - t0
        print("  耗时 %.1f s" % results["solve_seconds"])

        # ---- 净轴向力: 转子/定子 x OL/OC ----
        forces = {}
        for graph in ["Fr_Rotor_OL_Lumped", "Fr_Stator_OL_Lumped",
                      "Fr_Rotor_OC_Lumped", "Fr_Stator_OC_Lumped"]:
            series, meta = net_force_series(mc, graph)
            if series:
                forces[graph] = {"series_N": series, "stats": stats(series),
                                 "meta": meta}
                print("  %s: %s" % (graph, stats(series)))
            else:
                forces[graph] = None
                print("  [如实] %s 无数据" % graph)
        results["axial_forces"] = forces

        # ---- 校核 a: 定转子合力反号 ----
        checks = {}
        for case in ("OL", "OC"):
            fr = forces.get("Fr_Rotor_%s_Lumped" % case)
            fs = forces.get("Fr_Stator_%s_Lumped" % case)
            if fr and fs:
                mr, ms = fr["stats"]["mean"], fs["stats"]["mean"]
                checks["action_reaction_%s" % case] = {
                    "rotor_mean_N": mr, "stator_mean_N": ms,
                    "imbalance_pct": abs(mr + ms) / max(abs(mr), 1e-9) * 100}
        # ---- 校核 b: Σ(Ft×r) vs 转矩 ----
        tq_ft = torque_from_ft(mc, "Ft_Rotor_OL_Lumped")
        _, tq_graph = None, None
        txs, tys = read_2d(mc, 17)          # id17 = 总转矩 (第4轮已辨认)
        tq_graph = stats(tys)["mean"] if tys else None
        checks["torque_crosscheck"] = {"sum_Ft_x_r_Nm_t0": tq_ft,
                                       "torque_graph_mean_Nm": tq_graph}
        # ---- 校核 c: 解析 F ≈ A/(2μ0)·mean(B²), B 取气隙磁密图 ----
        bxs, bys = read_2d(mc, "FluxDensityAirgap")
        if bys:
            b2 = sum(b * b for b in bys) / len(bys)
            d_out = float(results["Stator_Lam_Dia_mm"]) * 1e-3
            d_in = float(results["Stator_Bore_mm"]) * 1e-3
            area = math.pi / 4.0 * (d_out ** 2 - d_in ** 2)
            checks["analytic"] = {"mean_B2_T2": b2, "area_m2": area,
                                  "F_est_N": area / (2 * MU0) * b2}
        results["checks"] = checks
        print("校核: %s" % json.dumps(checks, ensure_ascii=False, indent=1))

        # ---- CSV 波形 ----
        csv_path = os.path.join(OUT_DIR, "axial_force_%s.csv" % ts)
        with open(csv_path, "w", encoding="utf-8") as f:
            f.write("tstep,Fr_Rotor_OL_N,Fr_Stator_OL_N,"
                    "Fr_Rotor_OC_N,Fr_Stator_OC_N\n")
            nmax = max(len(v["series_N"]) if v else 0
                       for v in forces.values())
            for i in range(nmax):
                row = [str(i)]
                for g in ["Fr_Rotor_OL_Lumped", "Fr_Stator_OL_Lumped",
                          "Fr_Rotor_OC_Lumped", "Fr_Stator_OC_Lumped"]:
                    v = forces.get(g)
                    row.append("%.4f" % v["series_N"][i]
                               if v and i < len(v["series_N"]) else "")
                f.write(",".join(row) + "\n")
        results["csv"] = csv_path

        res_path = os.path.join(OUT_DIR, "axialforce_final_%s.json" % ts)
        with open(res_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print("RESULTS: %s" % res_path)
        print("CSV: %s" % csv_path)

        # ---- 结论摘要 ----
        for case, label in (("OL", "负载(RMS 21A)"), ("OC", "空载")):
            v = forces.get("Fr_Rotor_%s_Lumped" % case)
            if v:
                s = v["stats"]
                print("[结论] %s 转子净轴向力: 均值 %.1f N, 纹波峰峰 %.2f N"
                      % (label, s["mean"], s["pk2pk"]))
        return 0
    finally:
        if quit_after:
            try:
                mc.quit()
            except Exception:
                pass
        else:
            print("[提示] Motor-CAD 保持前台打开供检查。")


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
