#!/usr/bin/env python

# coding: utf-8

# In[5]:

import neurokit2 as nk

import numpy as np

import pandas as pd

import parameters as para

# In[26]:

def simulation(normal_N,abnormal_N, save_params = False):

    normal_data = []

    normal_params = []

    print('Creating normal dataset')

    for i in range(normal_N):

        ti = np.random.normal(para.mu_t_1, para.sigma_t_1)

        ai = np.random.normal(para.mu_a_1, para.sigma_a_1)

        bi = np.random.normal(para.mu_b_1, para.sigma_b_1)

        hr = np.random.normal(para.mu_hr_1, para.sigma_hr_1)

        noise = np.random.uniform(low=para.min_noise_1, high=para.max_noise_1)

        ecgs, _ = nk.ecg_simulate(duration=para.duration*2, sampling_rate=para.sampling_rate, noise=noise, Anoise=para.Anoise, heart_rate=hr, gamma=para.gamma, ti=ti, ai=ai, bi=bi)

        ecgs = np.array(ecgs)

        start_i = np.random.randint(len(ecgs[0])//4, len(ecgs[0])//2)

        normal_data.append(ecgs[:,start_i:start_i+para.sampling_rate*para.duration])

        normal_params.append({'ti':ti, 'ai':ai, 'bi':bi, 'hr':hr, 'noise':noise, 'gamma': para.gamma})

    abnormal_data = []

    abnormal_params = []

    print('Creating abnormal dataset')

    for i in range(abnormal_N):

        ti = np.random.normal(para.mu_t_2, para.sigma_t_2)

        ai = np.random.normal(para.mu_a_2, para.sigma_a_2)

        bi = np.random.normal(para.mu_b_2, para.sigma_b_2)

        hr = np.random.normal(para.mu_hr_2, para.sigma_hr_2)

        noise = np.random.uniform(low=para.min_noise_2, high=para.max_noise_2)

        ecgs, _ = nk.ecg_simulate(duration=para.duration*2, sampling_rate=para.sampling_rate, noise=noise, Anoise=para.Anoise, heart_rate=hr, gamma=para.gamma, ti=ti, ai=ai, bi=bi)

        ecgs = np.array(ecgs)

        start_i = np.random.randint(len(ecgs[0])//4, len(ecgs[0])//2)

        abnormal_data.append(ecgs[:,start_i:start_i+para.sampling_rate*para.duration])

        abnormal_params.append({'ti':ti, 'ai':ai, 'bi':bi, 'hr':hr, 'noise':noise, 'gamma': para.gamma})

    labels = np.array([0]*len(normal_data) + [1]*len(abnormal_data))

    permutation = np.random.permutation(len(labels))

    data = np.array(normal_data+abnormal_data)

    data_params = np.array(normal_params+abnormal_params)

    labels = labels[permutation]

    data = data[permutation]

    data_params = data_params[permutation]

    np.save('sim_ecg_data', data) # save ECG data

    np.save('sim_ecg_labels', labels) # save label

    if (save_params):

        np.save('sim_ecg_params',data_params) # save parameters for each ecg sample (12,2500)