import numpy as np
import pandas as pd
import datetime

import plotly.graph_objects as go
from wfdb import processing

from ecg import ECG
from ecg_rhythm_analyzer import RhythmAnalyzer
from wfdb_db_reader import WfdbDbReader

wfdb_db_reader = WfdbDbReader('./physionet.org/files/mitdb/1.0.0/')

ecg_inst = wfdb_db_reader.get_ecg('100')
ecg_ann_anomalies = [beat for beat in wfdb_db_reader.get_annotation('100') 
                     if beat[1] in ['A', 'a', 'V']]

Getting list of record names from './physionet.org/files/mitdb/1.0.0/RECORDS'...

rhythm_analyzer = RhythmAnalyzer(ecg_inst)
rr_anomalies_simple = rhythm_analyzer.detect_anomalies_simple(window_size=41, method='mean')

Learning initial signal parameters...
Found 8 beats during learning. Initializing using learned parameters
Running QRS detection...
QRS detection complete.

rolling_mean = pd.Series(ecg_inst.rr_intervals).rolling(window=41, center=True).mean().fillna(ecg_inst.rr_intervals.mean()).to_numpy()
rr_anomalies_simple = rhythm_analyzer.detect_anomalies_simple()

fig = go.Figure()

fig.add_trace(go.Scatter(x=[i for i in range(len(ecg_inst.rr_intervals))], 
                         y=[rr_interval for rr_interval in ecg_inst.rr_intervals],
                         mode='lines', name='RR intervals'))

fig.add_trace(go.Scatter(x=rr_anomalies_simple,
                         y=[ecg_inst.rr_intervals[i] for i in rr_anomalies_simple],
                         mode='markers', name='Detected Anomalies'))

fig.add_trace(go.Scatter(x=[i for i in range(len(ecg_inst.rr_intervals))], 
                         y=rolling_mean,
                         mode='lines', name='Rolling mean'))

fig.add_trace(go.Scatter(x=[i for i in range(len(ecg_inst.rr_intervals))], 
                         y=rolling_mean + (0.2 * rolling_mean),
                         mode='lines', name='Rolling 20% mean upper'))


fig.add_trace(go.Scatter(x=[i for i in range(len(ecg_inst.rr_intervals))], 
                         y=rolling_mean - (0.2 * rolling_mean),
                         mode='lines', name='Rolling 20% mean lower'))

fig.update_xaxes(title_text='RR interval index')
fig.update_yaxes(title_text='RR interval duration (ms)')

fig.show()