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()

fig = go.Figure()

fig.add_trace(go.Scatter(x=[datetime.datetime.fromtimestamp(seconds).time() for seconds in ecg_inst.time], 
                         y=ecg_inst.signal,
                         mode='lines', name='raw signal'))

fig.add_trace(go.Scatter(x=[datetime.datetime.fromtimestamp(ecg_inst.r_peaks[i][0]).time() 
                            for i in rr_anomalies_simple], 
                         y=[ecg_inst.r_peaks[i][1] for i in rr_anomalies_simple],
                         mode='markers', marker_size=9, name='anomalies'))

fig.add_trace(go.Scatter(x=[datetime.datetime.fromtimestamp(ann_anomaly[0] / ecg_inst.samp_freq).time() 
                            for ann_anomaly in ecg_ann_anomalies], 
                         y=[ecg_inst.signal[ann_anomaly[0]] 
                            for ann_anomaly in ecg_ann_anomalies],
                         mode='markers', marker_symbol='x',
                         marker_color='orange', name='annotations'))


fig.update_xaxes(title_text='Time (s)')
fig.update_yaxes(title_text='ECG Signal (mV)')

fig.show()