# -*- coding: utf-8 -*-
"""
第5轮: 力开关打开后筛 2D/3D 图名
=================================
修正第3轮漏洞: 图名可能在 ElectromagneticForcesCalc_* 打开后才注册, 先开
开关再筛。同时筛 3D 力图 (get_magnetic_3d_graph_point, 空间x时间), 并用
已知图 (转矩/气隙磁密) 反推 Graph Viewer 命名惯例。

用法: python axial_probe4.py [--quit]
"""
import json
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")

# 已知图命名惯例侦察 (id17=转矩0.522, id0-2=相电流, FluxDensityAirgap 已知)
KNOWN_PROBE = ["Torque", "Torque OL", "Torque (OL)", "TorqueOL", "Torque_OL",
               "Torque vs Angle", "TorqueVsAngle", "Cogging Torque",
               "Cogging Torque OC", "CoggingTorque", "Phase Current",
               "Phase Current OL", "Current", "CurrentOL", "Back EMF",
               "Phase EMF", "FluxDensityAirgap", "Airgap Flux Density"]

BASES = ["Axial Force", "Axial_Force", "AxialForce", "Fz", "Fa"]
ENTS = ["", " Rotor", " Stator", "_Rotor", "_Stator", "Rotor", "Stator"]
CASES = ["", " OL", " OC", "_OL", "_OC", " (OL)", " (OC)", " Load",
         " Open Circuit", "_OL_Lumped", "_OC_Lumped", "_Lumped"]

D3_EXTRA = ["Ft_Rotor_OL_Lumped", "Fr_Rotor_OL_Lumped",
            "Ft_Stator_OL_Lumped", "Fr_Stator_OL_Lumped",
            "OL_Axial_Force_Rotor", "OC_Axial_Force_Rotor",
            "OL_Axial_Force_Stator", "OC_Axial_Force_Stator",
            "Load_Axial_Force_Rotor", "Load_Axial_Force_Stator"]


def build_force_names():
    out = []
    for b in BASES:
        for e in ENTS:
            for c in CASES:
                n = b + e + c
                if n not in out:
                    out.append(n)
    for n in D3_EXTRA:
        if n not in out:
            out.append(n)
    return out


def classify2d(mc, g):
    try:
        mc.get_magnetic_graph_point(g, 0)
        return "OK"
    except Exception as e:
        m = str(e)
        if "does not exist" in m:
            return "ABSENT"
        if "No points exist" in m:
            return "EXISTS"
        return "ODD:" + m


def classify3d(mc, g):
    try:
        mc.get_magnetic_3d_graph_point(g, 1, 0, 0)
        return "OK"
    except Exception as e:
        m = str(e)
        if "does not exist" in m:
            return "ABSENT"
        if "No points exist" in m or "no points" in m.lower():
            return "EXISTS"
        return "ODD:" + m


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 read2d(mc, g, maxpts=128):
    xs, ys = [], []
    for i in range(maxpts):
        try:
            x, y = mc.get_magnetic_graph_point(g, i)
        except Exception:
            break
        xs.append(x)
        ys.append(y)
    return xs, ys


def read3d(mc, g, section, maxpts=256, tstep=0):
    xs, ys = [], []
    for i in range(maxpts):
        try:
            x, y = mc.get_magnetic_3d_graph_point(g, section, i, tstep)
        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
    print("启动 Motor-CAD (前台) ...")
    mc = MotorCAD()
    results = {"when": ts, "probe_round": 5}
    try:
        mc.load_from_file(MOT_SRC)
        for var in ["ElectromagneticForcesCalc_Load",
                    "ElectromagneticForcesCalc_OC"]:
            mc.set_variable(var, True)
        print("力开关已开, 开始筛名 (无求解) ...")

        known = {n: classify2d(mc, n) for n in KNOWN_PROBE}
        results["known_probe"] = known
        print("已知图命名侦察: %s" % json.dumps(
            {k: v for k, v in known.items() if v != "ABSENT"},
            ensure_ascii=False))

        force_names = build_force_names()
        hits2d = {}
        for n in force_names:
            k = classify2d(mc, n)
            if k != "ABSENT":
                hits2d[n] = k
        results["force_2d_hits"] = hits2d
        print("2D 力图命中: %s" % json.dumps(hits2d, ensure_ascii=False))

        hits3d = {}
        for n in force_names:
            k = classify3d(mc, n)
            if k != "ABSENT":
                hits3d[n] = k
        results["force_3d_hits"] = hits3d
        print("3D 力图命中: %s" % json.dumps(hits3d, ensure_ascii=False))

        promising2d = [n for n, k in hits2d.items()]
        promising3d = [n for n, k in hits3d.items()]
        if promising2d or promising3d:
            print("求解一次后读波形 ...")
            t0 = time.time()
            mc.do_magnetic_calculation()
            print("  耗时 %.1f s" % (time.time() - t0))
            waves = {}
            for n in promising2d:
                xs, ys = read2d(mc, n)
                if ys:
                    waves["2D:" + n] = {"stats": stats(ys), "x": xs, "y": ys}
                    print("  [2D] %s: %s" % (n, stats(ys)))
            for n in promising3d:
                for sec in (1, 2):
                    xs, ys = read3d(mc, n, sec)
                    if ys:
                        key = "3D:%s:sec%d" % (n, sec)
                        waves[key] = {"stats": stats(ys), "x": xs, "y": ys}
                        print("  [3D] %s sec%d: %s" % (n, sec, stats(ys)))
            results["waveforms"] = waves
        else:
            print("[如实] 力开关打开后仍无任何命中")

        res_path = os.path.join(OUT_DIR, "probe4_results_%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)
        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:]))