#!/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):
    if 1000 < dataset.mean() < 2000:
        dataset[dataset == 0] = 1825
        dataset[dataset < 1200] = 1225
        dataset[dataset > 2400] = 2425
    elif 2000 < dataset.mean() < 3000:
        dataset[dataset == 0] = 2375
        dataset[dataset < 1775] = 1775
        dataset[dataset > 2975] = 2975
        dataset -= 550
    else:
        print("something wrong")

    # select_dataset = preprocessing.Butterworth(select_dataset, "lowpass", low_cut=0.7, order=3)

    dataset -= dataset.mean()
    dataset = dataset.dot(1 / 600)

    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[::10]
    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