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

# ==========================================
# 0. 全局绘图参数与 CUDA 引擎初始化
# ==========================================
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['mathtext.fontset'] = 'stix' 

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

N_ant, N_RF, d_lambda = 24, 4, 0.5
clutter_angle = -15.0
noise_power = 1.0
INR_linear = 10**(30/10)

def gen_a(theta_deg):
    theta_rad = torch.deg2rad(torch.tensor(theta_deg, dtype=torch.float32, device=device))
    n_idx = torch.arange(-(N_ant-1)/2, (N_ant-1)/2 + 0.1, device=device)
    return torch.exp(1j * 2 * torch.pi * d_lambda * n_idx * torch.sin(theta_rad)).unsqueeze(1)

a_clutter = gen_a(clutter_angle)
R_interference = noise_power * torch.eye(N_ant, device=device) + INR_linear * (a_clutter @ a_clutter.mH)

dic_angles_full = torch.arange(-90, 90.5, 0.5, device=device)
A_dic_full = torch.cat([gen_a(th) / np.sqrt(N_ant) for th in dic_angles_full], dim=1)

angle_ranges = [np.arange(0, 11), np.arange(11, 21), np.arange(21, 31), np.arange(31, 41)]
snr_test_list = [0, 5, 10, 15, 20]
Monte_Carlo = 100
L_snapshots = 100
colors = ['#0072BD', '#D95319', '#EDB120', '#7E2F8E', '#77AC30']

def run_omp(w_opt, A_dic):
    res = w_opt.clone()
    F_RF_list = []
    for k in range(N_RF):
        idx = torch.argmax(torch.abs(A_dic.mH @ res))
        F_RF_list.append(A_dic[:, idx:idx+1])
        F_RF_mat = torch.cat(F_RF_list, dim=1)
        res = w_opt - F_RF_mat @ (torch.linalg.pinv(F_RF_mat) @ w_opt)
    return F_RF_mat

print(f"启动 CUDA 并行引擎，对 4 大扇区进行极速扫雷测试 (MC={Monte_Carlo})...")
start_time = time.time()

# 取消了这里原本的 fig, axs = plt.subplots(2, 2)
# 改为在循环内部动态创建单图

for fig_idx, test_angles in enumerate(angle_ranges):
    print(f"正在计算扇区 {test_angles[0]}° ~ {test_angles[-1]}° ...")
    
    # 【修改点 1】：为每一个扇区单独生成一张独立画布
    fig, ax = plt.subplots(figsize=(9, 7))
    
    RMSE_DFT_All = np.zeros((len(snr_test_list), len(test_angles)))
    RMSE_FULL_All = np.zeros((len(snr_test_list), len(test_angles)))
    RMSE_TSC_All = np.zeros((len(snr_test_list), len(test_angles)))
    
    for snr_idx, snr_dB in enumerate(snr_test_list):
        sig_p = 10**(snr_dB/10)
        
        for a_idx, t_angle in enumerate(test_angles):
            true_target_angle = t_angle + 0.04
            theta_prior = round(true_target_angle)
            window = 5
            
            dic_tsc = torch.arange(theta_prior - window, theta_prior + window + 0.1, 0.1, device=device)
            A_dic_tsc = torch.cat([gen_a(th) / np.sqrt(N_ant) for th in dic_tsc], dim=1)
            
            angles_dft = torch.linspace(theta_prior - window, theta_prior + window, N_RF, device=device)
            F_RF_dft = torch.cat([gen_a(th) / np.sqrt(N_ant) for th in angles_dft], dim=1)
            
            a_prior_broad = torch.zeros((N_ant, 1), dtype=torch.complex64, device=device)
            for tb in torch.arange(theta_prior - window, theta_prior + window + 0.1, 0.5, device=device):
                a_prior_broad += gen_a(tb)
            a_prior_broad /= torch.norm(a_prior_broad)
            
            w_dig_prior = torch.linalg.inv(R_interference) @ a_prior_broad
            w_dig_prior /= torch.norm(w_dig_prior)
            
            F_RF_full = run_omp(w_dig_prior, A_dic_full)
            F_RF_tsc = run_omp(w_dig_prior, A_dic_tsc)
            
            scan_grid = torch.arange(theta_prior - window, theta_prior + window + 0.01, 0.01, device=device)
            A_scan = torch.cat([gen_a(th) for th in scan_grid], dim=1)
            
            A_eff_dft = F_RF_dft.mH @ A_scan
            A_eff_full = F_RF_full.mH @ A_scan
            A_eff_tsc = F_RF_tsc.mH @ A_scan
            
            a_target = gen_a(true_target_angle)
            s_t = np.sqrt(sig_p/2) * (torch.randn(Monte_Carlo, 1, L_snapshots, device=device) + 1j * torch.randn(Monte_Carlo, 1, L_snapshots, device=device))
            s_c = np.sqrt(INR_linear/2) * (torch.randn(Monte_Carlo, 1, L_snapshots, device=device) + 1j * torch.randn(Monte_Carlo, 1, L_snapshots, device=device))
            n_t = np.sqrt(noise_power/2) * (torch.randn(Monte_Carlo, N_ant, L_snapshots, device=device) + 1j * torch.randn(Monte_Carlo, N_ant, L_snapshots, device=device))
            
            X = a_target.unsqueeze(0) * s_t + a_clutter.unsqueeze(0) * s_c + n_t
            
            def process_hybrid(F_RF, A_eff):
                Y = F_RF.mH.unsqueeze(0) @ X
                R = (Y @ Y.mH) / L_snapshots
                R += (0.05 * torch.diagonal(R, dim1=-2, dim2=-1).sum(-1).view(-1, 1, 1) / N_RF) * torch.eye(N_RF, device=device).unsqueeze(0)
                inv_R = torch.linalg.inv(R)
                P = 1.0 / torch.abs(torch.einsum('si, bij, js -> bs', A_eff.conj().T, inv_R, A_eff))
                return np.sqrt(np.mean(((scan_grid[torch.argmax(P, dim=1)] - true_target_angle)**2).cpu().numpy()))

            RMSE_DFT_All[snr_idx, a_idx] = process_hybrid(F_RF_dft, A_eff_dft)
            RMSE_FULL_All[snr_idx, a_idx] = process_hybrid(F_RF_full, A_eff_full)
            RMSE_TSC_All[snr_idx, a_idx] = process_hybrid(F_RF_tsc, A_eff_tsc)
            
        c_str = colors[snr_idx]
        lbl_suffix = f" (SNR={snr_dB}dB)"
        ax.semilogy(test_angles, RMSE_DFT_All[snr_idx, :], '--^', color=c_str, linewidth=1.5, label='DFT'+lbl_suffix)
        ax.semilogy(test_angles, RMSE_FULL_All[snr_idx, :], ':x', color=c_str, linewidth=2, label='传统OMP'+lbl_suffix)
        ax.semilogy(test_angles, RMSE_TSC_All[snr_idx, :], '-o', color=c_str, linewidth=2.5, label='TSC-OMP'+lbl_suffix)

    ax.set_xlim([test_angles[0], test_angles[-1]])
    ax.set_xticks(test_angles)
    ax.grid(True, which="both", ls="--", alpha=0.5)
    ax.set_ylim([1e-3, 20])
    ax.set_xlabel('目标真实出现角度 $\\theta (\\circ)$', fontsize=12, fontweight='bold')
    ax.set_ylabel(r'测角 RMSE $(^\circ)$', fontsize=12, fontweight='bold')
    ax.set_title(f'宽波束先验引导下，扇区 ({test_angles[0]}°~{test_angles[-1]}°) 鲁棒性评估', fontsize=13, fontweight='bold')


    # 【修改点 2】：独立设置每张图的图例。
    plt.tight_layout() 
    fig.subplots_adjust(bottom=0.25) 
    
    # 【关键调整】：将 bbox_to_anchor 的 Y 坐标从 -0.15 提上来，改为 -0.08 
    # 这样图例就会离 X 轴更近，避开底部操作条
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.10), ncol=3, fontsize=10)

    # 【修改点 3】：自动保存
    file_name = f'RMSE_Sector_{test_angles[0]}_{test_angles[-1]}.png'
    plt.savefig(file_name, dpi=300, bbox_inches='tight')
    print(f"  已保存独立图表: {file_name}")

print(f"全扇区计算完毕！总耗时: {time.time() - start_time:.4f} 秒。")

# 一次性在屏幕上弹出生成的4张独立图表
plt.show()