import numpy as np

class PhiBasisIterator:
    """
    逐步生成有理正交 (类 Muntz-Laguerre) 基函数 φ_p(s) / (或复对产生两列) 的迭代器。

    给定:
        s: 标量或数组 (复频率采样点)  s = j*2*pi*f
        poles: 极点序列 (允许包含复共轭对). 约定:
            - 纯实极点按自身顺序出现: a = -alpha (alpha>0)
            - 复极点以正虚部在前, 紧跟其共轭: a = -σ + jω,  a* = -σ - jω

    迭代:
        每次 __next__ 返回一个 list:
            - 实极点 -> [φ_p(s)]
            - 复成对 -> [φ_p^(1)(s), φ_p^(2)(s)]
        使用时可把返回列表中函数 append 到总基函数数组中。

    说明:
        这里实现的公式与用户原始代码意图相同, 但修正了:
            1) 原代码 sqrt(2*Re(a)) 在 Re(a)<0 (典型稳定极点) 下产生 NaN, 改为 alpha=-Re(a) >0.
            2) 原来在 __next__ 中使用 yield (非法); 改为标准迭代协议返回 list。
            3) 复极点的乘积与后续迭代用的累计乘积 product 分开维护。
        若需严格匹配论文中“多变量正交化”可再加数值正交 (QR) 步骤。

    注意:
        本实现假设 poles 已排序; 若不确定, 可预处理把正虚部极点放在其共轭之前。
    """
    def __init__(self, s, poles):
        self.s = np.asarray(s, dtype=complex)  # (Nf,) 或标量
        self.poles = list(poles)
        self.k = 0                # 当前处理到第 k 个极点(或复对的首元素)
        self.N = len(self.poles)
        # 累计乘积 Π_{已完成的块} ( (s - a_i*)/(s + a_i) )  (复对时使用 (s - a)(s - a*) / (s + a)(s + a*))
        self.product = np.ones_like(self.s, dtype=complex)

    def __iter__(self):
        return self

    def __next__(self):
        if self.k >= self.N:
            raise StopIteration

        a = self.poles[self.k]

        # 复对: 需要检查正虚部并且下一个是共轭
        if np.iscomplex(a) and np.imag(a) > 0:
            if self.k + 1 >= self.N or not np.isclose(self.poles[self.k + 1], np.conj(a)):
                raise ValueError(f"极点序列中复极点 {a} 未紧跟其共轭, 请排序.")
            a_conj = self.poles[self.k + 1]
            sigma = -np.real(a)          # >0
            alpha = sigma
            # (s + a)(s + a*) = (s + a)(s + conj(a))
            denom = (self.s + a) * (self.s + a_conj)
            # |a| 用于生成两组分子 (参考原代码的构造方式)
            r = np.abs(a)
            scale = np.sqrt(2 * alpha)

            numerator1 = scale * (self.s - r)
            numerator2 = scale * (self.s + r)

            phi1 = numerator1 / denom * self.product
            phi2 = numerator2 / denom * self.product

            # 更新累计乘积用于后续阶次：
            # 对复对的“整体”传统一致写成 ((s - a)(s - a*) / (s + a)(s + a*))
            self.product = self.product * ((self.s - a) * (self.s - a_conj)) / denom

            self.k += 2
            return [phi1, phi2]

        else:
            # 实极点 (允许传入 float 或实部为负的 complex(实数))
            a_real = np.real(a)
            # 稳定极点: a_real < 0, 令 alpha = -a_real > 0
            if a_real >= 0:
                raise ValueError(f"实极点需要在左半平面 (负实部), 得到 {a_real}")
            alpha = -a_real
            scale = np.sqrt(2 * alpha)
            # Laguerre 风格: (s - alpha)/(s + alpha)
            numerator = scale * (self.s - alpha)
            denom = (self.s + alpha)
            phi = numerator / denom * self.product
            # 更新累计乘积
            self.product = self.product * (self.s - alpha) / (self.s + alpha)
            self.k += 1
            return [phi]

# ------------------ 辅助函数: 批量生成全部基函数 ------------------
def generate_phi_basis(s, poles):
    """
    返回一个 list, 其中包含依次生成的所有基函数 (按 real → 1 函数, complex pair → 2 函数).
    """
    basis = []
    it = PhiBasisIterator(s, poles)
    for block in it:
        basis.extend(block)
    return basis

# ------------------ 示例与自测 ------------------
if __name__ == "__main__":
    # 示例极点: 1 个实极点 (-1), 一个复对 (-2+0.5j, -2-0.5j), 再一个复对 (-3+1.2j, -3-1.2j)
    poles = [-0.005+500j, -0.005-500j, -0.012+1200j, -0.012-1200j]
    f = np.linspace(1e9, 1e11, 200)            # 频率 (Hz)
    s = 1j * 2 * np.pi * f

    basis_funcs = generate_phi_basis(s, poles)
    print(f"生成基函数数量 = {len(basis_funcs)}")  # 实极点 1 →1, 两个复对 →2+2 共 5 个
    # 简单数值检查 (避免 NaN/Inf)
    for i, bf in enumerate(basis_funcs):
        if not np.all(np.isfinite(bf)):
            print(f"第 {i} 个基函数存在非有限值")
        else:
            # print(f"phi[{i}] max|.|={np.max(np.abs(bf)):.3e}, min|.|={np.min(np.abs(bf)):.3e}")
            print(f"phi[{i}] = {bf}")

    # 验证是否为正交基
    for i, bf in enumerate(basis_funcs):
        print(f"phi[{i}]*phi[{i}] = {np.vdot(bf, bf)}")  # 自身内积
        for j in range(i+1, len(basis_funcs)):
            bf2 = basis_funcs[j]
            ip = np.vdot(bf, bf2)
            print(f"phi[{i}] 与 phi[{j}] 内积 = {ip}")
            if np.abs(ip) > 1e-6:
                print(f"phi[{i}] 与 phi[{j}] 非正交, 内积={ip}")