import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pywt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
# 解決負號顯示問題
plt.rcParams['axes.unicode_minus'] = False  # 解決保存圖像是負號'-'顯示為方塊的問題
plt.rcParams.update({'text.usetex': False, 'font.family': 'serif', 'font.serif': 'cmr10', 'mathtext.fontset': 'cm'})
font1 = {'family': 'SimHei', 'weight': 'normal', 'size': 12}
font2 = {'family': 'Times New Roman', 'weight': 'normal', 'size': 18}
label = {'family': 'SimHei', 'weight': 'normal', 'size': 15}
xlsx_path = "../小波能量譜作圖.xlsx"
sheet_name = "表名"      
data_arr = pd.read_excel(xlsx_path, sheet_name=sheet_name)
column_name = '列名'     
row = 1024
y = data_arr[column_name][0:row]
x = data_arr['time'][0:row]
scale = np.arange(1, 50)
wavelet = 'gaus1'   # 'morl'  'gaus1'  小波基函數
# 時間-尺度小波能量譜
def time_scale_spectrum():
    coefs, freqs = pywt.cwt(y, scale, wavelet)  # np.arange(1, 31) 第一個參數必須 >=1     'morl'  'gaus1'
    scale_freqs = np.power(freqs, -1)  # 對頻率freqs 取倒數變為尺度
    fig = plt.figure(figsize=(5, 4))
    ax = Axes3D(fig)
    # X:time   Y:Scale   Z:Amplitude
    X = np.arange(0, row, 1)  # [0-1023]
    Y = scale_freqs
    X, Y = np.meshgrid(X, Y)
    Z = abs(coefs)
    # 繪制三維曲面圖
    ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='rainbow')
    # 設置三個坐標軸信息
    ax.set_xlabel('$Mileage/km$', color='b', fontsize=12)
    ax.set_ylabel('$Scale$', color='g', fontsize=12)
    ax.set_zlabel('$Amplitude/mm$', color='r', fontsize=12)
    plt.draw()
    plt.show()
# 時間小波能量譜
def time_spectrum():
    coefs, freqs = pywt.cwt(y, scale, wavelet)
    coefs_pow = np.power(coefs, 2)      # 對二維數組中的數平方
    spectrum_value = [0] * row    # len(freqs)
    # 將二維數組按照里程疊加每個里程上的所有scale值
    for i in range(row):
        sum = 0
        for j in range(len(freqs)):
            sum += coefs_pow[j][i]
        spectrum_value[i] = sum
    fig = plt.figure(figsize=(7, 2))
    line_width = 1
    line_color = 'dodgerblue'
    line_style = '-'
    T1 = fig.add_subplot(1, 1, 1)
    T1.plot(x, spectrum_value, label='模擬', linewidth=line_width, color=line_color, linestyle=line_style)
    # T1.legend(loc='upper right', prop=font1, frameon=True)  # lower ,left
    # 坐標軸名稱
    T1.set_xlabel('$time$', fontsize=15, fontdict=font1)  # fontdict設置子圖字體
    T1.set_ylabel('$E/mm^2$', fontsize=15, fontdict=font1)
    # 坐標刻度值字體大小
    T1.tick_params(labelsize=15)
    print(spectrum_value[269])
    plt.show()
# 尺度小波能量譜
def scale_spectrum():
    coefs, freqs = pywt.cwt(y, scale, wavelet)
    coefs_pow = np.power(coefs, 2)      # 對二維數組中的數平方
    scale_freqs = np.power(freqs, -1)   # 對頻率freqs 取倒數變為尺度
    spectrum_value = [0] * len(freqs)    # len(freqs)
    # 將二維數組按照里程疊加每個里程上的所有scale值
    for i in range(len(freqs)):
        sum = 0
        for j in range(row):
            sum += coefs_pow[i][j]
        spectrum_value[i] = sum
    fig = plt.figure(figsize=(7, 4))
    line_width = 1
    line_color1 = 'dodgerblue'
    line_style1 = '-'
    T1 = fig.add_subplot(1, 1, 1)
    T1.plot(scale_freqs, spectrum_value, label=column_name, linewidth=line_width, color=line_color1, linestyle=line_style1)
    # T1.legend(loc='upper right', prop=font1, frameon=True)  # lower ,left
    # 坐標軸名稱
    T1.set_xlabel('$Scale$', fontsize=15, fontdict=font1)  # fontdict設置子圖字體
    T1.set_ylabel('$E/mm^2$', fontsize=15, fontdict=font1)
    # 坐標刻度值字體大小
    T1.tick_params(labelsize=15)
    plt.show()
# 通過調用下面三個不同的函數選擇繪制能量譜
time_scale_spectrum()
# time_spectrum()
# scale_spectrum()