import plotly.graph_objects as go
import streamlit as st
import pandas as pd
from utils import *
import glob
import wfdb
import os

ANNOTATIONS_COL_NAME = 'annotations'

'''
# MIT-BIH Arrhythmia DB Exploration
'''

record_ids = [os.path.basename(file)[:-4] for file in glob.glob('data/*.dat')]

if len(record_ids) == 0:
    st.write('Warning ! No data could be found under the ./data/ directory.',
             '*\*.dat*, *\*.hea*, *\*.atr* files and such should be placed ',
             'immediately under the ./data/ directory')
else:
    record_ids.sort()
    record_id = st.selectbox('Select a record id', record_ids)
    record = wfdb.rdrecord(f'data/{record_id}')
    annotation = wfdb.rdann(f'data/{record_id}', 'atr')
    st.write('Signals found in this record :')
    for idx, signal in enumerate(record.sig_name):
        st.write(f'- `{signal}` : in {record.units[idx]}, with a frequency of '
                 f'{record.fs * record.samps_per_frame[idx]}hz')
    st.write(f'Comments for this record : {record.comments}')
    signals_df = pd.DataFrame(record.p_signal, columns=record.sig_name)
    annot_serie = pd.Series(annotation.symbol, index=annotation.sample,
                            name=ANNOTATIONS_COL_NAME)
    full_df = pd.concat([signals_df, annot_serie], axis=1)

    ''' ## Annotations '''
    beat_annot_count = annot_serie.isin(dict(beat_annotations)).sum()
    non_beat_annot_count = annot_serie.isin(dict(non_beat_annotations)).sum()
    unique_annot = annot_serie.value_counts().index.values
    st.write(f'This record contains `{annot_serie.size}` annotations '
             f'among which `{beat_annot_count}` beat annotations and '
             f'`{non_beat_annot_count}` non beat annotation(s).')
    st.write('The annotations are the followings :')
    for annot in unique_annot:
        st.write(f'- `{annot}` : {annotation_definitions[annot]}')
    st.write('More explanations on the annotations are available here : '
             'https://archive.physionet.org/physiobank/annotations.shtml')

    # Plot counts for each annotation
    annot_counts_df = annot_serie \
        .value_counts() \
        .rename_axis(ANNOTATIONS_COL_NAME) \
        .reset_index(name='counts')
    bar_fig = go.Figure(data=[go.Bar(x=annot_counts_df[ANNOTATIONS_COL_NAME],
                                     y=annot_counts_df['counts'],
                                     text=annot_counts_df['counts'],
                                     textposition='auto'
                                     )])
    bar_fig.update_layout(title='Annotations by count', yaxis_title='counts',
                          xaxis_title='annotations')
    st.write(bar_fig)

    ''' ## Explore full dataset '''
    signal = st.selectbox('Select a signal', record.sig_name)
    # Plot signals and annotations
    matching_rows_by_annot = {}
    for annot in unique_annot:
        matching_rows_by_annot[annot] = full_df[ANNOTATIONS_COL_NAME] == annot
    fig = go.Figure(layout=go.Layout(title=go.layout.Title(
        text='{} signal with annotations'.format(signal))))
    fig.add_trace(go.Scatter(x=full_df.index.values,
                             y=full_df[signal],
                             mode='lines',
                             name=signal))
    for annot, annot_matching_rows in matching_rows_by_annot.items():
        fig.add_trace(go.Scatter(x=full_df.index[annot_matching_rows].values,
                                 y=full_df[annot_matching_rows][signal].values,
                                 mode='markers',
                                 name='{} (annot)'.format(annot)))
    st.plotly_chart(fig)

[((18, 4, 20, 55), 'streamlit.write', 'st.write', ({(18, 13, 18, 76): '"""Warning ! No data could be found under the ./data/ directory."""', (19, 13, 19, 76): '"""*\\\\*.dat*, *\\\\*.hea*, *\\\\*.atr* files and such should be placed """', (20, 13, 20, 54): '"""immediately under the ./data/ directory"""'}, {}), "('Warning ! No data could be found under the ./data/ directory.',\n    '*\\\\*.dat*, *\\\\*.hea*, *\\\\*.atr* files and such should be placed ',\n    'immediately under the ./data/ directory')", True, 'import streamlit as st\n'), ((23, 16, 23, 62), 'streamlit.selectbox', 'st.selectbox', ({(23, 29, 23, 49): '"""Select a record id"""', (23, 51, 23, 61): 'record_ids'}, {}), "('Select a record id', record_ids)", True, 'import streamlit as st\n'), ((24, 13, 24, 47), 'wfdb.rdrecord', 'wfdb.rdrecord', ({(24, 27, 24, 46): 'f"""data/{record_id}"""'}, {}), "(f'data/{record_id}')", False, 'import wfdb\n'), ((25, 17, 25, 55), 'wfdb.rdann', 'wfdb.rdann', ({(25, 28, 25, 47): 'f"""data/{record_id}"""', (25, 49, 25, 54): '"""atr"""'}, {}), "(f'data/{record_id}', 'atr')", False, 'import wfdb\n'), ((26, 4, 26, 46), 'streamlit.write', 'st.write', ({(26, 13, 26, 45): '"""Signals found in this record :"""'}, {}), "('Signals found in this record :')", True, 'import streamlit as st\n'), ((30, 4, 30, 61), 'streamlit.write', 'st.write', ({(30, 13, 30, 60): 'f"""Comments for this record : {record.comments}"""'}, {}), "(f'Comments for this record : {record.comments}')", True, 'import streamlit as st\n'), ((31, 17, 31, 71), 'pandas.DataFrame', 'pd.DataFrame', (), '', True, 'import pandas as pd\n'), ((32, 18, 33, 54), 'pandas.Series', 'pd.Series', (), '', True, 'import pandas as pd\n'), ((34, 14, 34, 58), 'pandas.concat', 'pd.concat', (), '', True, 'import pandas as pd\n'), ((40, 4, 42, 65), 'streamlit.write', 'st.write', ({(40, 13, 42, 64): 'f"""This record contains `{annot_serie.size}` annotations among which `{beat_annot_count}` beat annotations and `{non_beat_annot_count}` non beat annotation(s)."""'}, {}), "(\n    f'This record contains `{annot_serie.size}` annotations among which `{beat_annot_count}` beat annotations and `{non_beat_annot_count}` non beat annotation(s).'\n    )", True, 'import streamlit as st\n'), ((43, 4, 43, 52), 'streamlit.write', 'st.write', ({(43, 13, 43, 51): '"""The annotations are the followings :"""'}, {}), "('The annotations are the followings :')", True, 'import streamlit as st\n'), ((46, 4, 47, 74), 'streamlit.write', 'st.write', ({(46, 13, 47, 73): '"""More explanations on the annotations are available here : https://archive.physionet.org/physiobank/annotations.shtml"""'}, {}), "(\n    'More explanations on the annotations are available here : https://archive.physionet.org/physiobank/annotations.shtml'\n    )", True, 'import streamlit as st\n'), ((61, 4, 61, 21), 'streamlit.write', 'st.write', ({(61, 13, 61, 20): 'bar_fig'}, {}), '(bar_fig)', True, 'import streamlit as st\n'), ((64, 13, 64, 61), 'streamlit.selectbox', 'st.selectbox', ({(64, 26, 64, 43): '"""Select a signal"""', (64, 45, 64, 60): 'record.sig_name'}, {}), "('Select a signal', record.sig_name)", True, 'import streamlit as st\n'), ((80, 4, 80, 24), 'streamlit.plotly_chart', 'st.plotly_chart', ({(80, 20, 80, 23): 'fig'}, {}), '(fig)', True, 'import streamlit as st\n'), ((15, 14, 15, 36), 'os.path.basename', 'os.path.basename', ({(15, 31, 15, 35): 'file'}, {}), '(file)', False, 'import os\n'), ((15, 54, 15, 77), 'glob.glob', 'glob.glob', ({(15, 64, 15, 76): '"""data/*.dat"""'}, {}), "('data/*.dat')", False, 'import glob\n'), ((28, 8, 29, 64), 'streamlit.write', 'st.write', ({(28, 17, 29, 63): 'f"""- `{signal}` : in {record.units[idx]}, with a frequency of {record.fs * record.samps_per_frame[idx]}hz"""'}, {}), "(\n    f'- `{signal}` : in {record.units[idx]}, with a frequency of {record.fs * record.samps_per_frame[idx]}hz'\n    )", True, 'import streamlit as st\n'), ((45, 8, 45, 66), 'streamlit.write', 'st.write', ({(45, 17, 45, 65): 'f"""- `{annot}` : {annotation_definitions[annot]}"""'}, {}), "(f'- `{annot}` : {annotation_definitions[annot]}')", True, 'import streamlit as st\n'), ((71, 18, 74, 41), 'plotly.graph_objects.Scatter', 'go.Scatter', (), '', True, 'import plotly.graph_objects as go\n'), ((54, 30, 58, 38), 'plotly.graph_objects.Bar', 'go.Bar', (), '', True, 'import plotly.graph_objects as go\n')]