# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/09_ecg_analysis.ipynb.

# %% auto 0
__all__ = ['vis_ecg', 'get_hrv_df']

# %% ../nbs/09_ecg_analysis.ipynb 3
from typing import Tuple

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import neurokit2 as nk

# %% ../nbs/09_ecg_analysis.ipynb 4
def vis_ecg(values_df: pd.DataFrame) -> None:
    """
    Visualize ECG data for 12 leads.

    Args:
        values_df (pd.DataFrame): A DataFrame containing ECG data with 12 columns, one for each lead.

    Returns:
        None: Displays a 3x4 grid of ECG plots for the 12 leads.
    """
    count = 0
    frequency = 1000
    time_data = np.arange(values_df.shape[0]) / frequency
    sns.set(rc={'figure.figsize':(40.7,18.27)})
    fig, axs = plt.subplots(3, 4)
    
    for y in range(0,4):    
        for x in range(0,3):
            col = values_df.columns[count]
            axs[x, y].plot(time_data, values_df[col].values)
            axs[x, y].set_title(col)
            count += 1
            plt.xlabel("time in seconds")
            plt.ylabel("ECG in uV")


# %% ../nbs/09_ecg_analysis.ipynb 5
def get_hrv_df(ECG_df: pd.DataFrame, sr: int = 1000) -> pd.DataFrame:
    """
    Compute the Heart Rate Variability (HRV) metrics for each ECG lead in the input DataFrame.

    Args:
        ECG_df (pd.DataFrame): A DataFrame containing ECG data with one column for each lead.
        sr (int, optional): The sampling rate of the ECG data. Defaults to 1000.

    Returns:
        pd.DataFrame: A DataFrame containing HRV metrics for each ECG lead.
    """
    hrv_list = []
    peaks_list = []
    
    for i in ECG_df.columns:
        peaks, info = nk.ecg_peaks(ECG_df[i], sampling_rate=sr, correct_artifacts=True)
        hrv = nk.hrv_time(peaks, sampling_rate=sr, show=False)
        hrv_list.append(hrv)
        peaks_list.append(peaks)

    hrv_df = pd.concat(hrv_list)
    hrv_df.index = ECG_df.columns
    hrv_df = hrv_df.dropna(axis=1, how='all')
    
    peaks_df = pd.concat(peaks_list, axis=1)
    peaks_df.columns = ECG_df.columns

    return hrv_df

