#!/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 - 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.7, order=6)
    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[::25]
    dataset_low = dataset_low.reshape(dataset_low.shape[0], 1)

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

    return dataset_low, dataset_high


if __name__ == '__main__':
    pass