# sampling rate
fs = 10000000 # 10MHz

# signal length
# s, sampling interval, time array
t = np.arange(0, 0.5, 1 / fs)

# generate signal
f1 = 60 # 60Hz
signal_f = 2*np.sin(2*np.pi*f1*t)  # amplitude = 2

# generate noise signal
np.random.seed(222)
n1 = np.random.uniform(1000000, 1500000, len(t)) # 1MHz ~ 1.5 MHz
signal_n = 0.2*np.sin(2*np.pi*n1*t)

# total signal
normal_signal = signal_f + signal_n

#  generate noise Arc signal

n3 = 4000000 # 4MHz

np.random.seed(333)
random_mag_Arc = np.random.uniform(0.8, 1.1, 1500)
signal_Arc = random_mag_Arc*np.sin(2*np.pi*n3*t[82500:84000])

signal_Arc # 아크 신호

# shoulder 범위만 아크 신호를 더해주기 위해 타입변환
temp = list(normal_signal)

# shoulder 범위만 아크 신호 더하고 shoulder 범위가 아닌 부분은 아크신호를 더하지 않음.

# shoulder 범위에 아크신호를 더해주기 위해 범위 설정
i,j = 1, 0
s0 = temp[82500*0+1500*0:82500*1+1500*0]
end = (82500*1)+(1500*0)

# 신호 길이(데이터 개수) 끝까지 반복해서 noise 섞인 sin파 신호와 아크 신호 합성을 반복해줌.
while (len(s0) != 5000000):
    start = end
    np.random.seed(0)
    # 아크 신호를 랜덤하기 더해주기 위한 규칙 생성
    num = random.randint(0,2)
    
    if i == j: # shoulder 범위가 아닐때 그냥 신호 더해줌.
        i += 1
        end = (82500*i)+(1500*j)
        s0 = s0+temp[start:end]
    # shoulder 범위 이면서 num == 2일 경우 아크 신호 더함.
    elif i!=j and num == 2:
        j += 1
        end = (82500*i)+(1500*j)
        x = temp[start:end] + signal_Arc
        s0 = s0+ list(x) 
    else: # shoulder 범위 이나 num은 2가 아닐 경우 노말 신호만 더해줌.
        j += 1
        end = (82500*i)+(1500*j)
        s0 = s0+ temp[start:end] 

# numpy array type으로 변환, 아크 신호가 합쳐진 signal 데이터 완성        
sig = np.array(s0)

# ratio 계수에 윈도우 shift하면서 중복된 값을 제거해주며, 순서는 그대로 유지해주는 함수 생성
def unique(array):
    uniq, index = np.unique(array, return_index=True)
    return uniq[index.argsort()]

# Wavelet Packet Transform 변환 후 시각화 함수
def WPT(signal):
    wp = pywt.WaveletPacket(data = signal, wavelet = 'db10', mode = 'symmetric') # 'db10'의 wavelet family 사용
    xx = wp['ddd'].data # 'ddd'(3 level)에서 고주파 성분 추출
    
    # 논문 I_ratio 수식 인용
    shift_size = 6000 #1/4 cycle
    window_length = 12000 # 1/2 cycle
    ratio = []
    I_max = xx[:12000].max() # 첫 번째 window max값(고정)
    
    # window shift
    for i in np.arange(0, len(xx), 6000):
        if len(ratio) < 104:
            window_2 = xx[i + shift_size: i + shift_size + window_length]

            I_n_max = window_2.max()

            I_ratio = I_n_max / I_max

            ratio.append(I_ratio)
            
    ratio= unique(ratio) # ratio 중복 값 제거
    tt = np.linspace(0, 0.5, len(ratio))
    
    fig = plt.figure(figsize = (15, 6))
    # 그래프 2개 출력
    spec = gridspec.GridSpec(ncols=1, nrows=2, width_ratios=[1], height_ratios=[1, 3])
    ax0 = fig.add_subplot(spec[0])
    ax0.plot(xx)
    
    ax1 = fig.add_subplot(spec[1])
    ax1.scatter(tt, ratio)
    ax1.plot(tt, ratio, 'yx-')
    plt.axhline(y = 1.6, color = 'r', linewidth = 2) # warning threshold
    plt.axhline(y = 1.2, color = 'g', linewidth = 2) # attention threshold
    
    global count
    count = 0
    
    # count 개수에 따라 현재 상태 출력
    def state(count):
        if count > 2:
            print("Warning!")
        elif count > 1:
            print("Attention!")
        elif 0 <= count < 2:
            print("Safety")
        elif count < 0: # 이 부분 빼고 위에 위험도 단계별 count 기준 수정 필요.
            count = 0
            print("Safety")
    
    for i in ratio:
        if i > 1.6: # warning threshold 값 넣어주면 됨.
            count += 1
            state(count)
        elif 1.2 < i < 1.6: # attention threshold 값 넣어주면 됨.
            state(count)
        else:
            if count < 0:
                count= 0
                state(count)
            else:
                count -= 1
                state(count)
    
WPT(sig)