# -*- coding: utf-8 -*-
"""
轴向磁拉力求解核心 — 从 axial_force_final.py / axial_compare.py 提取,
封装为可被 GUI 调用的 Solver 类。支持单点仿真与气隙扫描两种模式。
所有日志通过 log_cb 回调输出到 GUI。
"""
import json
import math
import os
import time

MU0 = 4e-7 * math.pi
MAX_TSTEPS = 64
MAX_NODES = 40

# 全局 keepalive: keep_open=True 时保存 MotorCAD 对象引用, 防止 solver 被 GC
# 时连带销毁 MotorCAD → 关闭 Motor-CAD 进程 (GUI 模式下必须保持)
_KEEPALIVE_MC = None


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):
    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 = {"sections": {}}
    for tstep in range(MAX_TSTEPS):
        total, got = 0.0, False
        for sec in (1, 2):
            xs, ys = _read_nodes(mc, graph, sec, tstep)
            if not ys:
                continue
            got = True
            nu = len(ys) - 1 if (len(xs) > 1 and
                                 abs(xs[-1] - xs[0] - 360.0) < 1e-6) else len(ys)
            total += sum(ys[:nu])
            if tstep == 0:
                meta["sections"][sec] = {"n_points": len(ys), "n_unique": nu}
        if not got:
            break
        series.append(total)
    return series, meta


def _torque_from_ft(mc, sec_radii_mm):
    tq = 0.0
    for sec, r_mm in zip((1, 2), sec_radii_mm):
        xs, ys = _read_nodes(mc, "Ft_Rotor_OL_Lumped", sec, 0)
        if not ys:
            return None
        nu = len(ys) - 1 if (len(xs) > 1 and
                             abs(xs[-1] - xs[0] - 360.0) < 1e-6) else len(ys)
        tq += sum(ys[:nu]) * (r_mm * 1e-3)
    return tq


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


class Solver:
    """轴向磁拉力求解器。参数通过构造函数传入, run_single() / run_sweep() 执行。"""

    def __init__(self, mot_path, out_dir, rms_current=21.0, speed_rpm=5000.0,
                 airgap_mm=1.0, magnet_temp_c=100.0, sec_radii_mm=None,
                 motorcad_exe=None, keep_open=True, log_cb=None):
        self.mot_path = mot_path
        self.out_dir = out_dir
        self.rms_current = rms_current
        self.speed_rpm = speed_rpm
        self.airgap_mm = airgap_mm
        self.magnet_temp_c = magnet_temp_c
        self.sec_radii_mm = sec_radii_mm or [28.25, 34.75]
        self.motorcad_exe = motorcad_exe
        self.keep_open = keep_open
        self.log = log_cb or (lambda msg: None)
        self.mc = None

    def _start_motorcad(self):
        from ansys.motorcad.core import MotorCAD, set_motorcad_exe
        if not os.environ.get("MOTORCAD_ACTIVEX") and self.motorcad_exe:
            if os.path.isfile(self.motorcad_exe):
                set_motorcad_exe(self.motorcad_exe)
                self.log("  未检测到 MOTORCAD_ACTIVEX, 已显式定位 exe")
        self.log("启动 Motor-CAD (前台) ...")
        self.mc = MotorCAD()
        try:
            self.mc.set_visible(True)
        except Exception as e:
            self.log("  [提示] set_visible 失败 (不影响计算): %s" % e)

    def _load_and_configure(self, tag):
        os.makedirs(self.out_dir, exist_ok=True)
        ts = time.strftime("%m%d_%H%M%S")
        self.mc.load_from_file(self.mot_path)
        out_mot = os.path.join(self.out_dir, "MARS_SSSR_%s_%s.mot" % (tag, ts))
        self.mc.save_to_file(out_mot)

        # 设置工况参数
        self.mc.set_variable("RMSCurrent", self.rms_current)
        self.mc.set_variable("ShaftSpeed", self.speed_rpm)
        self.mc.set_variable("Airgap", self.airgap_mm)
        self.mc.set_variable("Magnet_Temperature", self.magnet_temp_c)
        # 开力计算开关
        self.mc.set_variable("ElectromagneticForcesCalc_Load", True)
        self.mc.set_variable("ElectromagneticForcesCalc_OC", True)

        # 回读确认
        params = {}
        for var in ["RMSCurrent", "ShaftSpeed", "Airgap", "Magnet_Temperature",
                    "Stator_Lam_Dia", "Stator_Bore"]:
            try:
                params[var] = self.mc.get_variable(var)
            except Exception:
                params[var] = None
        self.log("  工况: RMS=%.1fA, %drpm, 气隙=%.2fmm, 磁钢%.0f°C"
                 % (params["RMSCurrent"], params["ShaftSpeed"],
                    params["Airgap"], params["Magnet_Temperature"]))
        return out_mot, params, ts

    def _solve_once(self, tag="axialF"):
        """执行一次求解, 返回 (results_dict, csv_path, json_path)"""
        out_mot, params, ts = self._load_and_configure(tag)
        self.log("求解中 (OC+OL 力同算) ...")
        t0 = time.time()
        self.mc.do_magnetic_calculation()
        dt = time.time() - t0
        self.log("  耗时 %.1f s" % dt)

        results = {"when": ts, "source_mot": os.path.basename(self.mot_path),
                   "work_mot": out_mot, "solve_seconds": dt,
                   "params": {"rms_current": params["RMSCurrent"],
                              "speed_rpm": params["ShaftSpeed"],
                              "airgap_mm": params["Airgap"],
                              "magnet_temp_c": params["Magnet_Temperature"]},
                   "convention_note": "AFM 2.5D 中 Fr(法向)=轴向力; 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(self.mc, graph)
            if series:
                forces[graph] = {"series_N": series, "stats": _stats(series),
                                 "meta": meta}
                s = _stats(series)
                self.log("  %s: 均值 %.1f N, 纹波 %.2f N"
                         % (graph, s["mean"], s["pk2pk"]))
            else:
                forces[graph] = None
                self.log("  [警告] %s 无数据" % graph)
        results["axial_forces"] = forces

        # 三判据
        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}
        tq_ft = _torque_from_ft(self.mc, self.sec_radii_mm)
        tys = _read_2d(self.mc, 17)
        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}
        bys = _read_2d(self.mc, "FluxDensityAirgap")
        if bys and params.get("Stator_Lam_Dia") and params.get("Stator_Bore"):
            b2 = sum(b * b for b in bys) / len(bys)
            d_out = float(params["Stator_Lam_Dia"]) * 1e-3
            d_in = float(params["Stator_Bore"]) * 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
        self.log("校核: %s" % json.dumps(checks, ensure_ascii=False))

        # CSV
        csv_path = os.path.join(self.out_dir, "axial_force_%s_%s.csv" % (tag, 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

        # JSON
        json_path = os.path.join(self.out_dir, "axialforce_%s_%s.json" % (tag, ts))
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)

        # 结论摘要
        for case, label in (("OL", "负载(RMS %.0fA)" % self.rms_current),
                            ("OC", "空载")):
            v = forces.get("Fr_Rotor_%s_Lumped" % case)
            if v:
                s = v["stats"]
                self.log("[结论] %s 转子净轴向力: 均值 %.1f N, 纹波峰峰 %.2f N"
                         % (label, s["mean"], s["pk2pk"]))
        return results, csv_path, json_path

    def run_single(self):
        """单点仿真。返回 results dict。"""
        self._start_motorcad()
        try:
            results, csv_path, json_path = self._solve_once("single")
            self.log("结果文件: %s" % json_path)
            return results
        finally:
            if not self.keep_open:
                try:
                    self.mc.quit()
                except Exception:
                    pass
            else:
                global _KEEPALIVE_MC
                _KEEPALIVE_MC = self.mc  # 保持引用, 防止 GC 关闭 Motor-CAD
                self.log("[提示] Motor-CAD 保持前台打开供检查。")

    def run_sweep(self, gaps):
        """气隙扫描。gaps: 气隙列表(mm)。返回汇总 dict。"""
        self._start_motorcad()
        try:
            os.makedirs(self.out_dir, exist_ok=True)
            ts = time.strftime("%m%d_%H%M%S")
            self.mc.load_from_file(self.mot_path)
            out_mot = os.path.join(self.out_dir, "MARS_SSSR_sweep_%s.mot" % ts)
            self.mc.save_to_file(out_mot)

            self.mc.set_variable("RMSCurrent", self.rms_current)
            self.mc.set_variable("ShaftSpeed", self.speed_rpm)
            self.mc.set_variable("Magnet_Temperature", self.magnet_temp_c)
            self.mc.set_variable("ElectromagneticForcesCalc_Load", True)
            self.mc.set_variable("ElectromagneticForcesCalc_OC", True)

            sweep = {"when": ts, "magnet_temp_C": self.magnet_temp_c,
                     "rms_current_A": self.rms_current, "work_mot": out_mot,
                     "cases": []}
            for g in gaps:
                self.mc.set_variable("Airgap", g)
                back = self.mc.get_variable("Airgap")
                self.log("== 气隙 %.2f mm (回读 %s), 求解 ..." % (g, back))
                t0 = time.time()
                self.mc.do_magnetic_calculation()
                dt = time.time() - t0
                case = {"airgap_mm": back, "solve_seconds": dt}
                for graph, key in [("Fr_Rotor_OC_Lumped", "F_OC"),
                                   ("Fr_Rotor_OL_Lumped", "F_OL")]:
                    s, _ = _net_force_series(self.mc, graph)
                    case[key] = _stats(s)
                bys = _read_2d(self.mc, "FluxDensityAirgap")
                if bys:
                    case["B2_mean_T2"] = sum(b * b for b in bys) / len(bys)
                self.log("   F_OC=%.1f N, F_OL=%.1f N (耗时 %.0fs)"
                         % (case["F_OC"]["mean"], case["F_OL"]["mean"], dt))
                sweep["cases"].append(case)

            # 磁负刚度 (有限差分, OC 口径)
            cs = sweep["cases"]
            if len(cs) >= 2:
                f = [c["F_OC"]["mean"] for c in cs]
                g = [c["airgap_mm"] for c in cs]
                kneg = {}
                for i in range(len(cs) - 1):
                    kneg["seg_%.2f_%.2f" % (g[i], g[i + 1])] = \
                        -(f[i + 1] - f[i]) / (g[i + 1] - g[i])
                if len(cs) >= 3:
                    kneg["central"] = -(f[-1] - f[0]) / (g[-1] - g[0])
                sweep["kneg_N_per_mm"] = kneg
                self.log("磁负刚度: %s N/mm" % json.dumps(kneg, ensure_ascii=False))

            json_path = os.path.join(self.out_dir, "sweep_results_%s.json" % ts)
            with open(json_path, "w", encoding="utf-8") as f:
                json.dump(sweep, f, ensure_ascii=False, indent=2)
            self.log("扫描结果: %s" % json_path)
            return sweep
        finally:
            if not self.keep_open:
                try:
                    self.mc.quit()
                except Exception:
                    pass
            else:
                global _KEEPALIVE_MC
                _KEEPALIVE_MC = self.mc  # 保持引用, 防止 GC 关闭 Motor-CAD
                self.log("[提示] Motor-CAD 保持前台打开供检查。")
