from TraditionalDetector.ml_detector import MLDetector
from TraditionalDetector.sd_detector import SDDetector
from TraditionalDetector.zf_detector import ZFDetector
from TraditionalDetector.mf_detector import MFDetector
from TraditionalDetector.mmse_detector import MMSEDetector
from TraditionalDetector.zfsic_detector import ZF_SICDetector
from TraditionalDetector.mmsesic_detector import MMSE_SICDetector
from TraditionalDetector.mfsic_detector import MF_SICDetector
from TraditionalDetector.lrzf_detector import LR_ZFDetector
from TraditionalDetector.lrmmse_detector import LR_MMSEDetector
from TraditionalDetector.mmselas_detector import MMSE_LASDetector
import matplotlib.pyplot as plt
from parameter import modinit
from tqdm import tqdm,trange
from dataset import GenData
import tensorflow as tf
import numpy as np
import time
import os

dt = time.localtime()
ft = '%Y%m%d%H%M%S'
nt = time.strftime(ft, dt)
print('==============当前时间：%s===============' % nt)

CUDA = input('GPU Device: ')
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
os.environ['CUDA_VISIBLE_DEVICES']=CUDA 

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

params = modinit()

# 对象实例化
gen_data = GenData(params)

# 实例化检测器
detector_name = params['algorithms']+'Detector'
detector = eval(detector_name)(params)

ser_dict = {}

# SNR 循环
for snrdb in trange(params['SNR_dB_MIN'], params['SNR_dB_MAX']+1):

    error_symbols, total_symbols = 0, 0
    snr = np.power(10.0, snrdb / 10.0)

    with tqdm(total=params['MIN_ERROR_SYMBOLS']) as pbar:

        # Batch 循环
        while error_symbols < params['MIN_ERROR_SYMBOLS']:

            # 生成数据
            x,H,y,sigma2 = gen_data.rayleigh_varying(snr,True)

            # 检测
            nodes = detector(x,H,y,sigma2)

            # 统计
            error_symbols += nodes['errornum']
            total_symbols += nodes['totalnum']

            # 进度条
            pbar.update(nodes['errornum'])

    # 计算 SER
    ser_dict[snrdb] = error_symbols / total_symbols

# 保存
np.save('Nt%d_Nr%d_mod%s_SNR%d_%d_%s_%s.npy'%(
            params['Nt'], params['Nr'], params['MOD_NAME'],
            params['SNR_dB_MIN'], params['SNR_dB_MAX'], 
            params['algorithms'], nt), ser_dict)

# 绘图
plt.plot(list(ser_dict.keys()), list(ser_dict.values()), 
            label=params['algorithms'])
plt.grid(True, which='both', linestyle='--')
plt.yscale('log')
plt.xlabel('SNR')
plt.ylabel('SER')
plt.ylim(1e-4, 1e-0)
plt.title('Nr%dNt%d_mod%s' % (params['Nr'], params['Nt'], params['MOD_NAME']))
plt.legend()
plt.show()

print(ser_dict)