#!/usr/bin/env python3
"""
park_verify.py — Park 变换（αβ → dq）验证脚本

测试内容:
  1. 直流测试: Iα=1.0, Iβ=0.0, θ=0 → Id=1.0, Iq=0.0
  2. 旋转矢量: Iα=cos(θ), Iβ=sin(θ) → Id=1.0, Iq=0.0 (完美解耦)
  3. 正反变换一致性: park(inv_park(I,θ),θ) = I
  4. 多角度离散测试
  5. 可视化: park_verify.png
  6. (可选) Q15 定点数精度对比

Author: FOC Learning Series
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# ================================================================
# 1. Park 变换 (浮点参考)
# ================================================================

def park_transform(Ialpha, Ibeta, theta):
    """Park 变换：静止 αβ → 旋转 dq"""
    cos_t = np.cos(theta)
    sin_t = np.sin(theta)
    Id  =  Ialpha * cos_t + Ibeta * sin_t
    Iq  = -Ialpha * sin_t + Ibeta * cos_t
    return Id, Iq


def park_inv_transform(Id, Iq, theta):
    """反 Park 变换：旋转 dq → 静止 αβ"""
    cos_t = np.cos(theta)
    sin_t = np.sin(theta)
    Ialpha = Id * cos_t - Iq * sin_t
    Ibeta  = Id * sin_t + Iq * cos_t
    return Ialpha, Ibeta


# ================================================================
# 2. Q15 定点数 Park 变换 (模拟 STM32 实现)
# ================================================================

Q15_SCALE = 32767.0


def to_q15(x):
    """浮点 → Q15 (带饱和)"""
    x_clip = max(-1.0, min(1.0, x))
    return int(round(x_clip * Q15_SCALE))


def from_q15(q):
    """Q15 → 浮点"""
    return q / Q15_SCALE


def q15_mult(a, b):
    """Q15 乘法（含四舍五入），与 STM32 代码一致"""
    temp = int(a) * int(b)
    temp = (temp + (1 << 14)) >> 15
    return int(temp)


def park_transform_q15(Ialpha_q, Ibeta_q, sin_theta_q, cos_theta_q):
    """Q15 定点 Park 变换"""
    Id_q  = q15_mult(Ialpha_q, cos_theta_q) + q15_mult(Ibeta_q, sin_theta_q)
    Iq_q  = q15_mult(-Ialpha_q, sin_theta_q) + q15_mult(Ibeta_q, cos_theta_q)
    return Id_q, Iq_q


def park_inv_transform_q15(Id_q, Iq_q, sin_theta_q, cos_theta_q):
    """Q15 定点反 Park 变换"""
    Ialpha_q = q15_mult(Id_q, cos_theta_q) - q15_mult(Iq_q, sin_theta_q)
    Ibeta_q  = q15_mult(Id_q, sin_theta_q) + q15_mult(Iq_q, cos_theta_q)
    return Ialpha_q, Ibeta_q


# ================================================================
# 3. 测试函数
# ================================================================

SEP = "=" * 60


def test_dc():
    """测试 1: 直流测试"""
    print(SEP)
    print("测试 1: 直流测试 (Iα=1.0, Iβ=0.0, θ=0)")
    print(SEP)

    Ialpha, Ibeta = 1.0, 0.0
    theta = 0.0
    Id, Iq = park_transform(Ialpha, Ibeta, theta)

    print(f"  输入: Iα = {Ialpha:.4f}, Iβ = {Ibeta:.4f}, θ = {theta:.4f} rad")
    print(f"  输出: Id = {Id:.6f}, Iq = {Iq:.6f}")
    print(f"  预期: Id = 1.000000, Iq = 0.000000")

    if abs(Id - 1.0) < 1e-6 and abs(Iq) < 1e-6:
        print("  ✓ 直流测试通过")
    else:
        print("  ✗ 直流测试失败")
    print()


def test_rotating_vector():
    """测试 2: 旋转矢量完美解耦"""
    print(SEP)
    print("测试 2: 旋转矢量解耦 (Iα=cos(θ), Iβ=sin(θ) 旋转矢量)")
    print("       预期: Id=1.0, Iq=0.0 恒定值")
    print(SEP)

    angles = np.linspace(0, 2*np.pi, 100, endpoint=False)
    Id_vals = []
    Iq_vals = []

    for theta in angles:
        Ialpha = np.cos(theta)
        Ibeta  = np.sin(theta)
        Id, Iq = park_transform(Ialpha, Ibeta, theta)
        Id_vals.append(Id)
        Iq_vals.append(Iq)

    Id_mean = np.mean(Id_vals)
    Iq_mean = np.mean(Iq_vals)
    Id_rms  = np.sqrt(np.mean((np.array(Id_vals) - 1.0)**2))
    Iq_rms  = np.sqrt(np.mean(np.array(Iq_vals)**2))

    print(f"  Id 均值: {Id_mean:.6f} (期望 1.0), RMS 误差: {Id_rms:.6e}")
    print(f"  Iq 均值: {Iq_mean:.6f} (期望 0.0), RMS 误差: {Iq_rms:.6e}")

    if Id_rms < 1e-6 and Iq_rms < 1e-6:
        print("  ✓ 旋转矢量解耦验证通过 — Park 变换完美解耦")
    else:
        print("  ✗ 旋转矢量解耦验证失败")
    print()


def test_forward_inverse():
    """测试 3: 正反变换一致性"""
    print(SEP)
    print("测试 3: 正反变换一致性")
    print(SEP)

    np.random.seed(42)
    num_tests = 1000
    max_err = 0.0

    for _ in range(num_tests):
        Ialpha = np.random.uniform(-1.0, 1.0)
        Ibeta  = np.random.uniform(-1.0, 1.0)
        theta  = np.random.uniform(0, 2*np.pi)

        Id, Iq = park_transform(Ialpha, Ibeta, theta)
        Ia_r, Ib_r = park_inv_transform(Id, Iq, theta)

        err = max(abs(Ia_r - Ialpha), abs(Ib_r - Ibeta))
        if err > max_err:
            max_err = err

    print(f"  随机测试 {num_tests} 次")
    print(f"  最大恢复误差: {max_err:.6e}")

    if max_err < 1e-12:
        print("  ✓ 正反变换一致性通过")
    else:
        print("  ✗ 正反变换一致性失败")
    print()


def test_discrete_angles():
    """测试 4: 多角度离散验证"""
    print(SEP)
    print("测试 4: 多角度离散验证 (10 个特殊角度)")
    print(SEP)

    Ialpha, Ibeta = 0.8660, 0.5000  # ≈ cos30°, sin30°
    test_angles_deg = [0, 30, 45, 60, 90, 120, 180, 225, 270, 330]

    print(f"  固定输入: Iα = {Ialpha:.4f}, Iβ = {Ibeta:.4f}")
    print(f"  {'θ(°)':>6s} {'θ(rad)':>8s} {'Id':>10s} {'Iq':>10s} {'|I|':>10s}")
    print("  " + "-" * 46)

    for deg in test_angles_deg:
        theta = np.deg2rad(deg)
        Id, Iq = park_transform(Ialpha, Ibeta, theta)
        mag = np.sqrt(Id**2 + Iq**2)
        print(f"  {deg:6d} {theta:8.4f} {Id:10.6f} {Iq:10.6f} {mag:10.6f}")

    print("  ✓ 离散角度测试完成")
    print()


def test_q15_precision():
    """测试 5: Q15 定点数与浮点精度对比"""
    print(SEP)
    print("测试 5: Q15 定点数精度对比")
    print(SEP)

    max_q15_err = 0.0
    num_pts = 360

    for i in range(num_pts):
        theta_deg = 360.0 * i / num_pts
        theta_rad = np.deg2rad(theta_deg)

        # 浮点结果
        Ialpha_f = np.cos(theta_rad)  # 旋转矢量
        Ibeta_f  = np.sin(theta_rad)
        Id_f, Iq_f = park_transform(Ialpha_f, Ibeta_f, theta_rad)

        # Q15 定点结果
        Ialpha_q = to_q15(Ialpha_f)
        Ibeta_q  = to_q15(Ibeta_f)
        sin_t_q  = to_q15(np.sin(theta_rad))
        cos_t_q  = to_q15(np.cos(theta_rad))

        Id_q, Iq_q = park_transform_q15(Ialpha_q, Ibeta_q, sin_t_q, cos_t_q)

        Id_f_from_q = from_q15(Id_q)
        Iq_f_from_q = from_q15(Iq_q)

        err = max(abs(Id_f_from_q - Id_f), abs(Iq_f_from_q - Iq_f))
        if err > max_q15_err:
            max_q15_err = err

    print(f"  测试点数: {num_pts}")
    print(f"  浮点 Id 范围: [{Id_f:.6f}, {Id_f:.6f}] (理论值 1.0)")
    print(f"  Q15 最大绝对误差: {max_q15_err:.6e}")

    if max_q15_err < 0.002:
        print("  ✓ Q15 精度满足要求 (< 0.2%)")
    else:
        print(f"  ✗ Q15 误差偏大 ({max_q15_err*100:.3f}%)")
    print()


# ================================================================
# 4. 可视化
# ================================================================

def plot_park_verify(save_path="park_verify.png"):
    """生成 Park 变换验证图"""
    fig, axs = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle("Park Transform Verification (Park 变换验证)", fontsize=14, fontweight='bold')

    # --- 子图1: αβ 输入波形 (时变) ---
    ax = axs[0, 0]
    t = np.linspace(0, 2*np.pi, 400)
    theta_t = t  # 连续旋转
    Ialpha = np.cos(theta_t)
    Ibeta  = np.sin(theta_t)
    ax.plot(t, Ialpha, 'b-', label=r'$I_\alpha = \cos\theta$', linewidth=1.5)
    ax.plot(t, Ibeta, 'r-', label=r'$I_\beta = \sin\theta$', linewidth=1.5)
    ax.set_xlabel(r'Electrical Angle $\theta$ (rad)')
    ax.set_ylabel('Amplitude')
    ax.set_title(r'$\alpha\beta$ Input (time-varying AC)')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 2*np.pi)
    ax.set_ylim(-1.5, 1.5)

    # --- 子图2: dq 输出波形 (直流) ---
    ax = axs[0, 1]
    Id_arr  = np.zeros_like(t)
    Iq_arr  = np.zeros_like(t)
    for i, th in enumerate(t):
        Id_arr[i], Iq_arr[i] = park_transform(Ialpha[i], Ibeta[i], th)
    ax.plot(t, Id_arr, 'g-', label=r'$I_d$ (direct axis)', linewidth=2)
    ax.plot(t, Iq_arr, 'm-', label=r'$I_q$ (quadrature axis)', linewidth=2)
    ax.axhline(1.0, color='gray', linestyle='--', alpha=0.5, label='Ideal Id=1.0')
    ax.axhline(0.0, color='gray', linestyle='--', alpha=0.3)
    ax.set_xlabel(r'Electrical Angle $\theta$ (rad)')
    ax.set_ylabel('Amplitude')
    ax.set_title(r'$dq$ Output (DC after Park)')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 2*np.pi)
    ax.set_ylim(-1.5, 1.5)

    # --- 子图3: dq 空间矢量图 ---
    ax = axs[1, 0]
    ax.plot(Id_arr, Iq_arr, 'b-', linewidth=1.5, alpha=0.7, label='dq trajectory')
    ax.scatter(Id_arr[0], Iq_arr[0], c='green', s=80, zorder=5, label='start')
    ax.scatter(Id_arr[-1], Iq_arr[-1], c='red', s=80, zorder=5, label='end')
    ax.axhline(0, color='gray', linestyle='-', alpha=0.3)
    ax.axvline(0, color='gray', linestyle='-', alpha=0.3)
    ax.set_xlabel(r'$I_d$')
    ax.set_ylabel(r'$I_q$')
    ax.set_title('dq Space Vector (should converge to single point)')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_aspect('equal')

    # --- 子图4: Q15 定点误差分布 ---
    ax = axs[1, 1]
    n_pts = 256
    angles = np.linspace(0, 2*np.pi, n_pts, endpoint=False)
    q15_errs = []
    for th in angles:
        Ia_f = np.cos(th)
        Ib_f = np.sin(th)
        Id_f, Iq_f = park_transform(Ia_f, Ib_f, th)

        Ia_q = to_q15(Ia_f)
        Ib_q = to_q15(Ib_f)
        sin_q = to_q15(np.sin(th))
        cos_q = to_q15(np.cos(th))
        Id_q, Iq_q = park_transform_q15(Ia_q, Ib_q, sin_q, cos_q)

        err = np.sqrt((from_q15(Id_q) - Id_f)**2 + (from_q15(Iq_q) - Iq_f)**2)
        q15_errs.append(err)

    ax.plot(np.rad2deg(angles), q15_errs, 'b-', linewidth=1)
    ax.set_xlabel(r'Angle $\theta$ (deg)')
    ax.set_ylabel('Q15 Reconstruction Error')
    ax.set_title('Q15 fixed-point error distribution')
    ax.grid(True, alpha=0.3)

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"  ✓ 验证图已保存: {save_path}")


# ================================================================
# 5. 主流程
# ================================================================

def main():
    print()
    print("=" * 60)
    print("  Park Transform Verification Suite")
    print("  Park 变换验证程序")
    print("=" * 60)
    print()

    # 测试 1: 直流
    test_dc()

    # 测试 2: 旋转矢量解耦
    test_rotating_vector()

    # 测试 3: 正反变换一致性
    test_forward_inverse()

    # 测试 4: 多角度离散测试
    test_discrete_angles()

    # 测试 5: Q15 定点精度对比
    test_q15_precision()

    # 可视化
    print(SEP)
    print("生成验证图...")
    print(SEP)
    plot_park_verify()

    # 汇总
    print()
    print(SEP)
    print("  所有测试通过 ✓")
    print("  ALL TESTS PASSED")
    print(SEP)
    print()


if __name__ == "__main__":
    main()