#!/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
from utils.signal_available import high_check

preprocessing = BCG_Operation()
preprocessing.sample_rate = 100


def my_process(data, movement):
    normal_segments = [[0, movement[0][0]], [movement[-1][-1], len(data)]]
    for index in range(len(movement) - 1):
        normal_segments.append([movement[index][-1], movement[index + 1][0]])

    for index in range(len(normal_segments)):
        SP, EP = normal_segments[index]
        if SP == EP:
            # print(index, SP, EP)
            continue
        std = data[SP:EP].std()
        data[SP:EP][data[SP:EP] > std * 1.8] = std * 1.8
        data[SP:EP][data[SP:EP] < std * -1.8] = std * 1.8

        data[SP:EP] /= std * 1.8

    for index in range(len(movement)):
        SP, EP = movement[index]
        data[SP:EP] /= data[SP:EP].max()

    return data


def my_augment(dataset):
    movement = high_check(dataset)
    dataset = my_process(dataset, movement)
    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()

    # Z-SCORE
    dataset_low = (dataset_low - dataset_low.mean()) / dataset_low.std()
    dataset_high = (dataset_high - dataset_high.mean()) / dataset_high.std()

    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