#!/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
"""
import torch.cuda
import yaml

from utils.Preprocessing import BCG_Operation
import numpy as np
from scipy.signal import stft
from torch import from_numpy
with open("./settings.yaml") as f:
    hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps

apply_samplerate = hyp["apply_samplerate"]
dataset_samplerate = hyp["dataset_samplerate"]
preprocessing = BCG_Operation()
preprocessing.sample_rate = dataset_samplerate


def my_augment(dataset, config):
    # dataset = preprocessing.Butterworth(dataset, "lowpass", low_cut=20, order=6)
    # dataset = (dataset - dataset.mean()) / dataset.std()
    # dataset_low = preprocessing.Butterworth(dataset, "lowpass", low_cut=0.7, order=6)
    # dataset_high = preprocessing.Butterworth(dataset, "highpass", high_cut=1, order=6)
    print(f"dataset sample_rate is {config['dataset_samplerate']}  down_ratio is {config['dataset_samplerate'] // config['apply_samplerate']}")
    dataset = dataset[::config['dataset_samplerate'] // config['apply_samplerate']]
    gpu = torch.cuda.is_available()
    dataset = {"raw": from_numpy(dataset).float().cuda() if gpu else from_numpy(dataset).float(),
               # "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_segment1 = dataset["low"][int(SP) * 100:int(EP) * 100].copy()
    # dataset_segment2 = dataset["high"][int(SP) * 100:int(EP) * 100].copy()

    # dataset_segment = np.concatenate(([dataset_segment1], [dataset_segment2]), axis=0)
    dataset_segment = dataset["raw"][int(SP):int(EP)].unsqueeze(dim=0)
    return [dataset_segment]


if __name__ == '__main__':
    pass