import numpy as np
from core.orthonormal_basis import generate_muntz_laguerre_basis

# ---------- 可选：离散再正交 (加权 QR) ----------
# def trapezoid_weights(freqs: np.ndarray):
#     if len(freqs) == 1:
#         return np.ones(1)
#     df = np.diff(freqs)
#     w = np.zeros_like(freqs, dtype=float)
#     w[0] = 0.5 * df[0]
#     w[-1] = 0.5 * df[-1]
#     if len(freqs) > 2:
#         w[1:-1] = 0.5 * (df[:-1] + df[1:])
#     return w

def weighted_qr_from_basis(basis_cols: list[np.ndarray], weights: np.ndarray | None = None):
    A = np.column_stack(basis_cols)  # (Nf, M)
    if weights is None:
        sw = np.ones(A.shape[0])
    else:
        sw = np.sqrt(weights.real)
    Aw = sw[:, None] * A
    Qw, R = np.linalg.qr(Aw)
    Phi = Qw / (sw[:, None] + 1e-30)
    return Phi, R  # Raw = Phi R

def check_discrete_orthogonality(Phi: np.ndarray, w: np.ndarray):
    G = Phi.conj().T @ (w[:, None] * Phi)
    off = np.max(np.abs(G - np.eye(G.shape[0])))
    return G, off

def verify_orthonormal(Phi: np.ndarray, w: np.ndarray, atol=1e-10, rtol=1e-8):
    """
    返回:
        G        : Gram 矩阵 (Φ^H W Φ)
        max_off  : 最大非对角幅值
        diag_err : max |diag(G)-1|
        passed   : 是否满足阈值
    """
    G = Phi.conj().T @ (w[:, None] * Phi)
    I = np.eye(G.shape[0])
    diag_err = np.max(np.abs(np.diag(G) - 1.0))
    max_off = np.max(np.abs(G - I + np.diag(np.diag(G) - 1.0)))
    passed = (diag_err <= atol) and (max_off <= rtol)
    return G, max_off, diag_err, passed

def omega_weights(freqs_hz: np.ndarray):
    """
    基于 ω=2πf 的梯形法得到 w_ω = Δω/(2π)，使得
    (1/2π) ∫_{-∞}^{∞} → Σ w_ω,k
    """
    f = freqs_hz
    if len(f) == 1:
        return np.ones(1)
    df = np.diff(f)
    w_f = np.zeros_like(f)
    w_f[0] = 0.5 * df[0]
    w_f[-1] = 0.5 * df[-1]
    if len(f) > 2:
        w_f[1:-1] = 0.5 * (df[:-1] + df[1:])
    # dω = 2π df,  (1/2π) * dω = df  ⇒ 直接 w_f 就是 w_ω
    return w_f  # 已等价于 Δω/(2π)

def evaluate(basis,freqs):
    w = omega_weights(freqs)
    Phi_num, R = weighted_qr_from_basis(basis, w)
    Gram, off = check_discrete_orthogonality(Phi_num, w)
    print("离散 Gram 最大非对角元素 =", off)
    print("R 形状:", R.shape)
    # 验证 Raw ≈ Phi R
    raw = np.column_stack(basis)
    err = np.max(np.abs(raw - Phi_num @ R))
    print("重构误差 ||Raw - Phi R||_∞ =", err)

    # 验证正交性
    print("离散 Gram 矩阵 (前5x5):")
    print(Gram[:5, :5])

    Gcheck, max_off, diag_err, ok = verify_orthonormal(Phi_num, w)
    print(f"Diag 误差={diag_err:.3e}, Max off={max_off:.3e}, Orthonormal={ok}")
    # 额外: 验证 R
    # raw = Φ R  =>  R ≈ Φ^H W raw  (因为 Φ^H W Φ = I)
    R_alt = Phi_num.conj().T @ (w[:,None] * raw)
    print("R 差异 ||R - R_alt||_max =", np.max(np.abs(R - R_alt)))

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

    basis = generate_muntz_laguerre_basis(s, init_poles)
    print("解析基函数数量 =", len(basis))
    print("基函数:")
    for i in range(len(basis)):
        print(f"φ_{i}:", basis[i][:])

    evaluate(basis, freqs)