#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
@author:Marques
@file:my_augment.py
@email:admin@marques22.com
@email:2021022362@m.scnu.edu.cn
@time:2022/07/26
"""
from utils.Preprocessing import BCG_Operation
import numpy as np
from scipy.signal import stft

preprocessing = BCG_Operation()
preprocessing.sample_rate = 100


def my_augment(dataset):
    dataset -= dataset.mean()
    dataset = preprocessing.Iirnotch(dataset)
    dataset = preprocessing.Butterworth(dataset, "lowpass", low_cut=20, order=6)
    dataset_low = preprocessing.Butterworth(dataset, "lowpass", low_cut=0.5, order=4)
    dataset_low = (dataset_low - dataset_low.mean()) / dataset_low.std()
    # dataset_high = preprocessing.Butterworth(dataset, "highpass", high_cut=1, order=6)

    dataset = {"low": dataset_low}
    # "high": dataset_high}
    return dataset


def get_stft(x, fs, n):
    print(len(x))
    f, t, amp = stft(x, fs, nperseg=n)
    z = np.abs(amp.copy())
    return f, t, z


def my_segment_augment(dataset, SP, EP):
    dataset_low = dataset["low"][int(SP) * 100:int(EP) * 100].copy()
    # dataset_high = dataset["high"][int(SP) * 100:int(EP) * 100].copy()

    dataset_low = dataset_low[::10]

    # 获取整段的特征 （3，1）
    # 按照十秒窗获取 （3，3）
    # 按照十秒窗步进两秒获取 （3，21）

    sub_windows_size = 30
    stride = 1

    manual_feature = [[], [], []]

    SP = 0
    EP = sub_windows_size
    while EP <= sub_windows_size:
        # mean
        manual_feature[0].append(abs(dataset_low[SP:EP]).mean())
        # var
        manual_feature[1].append(abs(dataset_low[SP:EP]).var())
        # RMS
        manual_feature[2].append(np.sqrt((dataset_low[SP:EP] ** 2).mean()))

        SP += stride
        EP += stride

    dataset_low = dataset_low.reshape(-1, 1)
    manual_feature = np.array(manual_feature)
    manual_feature = manual_feature.reshape(1, -1)

    # _, _, dataset_high = stft(dataset_high, 100, nperseg=50)
    # dataset_high = dataset_high.astype(np.float).T
    # dataset_high = dataset_high.reshape(dataset_high.shape[0], dataset_high.shape[1])

    # return dataset_low, dataset_high

    return dataset_low, manual_feature


if __name__ == '__main__':
    pass