import os
import math
import time
import numpy as np
import torch
import matplotlib.pyplot as plt

# ============================================================
# 0. 全局设置
# ============================================================
plt.rcParams['font.sans-serif'] = [
    'Microsoft YaHei', 'SimHei', 'PingFang SC',
    'Noto Sans CJK SC', 'Arial Unicode MS', 'DejaVu Sans'
]
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['mathtext.fontset'] = 'stix'

SAVE_DIR = "robust_tsc_omp_cuda_results"
os.makedirs(SAVE_DIR, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(2024)
np.random.seed(2024)

# ============================================================
# 1. 参数区
#   说明：
#   - 本脚本默认使用 CUDA（若可用）
#   - 为避免“卡死”，默认采用较稳妥的 Monte Carlo 与扫描步长
#   - 所有图都会生成并在脚本结束时统一以窗口形式显示
# ============================================================
ARRAY_TYPE = 'NULA'          # 'ULA' 或 'NULA'
N_ant = 24
N_RF = 4
d_lambda = 0.5

target_angle = 5.0
clutter_angle = -15.0
noise_power = 1.0
INR_dB = 30.0
INR_linear = 10 ** (INR_dB / 10)

L_snapshots = 100
Monte_Carlo = 300          
SNR_dB_range = np.arange(-10, 22, 2)

theta_prior = 5.0
window_half_width = 5.0
sector_min = theta_prior - window_half_width
sector_max = theta_prior + window_half_width

# TSC-OMP 字典与测角搜索栅格
tsc_grid_np = np.arange(sector_min, sector_max + 0.1, 0.1)
scan_grid_np = np.arange(sector_min, sector_max + 0.001, 0.05)  # 比 0.01 更快

# 对角加载系数
alpha_list = [0.00, 0.01, 0.03, 0.05, 0.10]

# 是否保存图片
SAVE_FIG = True

# ============================================================
# 2. 阵列位置
# ============================================================
def get_array_positions(array_type='ULA', N=24, d=0.5):
    if array_type.upper() == 'ULA':
        pos = np.arange(-(N - 1) / 2, (N - 1) / 2 + 1) * d
        return pos.astype(np.float32)
    elif array_type.upper() == 'NULA':
        # 与论文一致的示意性 NULA
        base = 0.7
        jump = 1.03
        pos = []
        x = 0.0
        for _ in range(12):
            pos.append(x)
            x += base
        x += (jump - base)
        for _ in range(12, 24):
            pos.append(x)
            x += base
        pos = np.array(pos, dtype=np.float32)
        pos = pos - np.mean(pos)
        return pos
    else:
        raise ValueError("ARRAY_TYPE 只能是 'ULA' 或 'NULA'")

array_pos_np = get_array_positions(ARRAY_TYPE, N_ant, d_lambda)
array_pos_t = torch.tensor(array_pos_np, dtype=torch.float32, device=device).view(-1, 1)

# ============================================================
# 3. 基础函数（CUDA）
# ============================================================
def sync():
    if device.type == 'cuda':
        torch.cuda.synchronize()

def steering_vector(theta_deg):
    theta = torch.as_tensor(theta_deg, dtype=torch.float32, device=device)
    theta_rad = torch.deg2rad(theta)
    phase = 2 * torch.pi * array_pos_t * torch.sin(theta_rad)
    return torch.exp(1j * phase)

def steering_matrix(theta_deg_array):
    theta = torch.as_tensor(theta_deg_array, dtype=torch.float32, device=device).view(1, -1)
    theta_rad = torch.deg2rad(theta)
    phase = 2 * torch.pi * array_pos_t @ torch.sin(theta_rad)
    return torch.exp(1j * phase)

def build_dictionary(angle_grid_np, normalize=True):
    A = steering_matrix(angle_grid_np)
    if normalize:
        A = A / math.sqrt(A.shape[0])
    return A

def build_prior_broad(sector_min, sector_max, step=0.5):
    angles = np.arange(sector_min, sector_max + 1e-12, step, dtype=np.float32)
    A = steering_matrix(angles)
    a = torch.sum(A, dim=1, keepdim=True)
    a = a / torch.linalg.norm(a)
    return a

def extract_omp(w_opt, A_dic, angle_grid_np, N_RF):
    """
    CUDA 版 OMP / TSC-OMP
    """
    res = w_opt.clone()
    F_RF = None
    chosen_angles = []

    for _ in range(N_RF):
        proj = torch.abs(A_dic.mH @ res).flatten()
        idx = int(torch.argmax(proj).item())
        atom = A_dic[:, idx:idx+1]
        F_RF = atom if F_RF is None else torch.cat((F_RF, atom), dim=1)
        chosen_angles.append(float(angle_grid_np[idx]))
        res = w_opt - F_RF @ (torch.linalg.pinv(F_RF) @ w_opt)

    return F_RF, chosen_angles

def build_prior_weight(target_angle, clutter_angle, sector_min, sector_max):
    a_clutter = steering_vector(clutter_angle)
    R_interf = noise_power * torch.eye(N_ant, dtype=torch.complex64, device=device) \
               + INR_linear * (a_clutter @ a_clutter.mH)
    a_prior = build_prior_broad(sector_min, sector_max, step=0.5)
    w_prior = torch.linalg.solve(R_interf, a_prior)
    w_prior = w_prior / torch.linalg.norm(w_prior)
    return w_prior

def calc_hybrid_weight_loaded(F_RF, R_in, a_target, alpha_dl=0.05):
    a_eff = F_RF.mH @ a_target
    R_eff = F_RF.mH @ R_in @ F_RF
    gamma = alpha_dl * torch.real(torch.trace(R_eff)) / R_eff.shape[0]
    R_eff_loaded = R_eff + gamma * torch.eye(R_eff.shape[0], dtype=torch.complex64, device=device)
    inv_R = torch.linalg.inv(R_eff_loaded)
    F_BB = (inv_R @ a_eff) / (a_eff.mH @ inv_R @ a_eff)
    w = F_RF @ F_BB
    w = w / torch.linalg.norm(w)
    cond_number = torch.linalg.cond(R_eff_loaded).real.item()
    return w, R_eff_loaded, cond_number

def robust_capon_spectrum_batch(F_RF, X_batch, A_scan, alpha_dl=0.05):
    """
    X_batch: [MC, N_ant, L]
    A_scan:  [N_ant, Ns]
    返回：
        P: [MC, Ns]
        mean_cond: 平均条件数
    """
    Y = torch.matmul(F_RF.mH.unsqueeze(0), X_batch)                     # [MC, N_RF, L]
    R_bb = torch.matmul(Y, Y.mH) / X_batch.shape[-1]                   # [MC, N_RF, N_RF]

    trace_R = torch.real(torch.diagonal(R_bb, dim1=-2, dim2=-1).sum(-1))
    gamma = alpha_dl * trace_R / F_RF.shape[1]
    eye = torch.eye(F_RF.shape[1], dtype=torch.complex64, device=device).unsqueeze(0)
    R_bb_loaded = R_bb + gamma[:, None, None] * eye

    inv_R = torch.linalg.inv(R_bb_loaded)                               # [MC, N_RF, N_RF]
    A_eff = F_RF.mH @ A_scan                                            # [N_RF, Ns]

    # P = 1 / |a^H invR a|
    denom = torch.einsum('rn,brs,sn->bn', A_eff.conj(), inv_R, A_eff)
    P = 1.0 / torch.clamp(torch.abs(denom), min=1e-12)

    cond_vals = torch.linalg.cond(R_bb_loaded).real
    mean_cond = torch.mean(cond_vals).item()
    return P, mean_cond

def calc_sinr(w, a_target, a_clutter, sig_p):
    R_interf = noise_power * torch.eye(N_ant, dtype=torch.complex64, device=device) \
               + INR_linear * (a_clutter @ a_clutter.mH)
    num = sig_p * (torch.abs(w.mH @ a_target).item() ** 2)
    den = torch.real(w.mH @ R_interf @ w).item()
    return 10 * np.log10(max(num / max(den, 1e-12), 1e-12))

# ============================================================
# 4. 固定 TSC-OMP 模拟矩阵
# ============================================================
print(f"当前设备: {device}")
sync()
tic = time.perf_counter()

w_prior = build_prior_weight(
    target_angle=target_angle,
    clutter_angle=clutter_angle,
    sector_min=sector_min,
    sector_max=sector_max
)
A_dic_tsc = build_dictionary(tsc_grid_np, normalize=True)
F_RF_tsc, chosen_angles = extract_omp(w_prior, A_dic_tsc, tsc_grid_np, N_RF)

a_target = steering_vector(target_angle)
a_clutter = steering_vector(clutter_angle)

sync()
print(f"TSC-OMP 选中的模拟原子角度: {chosen_angles}")
print(f"固定模拟矩阵构建耗时: {time.perf_counter() - tic:.3f} s")

# ============================================================
# 5. 静态方向图：不同加载系数
# ============================================================
sig_power_static = 10 ** (0 / 10.0)
s_t = np.sqrt(sig_power_static / 2) * (
    torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
    + 1j * torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
)
s_c = np.sqrt(INR_linear / 2) * (
    torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
    + 1j * torch.randn(1, L_snapshots, dtype=torch.float32, device=device)
)
n_t = np.sqrt(noise_power / 2) * (
    torch.randn(N_ant, L_snapshots, dtype=torch.float32, device=device)
    + 1j * torch.randn(N_ant, L_snapshots, dtype=torch.float32, device=device)
)
X_static = a_target @ s_t + a_clutter @ s_c + n_t
R_in_static = (X_static @ X_static.mH) / L_snapshots

theta_scan_np = np.arange(-60, 60.1, 0.1)
A_bp = steering_matrix(theta_scan_np)

plt.figure(figsize=(10, 6))
for alpha_dl in alpha_list:
    w_hyb, _, _ = calc_hybrid_weight_loaded(F_RF_tsc, R_in_static, a_target, alpha_dl=alpha_dl)
    BP = torch.abs(w_hyb.mH @ A_bp).flatten() ** 2
    BP = 10 * torch.log10(BP / torch.max(BP) + 1e-12)
    plt.plot(theta_scan_np, BP.detach().cpu().numpy(), linewidth=2, label=f'α={alpha_dl:.2f}')

plt.axvline(target_angle, color='g', linestyle='--', label=f'目标 {target_angle:.2f}°')
plt.axvline(clutter_angle, color='r', linestyle='--', label=f'杂波 {clutter_angle:.2f}°')
plt.xlabel('角度 θ (°)')
plt.ylabel('归一化增益 (dB)')
plt.title(f'不同对角加载系数下 TSC-OMP 波束方向图 ({ARRAY_TYPE})')
plt.grid(True, linestyle='--', alpha=0.5)
plt.legend()
plt.tight_layout()
if SAVE_FIG:
    plt.savefig(os.path.join(SAVE_DIR, f'BP_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)

# ============================================================
# 6. 输出 SINR / RMSE / 条件数 对比
# ============================================================
SINR_loaded = {alpha: [] for alpha in alpha_list}
RMSE_loaded = {alpha: [] for alpha in alpha_list}
COND_loaded = {alpha: [] for alpha in alpha_list}

A_scan = steering_matrix(scan_grid_np)

for snr_db in SNR_dB_range:
    sync()
    tic = time.perf_counter()

    sig_p = 10 ** (snr_db / 10.0)

    # 一次性批量生成 Monte Carlo 数据
    s_t = np.sqrt(sig_p / 2) * (
        torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
        + 1j * torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
    )
    s_c = np.sqrt(INR_linear / 2) * (
        torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
        + 1j * torch.randn(Monte_Carlo, 1, L_snapshots, dtype=torch.float32, device=device)
    )
    n_t = np.sqrt(noise_power / 2) * (
        torch.randn(Monte_Carlo, N_ant, L_snapshots, dtype=torch.float32, device=device)
        + 1j * torch.randn(Monte_Carlo, N_ant, L_snapshots, dtype=torch.float32, device=device)
    )
    X_batch = a_target.unsqueeze(0) * s_t + a_clutter.unsqueeze(0) * s_c + n_t
    R_in_batch = torch.matmul(X_batch, X_batch.mH) / L_snapshots

    for alpha_dl in alpha_list:
        # --- SINR：使用每个 Monte Carlo 样本求权，最后取平均 ---
        sinr_vals = []
        cond_eff_vals = []

        for mc in range(Monte_Carlo):
            w_hyb, _, cond_eff = calc_hybrid_weight_loaded(
                F_RF_tsc, R_in_batch[mc], a_target, alpha_dl=alpha_dl
            )
            sinr_vals.append(calc_sinr(w_hyb, a_target, a_clutter, sig_p))
            cond_eff_vals.append(cond_eff)

        SINR_loaded[alpha_dl].append(float(np.mean(sinr_vals)))

        # --- RMSE 与 R_bb 条件数：批量 Capon 谱 ---
        P, mean_cond_bb = robust_capon_spectrum_batch(
            F_RF_tsc, X_batch, A_scan, alpha_dl=alpha_dl
        )
        est_idx = torch.argmax(P, dim=1)
        est_angles = torch.as_tensor(scan_grid_np, dtype=torch.float32, device=device)[est_idx]
        rmse = torch.sqrt(torch.mean((est_angles - target_angle) ** 2)).item()

        RMSE_loaded[alpha_dl].append(rmse)
        COND_loaded[alpha_dl].append(mean_cond_bb)

    sync()
    print(f"SNR={snr_db:>3} dB 完成，耗时 {time.perf_counter() - tic:.2f} s")

# ============================================================
# 7. 输出 SINR 图
# ============================================================
plt.figure(figsize=(10, 6))
for alpha_dl in alpha_list:
    plt.plot(SNR_dB_range, SINR_loaded[alpha_dl], marker='o', linewidth=2, label=f'α={alpha_dl:.2f}')
plt.xlabel('输入 SNR (dB)')
plt.ylabel('输出 SINR (dB)')
plt.title(f'不同对角加载系数下 TSC-OMP 输出 SINR 对比 ({ARRAY_TYPE})')
plt.grid(True, linestyle='--', alpha=0.5)
plt.legend()
plt.tight_layout()
if SAVE_FIG:
    plt.savefig(os.path.join(SAVE_DIR, f'SINR_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)

# ============================================================
# 8. RMSE 图
# ============================================================
plt.figure(figsize=(10, 6))
for alpha_dl in alpha_list:
    plt.semilogy(SNR_dB_range, RMSE_loaded[alpha_dl], marker='s', linewidth=2, label=f'α={alpha_dl:.2f}')
plt.xlabel('输入 SNR (dB)')
plt.ylabel('测角 RMSE (°)')
plt.title(f'不同对角加载系数下 TSC-OMP 测角 RMSE 对比 ({ARRAY_TYPE})')
plt.grid(True, which='both', linestyle='--', alpha=0.5)
plt.legend()
plt.tight_layout()
if SAVE_FIG:
    plt.savefig(os.path.join(SAVE_DIR, f'RMSE_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)

# ============================================================
# 9. 条件数图
# ============================================================
plt.figure(figsize=(10, 6))
for alpha_dl in alpha_list:
    plt.semilogy(SNR_dB_range, COND_loaded[alpha_dl], marker='d', linewidth=2, label=f'α={alpha_dl:.2f}')
plt.xlabel('输入 SNR (dB)')
plt.ylabel('低维协方差矩阵条件数')
plt.title(f'不同对角加载系数下低维协方差矩阵条件数 ({ARRAY_TYPE})')
plt.grid(True, which='both', linestyle='--', alpha=0.5)
plt.legend()
plt.tight_layout()
if SAVE_FIG:
    plt.savefig(os.path.join(SAVE_DIR, f'COND_loaded_compare_{ARRAY_TYPE}.png'), dpi=300)

# ============================================================
# 10. 低SNR点 α 扫描汇总图
# ============================================================
low_snr_pick = 0
alpha_arr = np.array(alpha_list, dtype=float)
rmse_pick = np.array([RMSE_loaded[a][np.where(SNR_dB_range == low_snr_pick)[0][0]] for a in alpha_list], dtype=float)
cond_pick = np.array([COND_loaded[a][np.where(SNR_dB_range == low_snr_pick)[0][0]] for a in alpha_list], dtype=float)

fig = plt.figure(figsize=(10, 6))
ax1 = fig.add_subplot(111)
ax1.plot(alpha_arr, rmse_pick, marker='o', linewidth=2, color='#0072BD', label='RMSE')
ax1.set_xlabel('对角加载系数 α')
ax1.set_ylabel('RMSE (°)', color='#0072BD')
ax1.tick_params(axis='y', labelcolor='#0072BD')
ax1.grid(True, linestyle='--', alpha=0.5)

ax2 = ax1.twinx()
ax2.semilogy(alpha_arr, cond_pick, marker='s', linewidth=2, color='#D95319', label='条件数')
ax2.set_ylabel('条件数', color='#D95319')
ax2.tick_params(axis='y', labelcolor='#D95319')

plt.title(f'低SNR={low_snr_pick} dB 下加载系数对 RMSE 与条件数的影响 ({ARRAY_TYPE})')
fig.tight_layout()
if SAVE_FIG:
    plt.savefig(os.path.join(SAVE_DIR, f'alpha_sweep_lowSNR_{ARRAY_TYPE}.png'), dpi=300)

# ============================================================
# 11. 打印总结
# ============================================================
print("\n===== 鲁棒性提升实验完成 =====")
for alpha_dl in alpha_list:
    print(f"α={alpha_dl:.2f} : "
          f"平均SINR={np.mean(SINR_loaded[alpha_dl]):.3f} dB, "
          f"平均RMSE={np.mean(RMSE_loaded[alpha_dl]):.4f}°, "
          f"平均条件数={np.mean(COND_loaded[alpha_dl]):.3e}")

print(f"\n结果图片已保存到: {SAVE_DIR}")
print("即将以窗口形式显示所有图片……")
plt.show()