import numpy as np

# ------------------------------------------------------------
# 按论文公式 (9)(10)(11) 生成 Muntz–Laguerre 正交有理基 (解析形式)：
#
# 给定稳定极点集合 {p_k} (Re(p_k)<0)。论文记法中使用 -a_k，其中 Re(a_k)>0。
# 对应关系：p_k = -a_k  ⇒  a_k = -p_k,  Re(a_k)= -Re(p_k) >0
#
# 连续内积意义下（沿 jω 轴积分）这些 φ_k 解析正交。离散频率采样后数值上
# 可能偏离，可再用加权 QR 做数值再正交（可选）。
#
# 公式在“稳定极点 p 表达”下的改写：
#   (原) 实极点:  φ_p(s) = sqrt(2 Re(a_p)) / (s + a_p) * Π (s - a_i^*)/(s + a_i)
#   变换 a_p = -p ⇒ Re(a_p)= -Re(p) = σ >0 且 (s + a_p) = (s - p)
#   且乘积 (s - a_i^*)/(s + a_i) = (s + p_i^*)/(s - p_i)
#   ⇒ φ_p(s) = sqrt(-2 Re(p)) / (s - p) * Π_{i<p} (s + p_i^*)/(s - p_i)
#
# 复对 (p, p*)，取 imag(p)>0 的 p 作为首：
#   (原) φ_p   = sqrt(2 Re(a_p)) (s - |a_p|)/[(s + a_p)(s + a_p^*)] * Π(...)
#   (原) φ_{p+1}= sqrt(2 Re(a_p)) (s + |a_p|)/[(s + a_p)(s + a_p^*)] * Π(...)
#   代入 a_p=-p：
#       Re(a_p)= -Re(p)=σ>0,  (s + a_p) = (s - p), (s + a_p^*)=(s - p^*)
#       |a_p| = |p|
#       乘积同上 ⇒ Π_{i<p} (s + p_i^*)/(s - p_i)
#
#   ⇒ φ_p(s)   = sqrt(-2 Re(p)) (s - |p|)/[(s - p)(s - p^*)] * Π_{i<p} (s + p_i^*)/(s - p_i)
#      φ_{p^*}(s)= sqrt(-2 Re(p)) (s + |p|)/[(s - p)(s - p^*)] * Π_{i<p} (s + p_i^*)/(s - p_i)
#
# product 递推维护：prod_k = Π_{i≤k} (s + p_i^*)/(s - p_i)
# 对复对要顺序乘两次（p 与 p*）。
# ------------------------------------------------------------


def generate_laguerre_basis(poles:list|np.ndarray, s: np.ndarray):
    basis = np.zeros((len(poles)+1,len(s)),dtype=complex)

    product = np.ones(len(s),dtype=complex)

    basis[0] = np.ones(len(s),dtype=complex)  # φ_0 = 1
    i = 0
    while i < len(poles):
        if np.real(poles[i]) >= 0:
            raise ValueError(f"极点必须在左半平面: {poles[i]}")

        # 复对首 (正虚部)
        if np.iscomplex(poles[i]) and np.imag(poles[i]) > 0:
            if i + 1 >= len(poles):
                raise ValueError("复极点缺少共轭")
            pc = poles[i + 1]
            if not np.isclose(pc, np.conj(poles[i])):
                raise ValueError("复极点未按 (p, p*) 顺序排列 (正虚部在前)")
            sigma = -np.real(poles[i])           # >0
            scale = np.sqrt(2 * sigma)
            r = np.abs(poles[i])
            denom = (s - poles[i]) * (s - pc)

            # 两个基函数
            phi_p  = scale * (s - r) / denom * product
            phi_pc = scale * (s + r) / denom * product

            # product 先乘 (s + p^*)/(s - p)，再乘 (s + p)/(s - p^*)
            product = product * (s + pc) / (s - poles[i])
            product = product * (s + poles[i])  / (s - pc)

            basis[i + 1] = phi_p
            basis[i + 2] = phi_pc
            i += 2
            continue

        # 复对次 (负虚部) —— 应该被首元素处理，出现表示顺序错误
        if np.iscomplex(poles[i]) and np.imag(poles[i]) < 0:
            raise ValueError("检测到负虚部复极点但其共轭尚未处理，请将正虚部成员放在前面。")

        # 实极点
        sigma = -np.real(poles[i])
        if sigma <= 0:
            raise ValueError("实极点实部应为负 (稳定)。")
        scale = np.sqrt(2 * sigma)
        phi = scale / (s - poles[i]) * product
        # 更新乘积
        product = product * (s + poles[i]) / (s - poles[i])
        i += 1
        basis[i + 1] = phi
    return basis





# class MuntzLaguerreIterator:
#     def __init__(self, s: np.ndarray, stable_poles: list | np.ndarray):
#         """
#         s: 复频率数组 (Nf,), s = j 2π f
#         stable_poles: 稳定极点列表 (Re<0). 复共轭对要求正虚部在前 (p, p*).
#         """
#         self.s = np.asarray(s, dtype=complex)
#         self.poles = list(stable_poles)
#         self.N = len(self.poles)
#         self.k = 0
#         # 初始化乘积 Π_{i<p} (s + p_i^*)/(s - p_i)
#         self.product = np.ones_like(self.s, dtype=complex)

#     def __iter__(self):
#         return self

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

#         p = self.poles[self.k]
#         if np.real(p) >= 0:
#             raise ValueError(f"极点必须在左半平面: {p}")

#         # 复对首 (正虚部)
#         if np.iscomplex(p) and np.imag(p) > 0:
#             if self.k + 1 >= self.N:
#                 raise ValueError("复极点缺少共轭")
#             pc = self.poles[self.k + 1]
#             if not np.isclose(pc, np.conj(p)):
#                 raise ValueError("复极点未按 (p, p*) 顺序排列 (正虚部在前)")
#             sigma = -np.real(p)           # >0
#             scale = np.sqrt(2 * sigma)
#             r = np.abs(p)
#             denom = (self.s - p) * (self.s - pc)

#             # 两个基函数
#             phi_p  = scale * (self.s - r) / denom * self.product
#             phi_pc = scale * (self.s + r) / denom * self.product

#             # product 先乘 (s + p^*)/(s - p)，再乘 (s + p)/(s - p^*)
#             self.product = self.product * (self.s + pc) / (self.s - p)
#             self.product = self.product * (self.s + p)  / (self.s - pc)

#             self.k += 2
#             return [phi_p, phi_pc]

#         # 复对次 (负虚部) —— 应该被首元素处理，出现表示顺序错误
#         if np.iscomplex(p) and np.imag(p) < 0:
#             raise ValueError("检测到负虚部复极点但其共轭尚未处理，请将正虚部成员放在前面。")

#         # 实极点
#         sigma = -np.real(p)
#         if sigma <= 0:
#             raise ValueError("实极点实部应为负 (稳定)。")
#         scale = np.sqrt(2 * sigma)
#         phi = scale / (self.s - p) * self.product
#         # 更新乘积
#         self.product = self.product * (self.s + p) / (self.s - p)

#         self.k += 1
#         return [phi]
    
# def generate_muntz_laguerre_basis(s: np.ndarray, init_poles: list | np.ndarray):
#     """
#     生成完整基函数列表: [φ_0=1, φ_1, φ_2, ...]
#     """
#     basis = [np.ones_like(s, dtype=complex)]
#     for block in MuntzLaguerreIterator(s, init_poles):
#         basis.extend(block)
#     return basis

# if __name__ == "__main__":
#     # 示例稳定极点 (复对正虚部在前)
#     stable_poles = [
#         -0.8e9,
#         -1.0e9 + 2.5e9j,
#         -1.0e9 - 2.5e9j,
#         -2.2e9
#     ]
#     freqs = np.linspace(1e8, 8e9, 400)
#     s = 1j * 2 * np.pi * freqs

#     basis = generate_muntz_laguerre_basis(s, stable_poles)
#     print(f"生成 {len(basis)} 个基函数,{basis}")