text:
import numpy as np
import pandas as pd
import re
import os
import copy
from . import records
from . import _headers
from . import downloads

# Class for WFDB annotations
class Annotation(object):
    """
    The class representing WFDB annotations.

    Annotation objects can be created using the constructor, or by reading a WFDB annotation
    file with 'rdann'.

    The attributes of the Annotation object give information about the annotation as specified
    by https://www.physionet.org/physiotools/wag/annot-5.htm:
    - recordname: The base file name (without extension) of the record that the annotation
      is attached to.
    - extension: The file extension of the file the annotation is stored in.
    - sample: The annotation locations in samples relative to the beginning of the record.
    - symbol: The symbol used to display each annotation label.
    - subtype: The marked class/category of each annotation.
    - chan: The signal channel associated with each annotations.
    - num: The labelled annotation number for each annotation.
    - aux_note: The auxiliary information string for each annotation.
    - fs: The sampling frequency of the record, if available.
    - label_store: The integer value used to store/encode each annotation label
    - description: The descriptive string of each annotation label
    - custom_labels: The custom annotation labels defined in the annotation file.
      Maps the relationship between the three label fields.
      The data type is a pandas DataFrame with three columns: ['label_store', 'symbol', 'description']
    - contained_labels: The unique labels contained in this annotation. Same structure
      as custom_labels.

    Constructor function:
    def __init__(self, recordname, extension, sample, symbol=None, subtype=None,
                 chan=None, num=None, aux_note=None, fs=None, label_store=None,
                 description=None, custom_labels=None, contained_labels=None)

    Call 'show_ann_labels()' to see the list of standard annotation codes. Any text used to label
    annotations that are not one of these codes should go in the 'aux_note' field rather than the
    'sym' field.

    Example usage:
    import wfdb
    ann1 = wfdb.Annotation(recordname='rec1', extension='atr', sample=[10,20,400],
                           symbol = ['N','N','['], aux_note=[None, None, 'Serious Vfib'])
    """

    def __init__(self, recordname, extension, sample, symbol=None, subtype=None,
                 chan=None, num=None, aux_note=None, fs=None, label_store=None,
                 description=None, custom_labels=None, contained_labels=None):
        self.recordname = recordname
        self.extension = extension

        self.sample = sample
        self.symbol = symbol

        self.subtype = subtype
        self.chan = chan
        self.num = num
        self.aux_note = aux_note
        self.fs = fs

        self.label_store = label_store
        self.description = description

        self.custom_labels = custom_labels
        self.contained_labels = contained_labels

        #__label_map__: (storevalue, symbol, description) hidden attribute

    # Equal comparison operator for objects of this type
    def __eq__(self, other):
        att1 = self.__dict__
        att2 = other.__dict__

        if set(att1.keys()) != set(att2.keys()):
            print('keyset')
            return False

        for k in att1.keys():
            v1 = att1[k]
            v2 = att2[k]

            if type(v1) != type(v2):
                print(k)
                return False

            if isinstance(v1, np.ndarray):
                if not np.array_equal(v1, v2):
                    print(k)
                    return False
            elif isinstance(v1, pd.DataFrame):
                if not v1.equals(v2):
                    print(k)
                    return False
            else:
                if v1 != v2:
                    print(k)
                    return False

        return True

    # Write an annotation file
    def wrann(self, writefs=False):
        """
        Instance method to write a WFDB annotation file from an Annotation object.

        def wrann(self, writefs=False)

        Input Parameters:
        - writefs (default=False): Flag specifying whether to write the fs
          attribute to the file.

        """
        for field in ['recordname', 'extension']:
            if getattr(self, field) is None:
                raise Exception('Missing required field for writing annotation file: ',field)

        present_label_fields = self.get_label_fields()
        if not present_label_fields:
            raise Exception('At least one annotation label field is required to write the annotation: ', ann_label_fields)

        # Check the validity of individual fields
        self.checkfields()

        # Standardize the format of the custom_labels field
        self.standardize_custom_labels()

        # Create the label map used in this annotaion
        self.create_label_map()

        # Check the cohesion of fields
        self.check_field_cohesion(present_label_fields)

        # Calculate the label_store field if necessary
        if 'label_store' not in present_label_fields:
            self.convert_label_attribute(source_field=present_label_fields[0], target_field='label_store')

        # Write the header file using the specified fields
        self.wrannfile(writefs=writefs)

        return

    def get_label_fields(self):
        """
        Get the present label fields in the object
        """
        present_label_fields = []
        for field in ann_label_fields:
            if getattr(self, field) is not None:
                present_label_fields.append(field)

        return present_label_fields

    # Check the set fields of the annotation object
    def checkfields(self):
        # Check all set fields
        for field in ann_field_types:
            if getattr(self, field) is not None:
                # Check the type of the field's elements
                self.checkfield(field)
        return

    # Check a particular annotation field
    def checkfield(self, field):

        item = getattr(self, field)

        if not isinstance(item, ann_field_types[field]):
            raise TypeError('The '+field+' field must be one of the following types:', ann_field_types[field])

        if field in int_ann_fields:
            if item.dtype not in _headers.int_dtypes:
                raise TypeError('The '+field+' field must have an integer-based dtype.')

        # Field specific checks
        if field == 'recordname':
            if bool(re.search('[^-\w]', self.recordname)):
                raise ValueError('recordname must only comprise of letters, digits, hyphens, and underscores.')
        elif field == 'extension':
            if bool(re.search('[^a-zA-Z]', self.extension)):
                raise ValueError('extension must only comprise of letters.')
        elif field == 'fs':
            if self.fs <=0:
                raise ValueError('The fs field must be a non-negative number')
        elif field == 'custom_labels':
            """
            The role of this section is just to check the
            elements of this item, without utilizing
            any other fields. No format conversion
            or free value looksups etc are done.
            """

            # Check the structure of the subelements
            if isinstance(item, pd.DataFrame):
                column_names = list(item)
                if 'symbol' in column_names and 'description' in column_names:
                    if 'label_store' in column_names:
                        label_store = list(item['label_store'].values)
                    else:
                        label_store = None
                    symbol = item['symbol'].values
                    description = item['description'].values
                else:
                    raise ValueError(''.join(['If the '+field+' field is pandas dataframe, its columns',
                                             ' must be one of the following:\n-[label_store, symbol, description]',
                                             '\n-[symbol, description]']))
            else:
                if set([len(i) for i in item]) == {2}:
                    label_store = None
                    symbol = [i[0] for i in item]
                    description = [i[1] for i in item]
                elif set([len(i) for i in item]) == {3}:
                    label_store = [i[0] for i in item]
                    symbol = [i[1] for i in item]
                    description = [i[2] for i in item]
                else:
                    raise ValueError(''.join(['If the '+field+' field is an array-like object, its subelements',
                                             ' must be one of the following:\n- tuple triplets storing: ',
                                             '(label_store, symbol, description)\n- tuple pairs storing: ',
                                             '(symbol, description)']))

            # Check the values of the subelements
            if label_store:
                if len(item) != len(set(label_store)):
                    raise ValueError('The label_store values of the '+field+' field must be unique')

                if min(label_store) < 1 or max(label_store) > 49:
                    raise ValueError('The label_store values of the custom_labels field must be between 1 and 49')

            if len(item) != len(set(symbol)):
                raise ValueError('The symbol values of the '+field+' field must be unique')

            for i in range(len(item)):
                if label_store:
                    if not hasattr(label_store[i], '__index__'):
                        raise TypeError('The label_store values of the '+field+' field must be integer-like')

                if not isinstance(symbol[i], strtypes) or len(symbol[i]) not in [1,2,3]:
                    raise ValueError('The symbol values of the '+field+' field must be strings of length 1 to 3')

                if bool(re.search('[ \t\n\r\f\v]', symbol[i])):
                    raise ValueError('The symbol values of the '+field+' field must not contain whitespace characters')

                if not isinstance(description[i], strtypes):
                    raise TypeError('The description values of the '+field+' field must be strings')

                # Would be good to enfore this but existing garbage annotations have tabs and newlines...
                #if bool(re.search('[\t\n\r\f\v]', description[i])):
                #    raise ValueError('The description values of the '+field+' field must not contain tabs or newlines')

        # The string fields
        elif field in ['symbol', 'description', 'aux_note']:
            uniq_elements = set(item)

            for e in uniq_elements:
                if not isinstance(e, strtypes):
                    raise TypeError('Subelements of the '+field+' field must be strings')

            if field == 'symbol':
                for e in uniq_elements:
                    if len(e) not in [1,2,3]:
                        raise ValueError('Subelements of the '+field+' field must be strings of length 1 to 3')
                    if bool(re.search('[ \t\n\r\f\v]', e)):
                        raise ValueError('Subelements of the '+field+' field may not contain whitespace characters')
            else:
                for e in uniq_elements:
                    if bool(re.search('[\t\n\r\f\v]', e)):
                        raise ValueError('Subelements of the '+field+' field must not contain tabs or newlines')

        elif field == 'sample':
            if len(self.sample) == 1:
                sampdiffs = np.array([self.sample[0]])
            elif len(self.sample) > 1:
                sampdiffs = np.concatenate(([self.sample[0]], np.diff(self.sample)))
            else:
                raise ValueError("The 'sample' field must be a numpy array with length greater than 0")
            if min(self.sample) < 0 :
                raise ValueError("The 'sample' field must only contain non-negative integers")
            if min(sampdiffs) < 0 :
                raise ValueError("The 'sample' field must contain monotonically increasing sample numbers")
            if max(sampdiffs) > 2147483648:
                raise ValueError('WFDB annotation files cannot store sample differences greater than 2**31')

        elif field == 'label_store':
            if min(item) < 1 or max(item) > 49:
                raise ValueError('The label_store values must be between 1 and 49')

        # The C WFDB library stores num/sub/chan as chars.
        elif field == 'subtype':
            # signed character
            if min(self.subtype) < 0 or max(self.subtype) >127:
                raise ValueError("The 'subtype' field must only contain non-negative integers up to 127")
        elif field == 'chan':
            # unsigned character
            if min(self.chan) < 0 or max(self.chan) >255:
                raise ValueError("The 'chan' field must only contain non-negative integers up to 255")
        elif field == 'num':
            # signed character
            if min(self.num) < 0 or max(self.num) >127:
                raise ValueError("The 'num' field must only contain non-negative integers up to 127")

        return

    def check_field_cohesion(self, present_label_fields):
        """
        Check that the content and structure of different fields are consistent
        with one another.
        """
        # Ensure all written annotation fields have the same length
        nannots = len(self.sample)

        for field in ['sample', 'num', 'subtype', 'chan', 'aux_note']+present_label_fields:
            if getattr(self, field) is not None:
                if len(getattr(self, field)) != nannots:
                    raise ValueError("The lengths of the 'sample' and '"+field+"' fields do not match")

        # Ensure all label fields are defined by the label map. This has to be checked because
        # it is possible the user defined (or lack of) custom_labels does not capture all the
        # labels present.
        for field in present_label_fields:
            defined_values = self.__label_map__[field].values

            if set(getattr(self, field)) - set(defined_values) != set():
                raise ValueError('\n'.join(['\nThe '+field+' field contains elements not encoded in the stardard WFDB annotation labels, or this object\'s custom_labels field',
                                      '- To see the standard WFDB annotation labels, call: show_ann_labels()',
                                      '- To transfer non-encoded symbol items into the aux_note field, call: self.sym_to_aux()',
                                      '- To define custom labels, set the custom_labels field as a list of tuple triplets with format: (label_store, symbol, description)']))

        return

    def standardize_custom_labels(self):
        """
        Set the custom_labels field of the object to a standardized format:
        3 column pandas df with ann_label_fields as columns.

        Does nothing if there are no custom labels defined.
        Does nothing if custom_labels is already a df with all 3 columns

        If custom_labels is an iterable of pairs/triplets, this
        function will convert it into a df.

        If the label_store attribute is not already defined, this
        function will automatically choose values by trying to use:
        1. The undefined store values from the standard wfdb annotation
           label map.
        2. The unused label store values. This is extracted by finding the
           set of all labels contained in this annotation object and seeing
           which symbols/descriptions are not used.

        If there are more custom labels defined than there are enough spaces,
        even in condition 2 from above, this function will raise an error.

        This function must work when called as a standalone.
        """
        custom_labels = self.custom_labels

        if custom_labels is None:
            return

        self.checkfield('custom_labels')

        # Convert to dataframe if not already
        if not isinstance(custom_labels, pd.DataFrame):
            if len(self.custom_labels[0]) == 2:
                symbol = self.get_custom_label_attribute('symbol')
                description = self.get_custom_label_attribute('description')
                custom_labels = pd.DataFrame({'symbol': symbol, 'description': description})
            else:
                label_store = self.get_custom_label_attribute('label_store')
                symbol = self.get_custom_label_attribute('symbol')
                description = self.get_custom_label_attribute('description')
                custom_labels = pd.DataFrame({'label_store':label_store, 'symbol': symbol, 'description': description})

        # Assign label_store values to the custom labels if not defined
        if 'label_store' not in list(custom_labels):
            undefined_label_stores = self.get_undefined_label_stores()

            if len(custom_labels) > len(undefined_label_stores):
                available_label_stores = self.get_available_label_stores()
            else:
                available_label_stores = undefined_label_stores

            n_custom_labels = custom_labels.shape[0]

            if n_custom_labels > len(available_label_stores):
                raise ValueError('There are more custom_label definitions than storage values available for them.')

            custom_labels['label_store'] = available_label_stores[:n_custom_labels]

        custom_labels.set_index(custom_labels['label_store'].values, inplace=True)
        custom_labels = custom_labels[list(ann_label_fields)]

        self.custom_labels = custom_labels

        return

    def get_undefined_label_stores(self):
        """
        Get the label_store values not defined in the
        standard wfdb annotation labels.
        """
        return list(set(range(50)) - set(ann_label_table['label_store']))

    def get_available_label_stores(self, usefield = 'tryall'):
        """
        Get the label store values that may be used
        for writing this annotation.

        Available store values include:
        - the undefined values in the standard wfdb labels
        - the store values not used in the current
          annotation object.
        - the store values whose standard wfdb symbols/descriptions
          match those of the custom labels (if custom_labels exists)

        If 'usefield' is explicitly specified, the function will use that
        field to figure out available label stores. If 'usefield'
        is set to 'tryall', the function will choose one of the contained
        attributes by checking availability in the order: label_store, symbol, description
        """

        # Figure out which field to use to get available labels stores.
        if usefield == 'tryall':
            if self.label_store is not None:
                usefield = 'label_store'
            elif self.symbol is not None:
                usefield = 'symbol'
            elif self.description is not None:
                usefield = 'description'
            else:
                raise ValueError('No label fields are defined. At least one of the following is required: ', ann_label_fields)
            return self.get_available_label_stores(usefield = usefield)
        # Use the explicitly stated field to get available stores.
        else:
            # If usefield == 'label_store', there are slightly fewer/different steps
            # compared to if it were another option

            contained_field = getattr(self, usefield)

            # Get the unused label_store values
            if usefield == 'label_store':
                unused_label_stores = set(ann_label_table['label_store'].values) - contained_field
            else:
                # the label_store values from the standard wfdb annotation labels
                # whose symbols are not contained in this annotation
                unused_field = set(ann_label_table[usefield].values) - contained_field
                unused_label_stores = ann_label_table.loc[ann_label_table[usefield] in unused_field, 'label_store'].values

            # Get the standard wfdb label_store values overwritten by the
            # custom_labels if any
            if self.custom_symbols is not None:
                custom_field = set(self.get_custom_label_attribute(usefield))
                if usefield == 'label_store':
                    overwritten_label_stores = set(custom_field).intersection(set(ann_label_table['label_store']))
                else:
                    overwritten_fields = set(custom_field).intersection(set(ann_label_table[usefield]))
                    overwritten_label_stores = ann_label_table.loc[ann_label_table[usefield] in overwritten_fields, 'label_store'].values
            else:
                overwritten_label_stores = set()

            # The undefined values in the standard wfdb labels
            undefined_label_stores = self.get_undefined_label_stores()
            # Final available label stores = undefined + unused + overwritten
            available_label_stores = set(undefined_label_stores).union(set(unused_label_stores)).union(overwritten_label_stores)

            return available_label_stores

    def get_custom_label_attribute(self, attribute):
        """
        Get a list of the custom_labels attribute.
        ie. label_store, symbol, or description.

        The custom_labels variable could be in
        a number of formats
        """

        if attribute not in ann_label_fields:
            raise ValueError('Invalid attribute specified')

        if isinstance(self.custom_labels, pd.DataFrame):
            if 'label_store' not in list(self.custom_labels):
                raise ValueError('label_store not defined in custom_labels')
            a = list(self.custom_labels[attribute].values)
        else:
            if len(self.custom_labels[0]) == 2:
                if attribute == 'label_store':
                    raise ValueError('label_store not defined in custom_labels')
                elif attribute == 'symbol':
                    a = [l[0] for l in self.custom_labels]
                elif attribute == 'description':
                    a = [l[1] for l in self.custom_labels]
            else:
                if attribute == 'label_store':
                    a = [l[0] for l in self.custom_labels]
                elif attribute == 'symbol':
                    a = [l[1] for l in self.custom_labels]
                elif attribute == 'description':
                    a = [l[2] for l in self.custom_labels]

        return a

    def create_label_map(self, inplace=True):
        """
        Creates mapping df based on ann_label_table and self.custom_labels.

        Table composed of entire WFDB standard annotation table, overwritten/appended
        with custom_labels if any. Sets __label_map__ attribute, or returns value.
        """

        label_map =  ann_label_table.copy()

        if self.custom_labels is not None:
            self.standardize_custom_labels()
            for i in self.custom_labels.index:
                label_map.loc[i] = self.custom_labels.loc[i]

        if inplace:
            self.__label_map__ = label_map
        else:
            return label_map

    def wrannfile(self, writefs):
        """
        Calculate the bytes used to encode an annotation set and
        write them to an annotation file
        """

        # Calculate the fs bytes to write if present and desired to write
        if writefs:
            fs_bytes = self.calc_fs_bytes()
        else:
            fs_bytes = []
        # Calculate the custom_labels bytes to write if present
        cl_bytes = self.calc_cl_bytes()
        # Calculate the core field bytes to write
        core_bytes = self.calc_core_bytes()

        # Mark the end of the special annotation types if needed
        if fs_bytes == [] and cl_bytes == []:
            end_special_bytes = []
        else:
            end_special_bytes = [0, 236, 255, 255, 255, 255, 1, 0]

        # Write the file
        with open(self.recordname+'.'+self.extension, 'wb') as f:
            # Combine all bytes to write: fs (if any), custom annotations (if any), main content, file terminator
            np.concatenate((fs_bytes, cl_bytes, end_special_bytes, core_bytes, np.array([0,0]))).astype('u1').tofile(f)

        return

    # Calculate the bytes written to the annotation file for the fs field
    def calc_fs_bytes(self):

        if self.fs is None:
            return []

        # Initial indicators of encoding fs
        data_bytes = [0,88, 0, 252,35,35,32,116,105,109,101,32,114,101,115,111,108,117,116,105,111,110,58,32]

        # Check if fs is close enough to int
        if isinstance(self.fs, float):
            if round(self.fs,8) == float(int(self.fs)):
                self.fs = int(self.fs)

        fschars = str(self.fs)
        ndigits = len(fschars)

        for i in range(ndigits):
            data_bytes.append(ord(fschars[i]))

        # Fill in the aux_note length
        data_bytes[2] = ndigits + 20

        # odd number of digits
        if ndigits % 2:
            data_bytes.append(0)

        return np.array(data_bytes).astype('u1')

    # Calculate the bytes written to the annotation file for the custom_labels field
    def calc_cl_bytes(self):

        if self.custom_labels is None:
            return []

        # The start wrapper: '0 NOTE length aux_note ## annotation type definitions'
        headbytes = [0,88,30,252,35,35,32,97,110,110,111,116,97,116,105,111,110,32,116,
                     121,112,101,32,100,101,102,105,110,105,116,105,111,110,115]

        # The end wrapper: '0 NOTE length aux_note ## end of definitions' followed by SKIP -1, +1
        tailbytes =  [0,88,21,252,35,35,32,101,110,100,32,111,102,32,100,101,102,105,110,
                      105,116,105,111,110,115,0]

        custom_bytes = []

        for i in self.custom_labels.index:
            custom_bytes += custom_triplet_bytes(list(self.custom_labels.loc[i, list(ann_label_fields)]))

        # writecontent = []
        # for i in range(len(self.custom_labels)):
        #     writecontent.append([freenumbers[i],list(custom_labels.keys())[i],list(custom_labels.values())[i]])

        # custombytes = [customcode2bytes(triplet) for triplet in writecontent]
        # custombytes = [item for sublist in custombytes for item in sublist]

        return np.array(headbytes+custom_bytes+tailbytes).astype('u1')

    def calc_core_bytes(self):
        """
        Convert all used annotation fields into bytes to write
        """
        # The difference sample to write
        if len(self.sample) == 1:
            sampdiff = np.array([self.sample[0]])
        else:
            sampdiff = np.concatenate(([self.sample[0]], np.diff(self.sample)))

        # Create a copy of the annotation object with a
        # compact version of fields to write
        compact_annotation = copy.deepcopy(self)
        compact_annotation.compact_fields()

        # The optional fields to be written. Write if they are not None or all empty
        extra_write_fields = []

        for field in ['num', 'subtype', 'chan', 'aux_note']:
            if not isblank(getattr(compact_annotation, field)):
                extra_write_fields.append(field)

        data_bytes = []

        # Iterate across all fields one index at a time
        for i in range(len(sampdiff)):

            # Process the samp (difference) and sym items
            data_bytes.append(field2bytes('samptype', [sampdiff[i], self.symbol[i]]))

            # Process the extra optional fields
            for field in extra_write_fields:
                value = getattr(compact_annotation, field)[i]
                if value is not None:
                    data_bytes.append(field2bytes(field, value))

        # Flatten and convert to correct format
        data_bytes = np.array([item for sublist in data_bytes for item in sublist]).astype('u1')

        return data_bytes

    # Compact all of the object's fields so that the output
    # writing annotation file writes as few bytes as possible
    def compact_fields(self):

        # Number of annotations
        nannots = len(self.sample)

        # Chan and num carry over previous fields. Get lists of as few
        # elements to write as possible
        self.chan = compact_carry_field(self.chan)
        self.num = compact_carry_field(self.num)

        # Elements of 0 (default) do not need to be written for subtype.
        # num and sub are signed in original c package...
        if self.subtype is not None:
            if isinstance(self.subtype, list):
                for i in range(nannots):
                    if self.subtype[i] == 0:
                        self.subtype[i] = None
                if np.array_equal(self.subtype, [None]*nannots):
                    self.subtype = None
            else:
                zero_inds = np.where(self.subtype==0)[0]
                if len(zero_inds) == nannots:
                    self.subtype = None
                else:
                    self.subtype = list(self.subtype)
                    for i in zero_inds:
                        self.subtype[i] = None

        # Empty aux_note strings are not written
        if self.aux_note is not None:
            for i in range(nannots):
                if self.aux_note[i] == '':
                    self.aux_note[i] = None
            if np.array_equal(self.aux_note, [None]*nannots):
                self.aux_note = None

    def sym_to_aux(self):
        # Move non-encoded symbol elements into the aux_note field
        self.checkfield('symbol')

        # Non-encoded symbols
        label_table_map = self.create_label_map(inplace=False)
        external_syms = set(self.symbol) - set(label_table_map['symbol'].values)

        if external_syms == set():
            return

        if self.aux_note is None:
            self.aux_note = ['']*len(self.sample)

        for ext in external_syms:
            for i in [i for i,x in enumerate(self.symbol) if x == ext]:
                if not self.aux_note[i]:
                    self.aux_note[i] = self.symbol[i]
                else:
                    self.aux_note[i] = self.symbol[i]+' '+self.aux_note[i]
                self.symbol[i] = '"'
        return

    def get_contained_labels(self, inplace=True):
        """
        Get the set of unique labels contained in this annotation.
        Returns a pandas dataframe or sets the __contained__ labels
        attribute of the object.

        Requires the label_store field to be set.

        Function will try to use attributes contained in the order:
        1. label_store
        2. symbol
        3. description

        This function should also be called
        to summarize information about an
        annotation after it has been
        read. Should not be a helper function
        to others except rdann.
        """
        if self.custom_labels is not None:
            self.checkfield('custom_labels')

        # Create the label map
        label_map = ann_label_table.copy()

        # Convert the tuple triplets into a pandas dataframe if needed
        if isinstance(self.custom_labels, (list, tuple)):
            custom_labels = label_triplets_to_df(self.custom_labels)
        elif isinstance(self.custom_labels, pd.DataFrame):
            # Set the index just in case it doesn't already match the label_store
            self.custom_labels.set_index(self.custom_labels['label_store'].values, inplace=True)
            custom_labels = self.custom_labels
        else:
            custom_labels = None

        # Merge the standard wfdb labels with the custom labels.
        # custom labels values overwrite standard wfdb if overlap.
        if custom_labels is not None:
            for i in custom_labels.index:
                label_map.loc[i] = custom_labels.loc[i]
            # This doesn't work...
            # label_map.loc[custom_labels.index] = custom_labels.loc[custom_labels.index]

        # Get the labels using one of the features
        if self.label_store is not None:
            index_vals = set(self.label_store)
            reset_index = False
        elif self.symbol is not None:
            index_vals = set(self.symbol)
            label_map.set_index(label_map['symbol'].values, inplace=True)
            reset_index = True
        elif self.description is not None:
            index_vals = set(self.description)
            label_map.set_index(label_map['description'].values, inplace=True)
            reset_index = True
        else:
            raise Exception('No annotation labels contained in object')

        contained_labels = label_map.loc[index_vals, :]

        if reset_index:
            contained_labels.set_index(contained_labels['label_store'].values, inplace=True)

        if inplace:
            self.contained_labels = contained_labels
            return
        else:
            return contained_labels

    def set_label_elements(self, wanted_label_elements):
        """
        Set one or more label elements based on
        at least one of the others
        """
        if isinstance(wanted_label_elements, str):
            wanted_label_elements = [wanted_label_elements]

        # Figure out which desired label elements are missing
        missing_elements = [e for e in wanted_label_elements if getattr(self, e) is None]

        contained_elements = [e for e in ann_label_fields if getattr(self, e )is not None]

        if not contained_elements:
            raise Exception('No annotation labels contained in object')

        for e in missing_elements:
            self.convert_label_attribute(contained_elements[0], e)

        unwanted_label_elements = list(set(ann_label_fields) - set(wanted_label_elements))

        self.rm_attributes(unwanted_label_elements)

        return

    def rm_attributes(self, attributes):
        if isinstance(attributes, str):
            attributes = [attributes]
        for a in attributes:
            setattr(self, a, None)
        return

    def convert_label_attribute(self, source_field, target_field, inplace=True, overwrite=True):
        """
        Convert one label attribute (label_store, symbol, or description) to another.
        Input arguments:
        - inplace - If True, sets the object attribute. If False, returns the value.
        - overwrite - if True, performs conversion and replaces target field attribute even if the
          target attribute already has a value. If False, does not perform conversion in the aforementioned case.
          Set to True (do conversion) if inplace=False.

        Creates mapping df on the fly based on ann_label_table and self.custom_labels
        """
        if inplace and not overwrite:
            if getattr(self, target_field) is not None:
                return

        label_map = self.create_label_map(inplace=False)
        label_map.set_index(label_map[source_field].values, inplace=True)

        target_item = label_map.loc[getattr(self, source_field), target_field].values

        if target_field != 'label_store':
            # Should already be int64 dtype if target is label_store
            target_item = list(target_item)

        if inplace:
            setattr(self, target_field, target_item)
        else:
            return target_item

def label_triplets_to_df(triplets):
    """
    Get a pd dataframe from a tuple triplets
    used to define annotation labels.

    The triplets should come in the
    form: (label_store, symbol, description)
    """

    label_df = pd.DataFrame({'label_store':np.array([t[0] for t in triplets], dtype='int'),
                             'symbol':[t[1] for t in triplets],
                             'description':[t[2] for t in triplets]})

    label_df.set_index(label_df['label_store'].values, inplace=True)
    label_df = label_df[list(ann_label_fields)]

    return label_df

def custom_triplet_bytes(custom_triplet):
    """
    Convert triplet of [label_store, symbol, description] into bytes
    for defining custom labels in the annotation file
    """
    # Structure: 0, NOTE, len(aux_note), aux_note, codenumber, space, codesymbol, space, description, (0 null if necessary)
    # Remember, aux_note string includes 'number(s)<space><symbol><space><description>''
    annbytes = [0, 88, len(custom_triplet[2]) + 3 + len(str(custom_triplet[0])), 252] + [ord(c) for c in str(custom_triplet[0])] \
               + [32] + [ord(custom_triplet[1])] + [32] + [ord(c) for c in custom_triplet[2]]

    if len(annbytes) % 2:
        annbytes.append(0)

    return annbytes

# Tests whether the item is blank
def isblank(x):
    if x is None:
        return True
    elif isinstance(x, list):
        if set(x) == set([None]):
            return True
    return False

def compact_carry_field(full_field):
    """
    Return the compact list version of a list/array of an
    annotation field that has previous values carried over
    (chan or num)
    - The first sample is 0 by default. Only set otherwise
      if necessary.
    - Only set fields if they are different from their prev
      field
    """

    # Keep in mind that the field may already be compact or None

    if full_field is None:
        return None

    # List of same length. Place None where element
    # does not need to be written
    compact_field = [None]*len(full_field)

    prev_field = 0

    for i in range(len(full_field)):
        current_field = full_field[i]
        if current_field != prev_field:
            compact_field[i] = current_field
            prev_field = current_field

    # May further simplify
    if np.array_equal(compact_field, [None]*len(full_field)):
        compact_field = None

    return compact_field

# Convert an annotation field into bytes to write
def field2bytes(field, value):

    data_bytes = []

    # samp and sym bytes come together
    if field == 'samptype':
        # Numerical value encoding annotation symbol
        typecode = ann_label_table.loc[ann_label_table['symbol']==value[1], 'label_store'].values[0]

        # sample difference
        sd = value[0]

        # Add SKIP element if value is too large for single byte
        if sd>1023:
            # 8 bytes in total:
            # - [0, 59>>2] indicates SKIP
            # - Next 4 gives sample difference
            # - Final 2 give 0 and sym
            data_bytes = [0, 236, (sd&16711680)>>16, (sd&4278190080)>>24, sd&255, (sd&65280)>>8, 0, 4*typecode]
        # Just need samp and sym
        else:
            # - First byte stores low 8 bits of samp
            # - Second byte stores high 2 bits of samp
            #   and sym
            data_bytes = [sd & 255, ((sd & 768) >> 8) + 4*typecode]

    elif field == 'num':
        # First byte stores num
        # second byte stores 60*4 indicator
        data_bytes = [value, 240]
    elif field == 'subtype':
        # First byte stores subtype
        # second byte stores 61*4 indicator
        data_bytes = [value, 244]
    elif field == 'chan':
        # First byte stores num
        # second byte stores 62*4 indicator
        data_bytes = [value, 248]
    elif field == 'aux_note':
        # - First byte stores length of aux_note field
        # - Second byte stores 63*4 indicator
        # - Then store the aux_note string characters
        data_bytes = [len(value), 252] + [ord(i) for i in value]
        # Zero pad odd length aux_note strings
        if len(value) % 2:
            data_bytes.append(0)

    return data_bytes

# Function for writing annotations
def wrann(recordname, extension, sample, symbol=None, subtype=None, chan=None, num=None,
          aux_note=None, label_store=None, fs=None, custom_labels=None):
    """
    Write a WFDB annotation file.

    Specify at least the following:
    - The record name of the WFDB record (recordname)
    - The annotation file extension (extension)
    - The annotation locations in samples relative to the beginning of the record (sample)
    - Either the numerical values used to store the labels (label_store), or more commonly,
      the display symbols of each label (symbol).

    Usage:
    wrann(recordname, extension, sample, symbol=None, subtype=None, chan=None, num=None,
          aux_note=None, label_store=None, fs=None, custom_labels=None):

    Input arguments:
    - recordname (required): The string name of the WFDB record to be written (without any file extensions).
    - extension (required): The string annotation file extension.
    - sample (required): The annotation location in samples relative to the beginning of the record. Numpy array.
    - symbol (default=None): The symbols used to display the annotation labels. List or numpy array. If this field
      is present, 'label_store' must not be present.
    - subtype (default=None): The marked class/category of each annotation. Numpy array.
    - chan (default=None): The signal channel associated with each annotation. Numpy array.
    - num (default=None): The labelled annotation number of each annotation. Numpy array.
    - aux_note (default=None): The auxiliary information strings. List or numpy array.
    - label_store (default=None): The integer values used to store the annotation labels. Numpy array.
      If this field is present, 'symbol' must not be present.
    - fs (default=None): The numerical sampling frequency of the record to be written to the file.
    - custom_labels (default=None): The map of custom defined annotation labels used for this annotation, in
      addition to the standard WFDB annotation labels. The custom labels are defined by two or three fields:
      - The integer values used to store custom annotation labels in the file (optional)
      - Their short display symbols
      - Their long descriptions.
      This input argument may come in four formats:
      1. A pandas.DataFrame object with columns: ['label_store', 'symbol', 'description']
      2. A pandas.DataFrame object with columns: ['symbol', 'description']
         If this option is chosen, label_store values are automatically chosen.
      3. A list or tuple of tuple triplets, with triplet elements representing: (label_store, symbol, description).
      4. A list or tuple of tuple pairs, with pair elements representing: (symbol, description).
         If this option is chosen, label_store values are automatically chosen.
      If the 'label_store' field is given for this function, and 'custom_labels' is defined, 'custom_labels'
      must contain 'label_store' in its mapping. ie. it must come in format 1 or 3 above.

    Note: This gateway function was written to enable a simple way to write WFDB annotation files without
          needing to explicity create an Annotation object beforehand.

          You may also create an Annotation object, manually set its attributes, and call its wrann() instance method.

    Note: Each annotation stored in a WFDB annotation file contains a sample field and a label field. All other fields
          may or may not be present.

    Example Usage:
    import wfdb
    # Read an annotation as an Annotation object
    annotation = wfdb.rdann('b001', 'atr', pbdir='cebsdb')
    # Call the gateway wrann function, manually inserting fields as function input parameters
    wfdb.wrann('b001', 'cpy', annotation.sample, annotation.symbol)
    """

    # Create Annotation object
    annotation = Annotation(recordname=recordname, extension=extension, sample=sample, symbol=symbol,
                            subtype=subtype, chan=chan, num=num, aux_note=aux_note, label_store=label_store,
                            fs=fs, custom_labels=custom_labels)

    # Find out which input field describes the labels
    if symbol is None:
        if label_store is None:
            raise Exception("Either the 'symbol' field or the 'label_store' field must be set")
    else:
        if label_store is None:
            annotation.sym_to_aux()
        else:
            raise Exception("Only one of the 'symbol' and 'label_store' fields may be input, for describing annotation labels")

    # Perform field checks and write the annotation file
    annotation.wrann(writefs = True)

def show_ann_labels():
    """
    Display the standard wfdb annotation label mapping

    Usage:
    show_ann_labels()
    """
    print(ann_label_table)

def show_ann_classes():
    """
    Display the standard wfdb annotation classes

    Usage:
    show_ann_classes()
    """
    print(ann_class_table)

## ------------- Reading Annotations ------------- ##

# todo: return as df option?
def rdann(recordname, extension, sampfrom=0, sampto=None, shiftsamps=False,
          pbdir=None, return_label_elements=['symbol'], summarize_labels=False):
    """ Read a WFDB annotation file recordname.extension and return an
    Annotation object.

    Usage:
    annotation = rdann(recordname, extension, sampfrom=0, sampto=None, shiftsamps=False,
                 pbdir=None, return_label_elements=['symbol'], summarize_labels=False)

    Input arguments:
    - recordname (required): The record name of the WFDB annotation file. ie. for
      file '100.atr', recordname='100'
    - extension (required): The annotatator extension of the annotation file. ie. for
      file '100.atr', extension='atr'
    - sampfrom (default=0): The minimum sample number for annotations to be returned.
    - sampto (default=None): The maximum sample number for annotations to be returned.
    - shiftsamps (default=False): Boolean flag that specifies whether to return the
      sample indices relative to 'sampfrom' (True), or sample 0 (False). Annotation files
      store exact sample locations.
    - pbdir (default=None): Option used to stream data from Physiobank. The Physiobank database
      directory from which to find the required annotation file.
      eg. For record '100' in 'http://physionet.org/physiobank/database/mitdb', pbdir = 'mitdb'.
    - return_label_elements (default=['symbol']): The label elements that are to be returned
      from reading the annotation file. A list with at least one of the following: 'symbol',
      'label_store', 'description'.
    - summarize_labels (default=False): Assign a summary table of the set of annotation labels
      contained in the file to the 'contained_labels' attribute of the returned object.
      This table will contain the columns: ['label_store', 'symbol', 'description', 'n_occurences']

    Output argument:
    - annotation: The Annotation object. Call help(wfdb.Annotation) for the attribute
      descriptions.

    Note: For every annotation sample, the annotation file explictly stores the 'sample'
    and 'symbol' fields but not necessarily the others. When reading annotation files
    using this function, fields which are not stored in the file will either take their
    default values of 0 or None, or will be carried over from their previous values if any.

    Example usage:
    import wfdb
    ann = wfdb.rdann('sampledata/100', 'atr', sampto = 300000)
    """

    return_label_elements = check_read_inputs(sampfrom, sampto, return_label_elements)

    # Read the file in byte pairs
    filebytes = load_byte_pairs(recordname, extension, pbdir)

    # Get wfdb annotation fields from the file bytes
    sample, label_store, subtype, chan, num, aux_note = proc_ann_bytes(filebytes, sampto)

    # Get the indices of annotations that hold definition information about
    # the entire annotation file, and other empty annotations to be removed.
    potential_definition_inds, rm_inds = get_special_inds(sample, label_store, aux_note)

    # Try to extract information describing the annotation file
    fs, custom_labels = interpret_defintion_annotations(potential_definition_inds, aux_note)

    # Remove annotations that do not store actual sample and label information
    sample, label_store, subtype, chan, num, aux_note = rm_empty_indices(rm_inds, sample, label_store, subtype, chan, num, aux_note)

    # Convert lists to numpy arrays
    sample, label_store, subtype, chan, num= lists_to_arrays(sample, label_store, subtype, chan, num)

    # Obtain annotation sample relative to the starting signal index
    if shiftsamps and len(sample) > 0 and sampfrom:
        sample = sample - sampfrom

    # Try to get fs from the header file if it is not contained in the annotation file
    if fs is None:
        try:
            rec = records.rdheader(recordname, pbdir)
            fs = rec.fs
        except:
            pass

    # Create the annotation object
    annotation = Annotation(os.path.split(recordname)[1], extension, sample=sample, label_store=label_store,
                            subtype=subtype, chan=chan, num=num, aux_note=aux_note, fs=fs,
                            custom_labels=custom_labels)

    # Get the set of unique label definitions contained in this annotation
    if summarize_labels:
        annotation.get_contained_labels(inplace=True)

    # Set/unset the desired label values
    annotation.set_label_elements(return_label_elements)

    return annotation

def check_read_inputs(sampfrom, sampto, return_label_elements):

    if sampto and sampto <= sampfrom:
        raise ValueError("sampto must be greater than sampfrom")
    if sampfrom < 0:
        raise ValueError("sampfrom must be a non-negative integer")

    if isinstance(return_label_elements, str):
        return_label_elements = [return_label_elements]

    if set.union(set(ann_label_fields), set(return_label_elements))!=set(ann_label_fields):
        raise ValueError('return_label_elements must be a list containing one or more of the following elements:',label_types)

    return return_label_elements

# Load the annotation file 1 byte at a time and arrange in pairs
def load_byte_pairs(recordname, extension, pbdir):
    # local file
    if pbdir is None:
        with open(recordname + '.' + extension, 'rb') as f:
            filebytes = np.fromfile(f, '<u1').reshape([-1, 2])
    # physiobank file
    else:
        filebytes = downloads.streamannotation(recordname+'.'+extension, pbdir).reshape([-1, 2])

    return filebytes

#  Get regular annotation fields from the annotation bytes
def proc_ann_bytes(filebytes, sampto):

    # Base annotation fields
    sample, label_store, subtype, chan, num, aux_note = [], [], [], [], [], []

    # Indexing Variables
    # Total number of sample from beginning of record. Annotation bytes only store sample_diff
    sample_total = 0
    # Byte pair index
    bpi = 0

    # Process annotations. Iterate across byte pairs.
    # Sequence for one ann is:
    # - SKIP pair (if any)
    # - samp + sym pair
    # - other pairs (if any)
    # The last byte pair of the file is 0 indicating eof.
    while (bpi < filebytes.shape[0] - 1):

        # Get the sample and label_store fields of the current annotation
        sample_diff, current_label_store, bpi = proc_core_fields(filebytes, bpi)
        sample_total = sample_total + sample_diff
        sample.append(sample_total)
        label_store.append(current_label_store)

        # Process any other fields belonging to this annotation

        # Flags that specify whether the extra fields need to be updated
        update = {'subtype':True, 'chan':True, 'num':True, 'aux_note':True}
        # Get the next label store value - it may indicate additional
        # fields for this annotation, or the values of the next annotation.
        current_label_store = filebytes[bpi, 1] >> 2

        while (current_label_store > 59):
            subtype, chan, num, aux_note, update, bpi = proc_extra_field(current_label_store, filebytes,
                                                                         bpi, subtype, chan, num,
                                                                         aux_note, update)

            current_label_store = filebytes[bpi, 1] >> 2

        # Set defaults or carry over previous values if necessary
        subtype, chan, num, aux_note = update_extra_fields(subtype, chan, num, aux_note, update)

        if sampto and sampto<sample_total:
            sample, label_store, subtype, chan, num, aux_note = rm_last(sample, label_store, subtype, chan, num, aux_note)
            break

    return sample, label_store, subtype, chan, num, aux_note

# Get the sample difference and store fields of the current annotation
def proc_core_fields(filebytes, bpi):

    label_store = filebytes[bpi, 1] >> 2

    # The current byte pair will contain either the actual d_sample + annotation store value,
    # or 0 + SKIP.

    # Not a skip - it is the actual sample number + annotation type store value
    if label_store != 59:
        sample_diff = filebytes[bpi, 0] + 256 * (filebytes[bpi, 1] & 3)
        bpi = bpi + 1
    # Skip. Note: Could there be another skip after the first?
    else:
        # 4 bytes storing dt
        sample_diff = 65536 * filebytes[bpi + 1,0] + 16777216 * filebytes[bpi + 1,1] \
             + filebytes[bpi + 2,0] + 256 * filebytes[bpi + 2,1]

        # Data type is long integer (stored in two's complement). Range -2**31 to 2**31 - 1
        if sample_diff > 2147483647:
            sample_diff = sample_diff - 4294967296

        # After the 4 bytes, the next pair's samp is also added
        sample_diff = sample_diff + filebytes[bpi + 3, 0] + 256 * (filebytes[bpi + 3, 1] & 3)

        # The label is stored after the 4 bytes. Samples here should be 0.
        label_store = filebytes[bpi + 3, 1] >> 2
        bpi = bpi + 4

    return sample_diff, label_store, bpi

def proc_extra_field(label_store, filebytes, bpi, subtype, chan, num, aux_note, update):
    """
    Process extra fields belonging to the current annotation.
    Potential updated fields: subtype, chan, num, aux_note
    """

    # aux_note and sub are reset between annotations. chan and num copy over
    # previous value if missing.

    # SUB
    if label_store == 61:
        # sub is interpreted as signed char.
        subtype.append(filebytes[bpi, 0].astype('i1'))
        update['subtype'] = False
        bpi = bpi + 1
    # CHAN
    elif label_store == 62:
        # chan is interpreted as unsigned char
        chan.append(filebytes[bpi, 0])
        update['chan'] = False
        bpi = bpi + 1
    # NUM
    elif label_store == 60:
        # num is interpreted as signed char
        num.append(filebytes[bpi, 0].astype('i1'))
        update['num'] = False
        bpi = bpi + 1
    # aux_note
    elif label_store == 63:
        # length of aux_note string. Max 256? No need to check other bits of
        # second byte?
        aux_notelen = filebytes[bpi, 0]
        aux_notebytes = filebytes[bpi + 1:bpi + 1 + int(np.ceil(aux_notelen / 2.)),:].flatten()
        if aux_notelen & 1:
            aux_notebytes = aux_notebytes[:-1]
        # The aux_note string
        aux_note.append("".join([chr(char) for char in aux_notebytes]))
        update['aux_note'] = False
        bpi = bpi + 1 + int(np.ceil(aux_notelen / 2.))

    return subtype, chan, num, aux_note, update, bpi

def update_extra_fields(subtype, chan, num, aux_note, update):
    """
    Update the field if the current annotation did not
    provide a value.

    - aux_note and sub are set to default values if missing.
    - chan and num copy over previous value if missing.
    """

    if update['subtype']:
        subtype.append(0)

    if update['chan']:
        if chan == []:
            chan.append(0)
        else:
            chan.append(chan[-1])
    if update['num']:
        if num == []:
            num.append(0)
        else:
            num.append(num[-1])

    if update['aux_note']:
        aux_note.append('')

    return subtype, chan, num, aux_note

rx_fs = re.compile("## time resolution: (?P<fs>\d+\.?\d*)")
rx_custom_label = re.compile("(?P<label_store>\d+) (?P<symbol>\S+) (?P<description>.+)")

def get_special_inds(sample, label_store, aux_note):
    """
    Get the indices of annotations that hold definition information about
    the entire annotation file, and other empty annotations to be removed.

    Note: There is no need to deal with SKIP annotations (label_store=59)
          which were already dealt with in proc_core_fields and hence not
          included here.
    """

    s0_inds = np.where(sample == np.int64(0))[0]
    note_inds = np.where(label_store == np.int64(22))[0]

    # sample = 0 with aux_note means there should be an fs or custom label definition.
    # Either way, they are to be removed.
    potential_definition_inds = set(s0_inds).intersection(note_inds)

    # Other indices which are not actual annotations.
    notann_inds = np.where(label_store == np.int64(0))[0]

    rm_inds = potential_definition_inds.union(set(notann_inds))

    return potential_definition_inds, rm_inds

def interpret_defintion_annotations(potential_definition_inds, aux_note):
    """
    Try to extract annotation definition information from annotation notes.
    Information that may be contained:
    - fs - sample=0, label_state=22, aux_note='## time resolution: XXX'
    - custom annotation label definitions
    """

    fs = None
    custom_labels = []

    if len(potential_definition_inds) > 0:
        i = 0
        while i<len(potential_definition_inds):
            if aux_note[i].startswith('## '):
                if not fs:
                    search_fs = rx_fs.findall(aux_note[i])
                    if search_fs:
                        fs = float(search_fs[0])
                        if round(fs, 8) == float(int(fs)):
                            fs = int(fs)
                        i += 1
                        continue
                if aux_note[i] == '## annotation type definitions':
                    i += 1
                    while aux_note[i] != '## end of definitions':
                        label_store, symbol, description = rx_custom_label.findall(aux_note[i])[0]
                        custom_labels.append((int(label_store), symbol, description))
                        i += 1
                    i += 1
            else:
                i += 1

    if not custom_labels:
        custom_labels = None

    return fs, custom_labels

def rm_empty_indices(*args):
    """
    Remove unwanted list indices. First argument is the list
    of indices to remove. Other elements are the lists
    to trim.
    """
    rm_inds = args[0]

    if not rm_inds:
        return args[1:]

    keep_inds = [i for i in range(len(args[1])) if i not in rm_inds]

    return [[a[i] for i in keep_inds] for a in args[1:]]

def lists_to_arrays(*args):
    """
    Convert lists to numpy arrays
    """
    return [np.array(a, dtype='int') for a in args]

def rm_last(*args):
    """
    Remove the last index from each list
    """
    if len(args) == 1:
        return args[:-1]
    else:
        return [a[:-1] for a in args]
    return

## ------------- /Reading Annotations ------------- ##

# Allowed types of each Annotation object attribute.
ann_field_types = {'recordname': (str), 'extension': (str), 'sample': (np.ndarray),
                 'symbol': (list, np.ndarray),  'subtype': (np.ndarray), 'chan': (np.ndarray),
                 'num': (np.ndarray), 'aux_note': (list, np.ndarray), 'fs': _headers.floattypes,
                 'label_store': (np.ndarray), 'description':(list, np.ndarray), 'custom_labels': (pd.DataFrame, list, tuple),
                 'contained_labels':(pd.DataFrame, list, tuple)}

strtypes = (str, np.str_)

# Elements of the annotation label
ann_label_fields = ('label_store', 'symbol', 'description')

# Numerical integer annotation fields: sample, label_store, sub, chan, num
int_ann_fields = [field for field in ann_field_types if ann_field_types[field] == (np.ndarray)]

class AnnotationClass(object):
    def __init__(self, extension, description, human_reviewed):

        self.extension = extension
        self.description = description
        self.human_reviewed = human_reviewed

ann_classes = [
    AnnotationClass('atr', 'Reference ECG annotations', True),

    AnnotationClass('blh', 'Human reviewed beat labels', True),
    AnnotationClass('blm', 'Machine beat labels', False),

    AnnotationClass('alh', 'Human reviewed alarms', True),
    AnnotationClass('alm', 'Machine alarms', False),

    AnnotationClass('qrsc', 'Human reviewed qrs detections', True),
    AnnotationClass('qrs', 'Machine QRS detections', False),

    AnnotationClass('bph', 'Human reviewed BP beat detections', True),
    AnnotationClass('bpm', 'Machine BP beat detections', False),

    #AnnotationClass('alh', 'Human reviewed BP alarms', True),
    #AnnotationClass('alm', 'Machine BP alarms', False),
    # separate ecg and other signal category alarms?
    # Can we use signum to determine the channel it was triggered off?

    #ppg alarms?
    #eeg alarms?
]

ann_class_table = pd.DataFrame({'extension':[ac.extension for ac in ann_classes], 'description':[ac.description for ac in ann_classes],
                                 'human_reviewed':[ac.human_reviewed for ac in ann_classes]})
ann_class_table.set_index(ann_class_table['extension'].values, inplace=True)
ann_class_table = ann_class_table[['extension', 'description', 'human_reviewed']]

# Individual annotation labels
class AnnotationLabel(object):
    def __init__(self, label_store, symbol, short_description, description):
        self.label_store = label_store
        self.symbol = symbol
        self.short_description = short_description
        self.description = description

    def __str__(self):
        return str(self.label_store)+', '+str(self.symbol)+', '+str(self.short_description)+', '+str(self.description)

ann_labels = [
    AnnotationLabel(0, " ", 'NOTANN', 'Not an actual annotation'),
    AnnotationLabel(1, "N", 'NORMAL', 'Normal beat'),
    AnnotationLabel(2, "L", 'LBBB', 'Left bundle branch block beat'),
    AnnotationLabel(3, "R", 'RBBB', 'Right bundle branch block beat'),
    AnnotationLabel(4, "a", 'ABERR', 'Aberrated atrial premature beat'),
    AnnotationLabel(5, "V", 'PVC', 'Premature ventricular contraction'),
    AnnotationLabel(6, "F", 'FUSION', 'Fusion of ventricular and normal beat'),
    AnnotationLabel(7, "J", 'NPC', 'Nodal (junctional) premature beat'),
    AnnotationLabel(8, "A", 'APC', 'Atrial premature contraction'),
    AnnotationLabel(9, "S", 'SVPB', 'Premature or ectopic supraventricular beat'),
    AnnotationLabel(10, "E", 'VESC', 'Ventricular escape beat'),
    AnnotationLabel(11, "j", 'NESC', 'Nodal (junctional) escape beat'),
    AnnotationLabel(12, "/", 'PACE', 'Paced beat'),
    AnnotationLabel(13, "Q", 'UNKNOWN', 'Unclassifiable beat'),
    AnnotationLabel(14, "~", 'NOISE', 'Signal quality change'),
    # AnnotationLabel(15, None, None, None),
    AnnotationLabel(16, "|", 'ARFCT',  'Isolated QRS-like artifact'),
    # AnnotationLabel(17, None, None, None),
    AnnotationLabel(18, "s", 'STCH',  'ST change'),
    AnnotationLabel(19, "T", 'TCH',  'T-wave change'),
    AnnotationLabel(20, "*", 'SYSTOLE',  'Systole'),
    AnnotationLabel(21, "D", 'DIASTOLE',  'Diastole'),
    AnnotationLabel(22, '"', 'NOTE',  'Comment annotation'),
    AnnotationLabel(23, "=", 'MEASURE',  'Measurement annotation'),
    AnnotationLabel(24, "p", 'PWAVE',  'P-wave peak'),
    AnnotationLabel(25, "B", 'BBB',  'Left or right bundle branch block'),
    AnnotationLabel(26, "^", 'PACESP',  'Non-conducted pacer spike'),
    AnnotationLabel(27, "t", 'TWAVE',  'T-wave peak'),
    AnnotationLabel(28, "+", 'RHYTHM',  'Rhythm change'),
    AnnotationLabel(29, "u", 'UWAVE',  'U-wave peak'),
    AnnotationLabel(30, "?", 'LEARN',  'Learning'),
    AnnotationLabel(31, "!", 'FLWAV',  'Ventricular flutter wave'),
    AnnotationLabel(32, "[", 'VFON',  'Start of ventricular flutter/fibrillation'),
    AnnotationLabel(33, "]", 'VFOFF',  'End of ventricular flutter/fibrillation'),
    AnnotationLabel(34, "e", 'AESC',  'Atrial escape beat'),
    AnnotationLabel(35, "n", 'SVESC',  'Supraventricular escape beat'),
    AnnotationLabel(36, "@", 'LINK',  'Link to external data (aux_note contains URL)'),
    AnnotationLabel(37, "x", 'NAPC',  'Non-conducted P-wave (blocked APB)'),
    AnnotationLabel(38, "f", 'PFUS',  'Fusion of paced and normal beat'),
    AnnotationLabel(39, "(", 'WFON',  'Waveform onset'),
    AnnotationLabel(40, ")", 'WFOFF',  'Waveform end'),
    AnnotationLabel(41, "r", 'RONT',  'R-on-T premature ventricular contraction'),
    # AnnotationLabel(42, None, None, None),
    # AnnotationLabel(43, None, None, None),
    # AnnotationLabel(44, None, None, None),
    # AnnotationLabel(45, None, None, None),
    # AnnotationLabel(46, None, None, None),
    # AnnotationLabel(47, None, None, None),
    # AnnotationLabel(48, None, None, None),
    # AnnotationLabel(49, None, None, None),
]

ann_label_table = pd.DataFrame({'label_store':np.array([al.label_store for al in ann_labels], dtype='int'), 'symbol':[al.symbol for al in ann_labels],
                               'description':[al.description for al in ann_labels]})
ann_label_table.set_index(ann_label_table['label_store'].values, inplace=True)
ann_label_table = ann_label_table[['label_store','symbol','description']]


text:
import collections
import functools
import gzip
import itertools
import math
import numpy as np
import onmt.inputters
import os
import random
import torchtext

from scipy.special import softmax

from onmt.utils.logging import logger
from .vocab import Vocabulary
from .utils import UNDER

WARM_UP = 50000

class TransformModel():
    """ A model from which individual Transforms
    can be instantiated for each task """
    def __init__(self, data_config):
        self.data_config = data_config

    def warm_up(self, vocabs):
        pass

    def get_transform(self, transform, task):
        raise NotImplementedError()

class Transform():
    """ A preprocessing step, reapplied separately
    each time a minibatch is instantiated. """
    def __init__(self, data_config):
        pass

    def warm_up(self, vocabs):
        pass

    def set_train_opts(self, data_config):
        pass

    def get_specials(self):
        """ returns special tokens added by this transform """
        return set()

    def apply(self, tpl, task, is_train=True):
        return tpl

    def stats(self):
        return ()

    def __repr__(self):
        return self.__class__.__name__

class SimpleTransform(Transform):
    """ A Transform that is its own TransformModel """
    def __init__(self, data_config):
        pass

    def get_transform(self, transform, task):
        return self

class DuplicateMonoTransform(SimpleTransform):
    def apply(self, tpl, task, is_train=True):
        src = tuple(tpl[0])
        return src, src

class ReorderTransform(SimpleTransform):
    def __init__(self, data_config):
        super().__init__(data_config)

    def set_train_opts(self, data_config):
        self.max_dist = data_config['meta']['train'].get(
            'reorder_max_dist', 3)

    def _perturb(self, tokens):
        indices = [(i + random.uniform(0, self.max_dist), token)
                   for (i, token) in enumerate(tokens)]
        indices.sort()
        return tuple(token for (i, token) in indices)

    def apply(self, tpl, task, is_train=True):
        if not is_train:
            return tpl
        src, tgt = tpl
        src = self._perturb(src)
        return src, tgt

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.max_dist)

class DropTransform(SimpleTransform):
    def __init__(self, data_config):
        super().__init__(data_config)
        self._sum_draw = 0
        self._sum_toks = 0

    def set_train_opts(self, data_config):
        self.temperature = data_config['meta']['train'].get(
            'drop_temperature', 2.0) * -1

    def _drop(self, tokens):
        n_tokens = len(tokens)
        indices = np.arange(n_tokens)
        logits = indices * self.temperature
        logits = softmax(logits)
        n_drops = np.random.choice(indices, p=logits)
        self._sum_draw += n_drops
        self._sum_toks += n_tokens

        drop_indices = set(np.random.choice(indices,
                                            size=n_drops,
                                            replace=False))
        tokens = [tok for (i, tok) in enumerate(tokens)
                  if i not in drop_indices]
        return tokens

    def apply(self, tpl, task, is_train=True):
        if not is_train:
            return tpl
        src, tgt = tpl
        src = self._drop(src)
        return src, tgt

    def stats(self):
        if self._sum_toks == 0:
            yield('no tokens dropped')
        else:
            yield('tokens dropped {} / {} = {}'.format(
                self._sum_draw,
                self._sum_toks,
                self._sum_draw / self._sum_toks))

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.temperature)

class SwitchOutTransformModel(TransformModel):
    def __init__(self, data_config):
        super().__init__(data_config)
        self.vocab = None

    def warm_up(self, vocabs):
        # assumes shared vocab
        self.vocab = list(vocabs['shared'].keys())

    def get_transform(self, transform, task):
        return SwitchOutTransform(self.data_config, self.vocab)

class SwitchOutTransform(Transform):
    def __init__(self, data_config, vocab):
        super().__init__(data_config)
        self.vocab = vocab
        self._sum_draw = 0
        self._sum_toks = 0

    def set_train_opts(self, data_config):
        self.temperature = data_config['meta']['train'].get(
            'switchout_temperature', 2.0) * -1

    def _replace(self, token):
        return random.choice(self.vocab)

    def _switchout(self, tokens):
        n_tokens = len(tokens)
        indices = np.arange(n_tokens)
        logits = indices * self.temperature
        logits = softmax(logits)
        n_switchouts = np.random.choice(indices, p=logits)
        self._sum_draw += n_switchouts
        self._sum_toks += n_tokens

        switchout_indices = set(np.random.choice(indices,
                                                 size=n_switchouts,
                                                 replace=False))
        out = []
        for (i, tok) in enumerate(tokens):
            if i in switchout_indices:
                out.append(self._replace(tok))
            else:
                out.append(tok)
        return tuple(out)

    def apply(self, tpl, task, is_train=True):
        if not is_train:
            return tpl
        src, tgt = tpl
        src = self._switchout(src)
        tgt = self._switchout(tgt)
        return src, tgt

    def stats(self):
        if self._sum_toks == 0:
            yield('no tokens switched out')
        else:
            yield('switchout  {} / {} = {}'.format(
                self._sum_draw,
                self._sum_toks,
                self._sum_draw / self._sum_toks))

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.temperature)

class WbNoiseTransformModel(TransformModel):
    def __init__(self, data_config):
        super().__init__(data_config)
        self.vocab = None

    def warm_up(self, vocabs):
        # assumes shared vocab
        self.vocab = set(vocabs['shared'].keys())

    def get_transform(self, transform, task):
        return WbNoiseTransform(self.data_config, self.vocab)

class WbNoiseTransform(SwitchOutTransform):
    def set_train_opts(self, data_config):
        self.temperature = data_config['meta']['train'].get(
            'wb_noise_temperature', 2.0) * -1

    def _replace(self, token):
        if token[0] == UNDER:
            # wb -> no wb
            proposed = token[1:]
        else:
            # no wb -> wb
            proposed = UNDER + token
        if len(proposed) == 0:
            return token
        if proposed not in self.vocab:
            return token
        return proposed

    def stats(self):
        if self._sum_toks == 0:
            yield('no wb_noise')
        else:
            yield('wb_noise  {} / {} = {}'.format(
                self._sum_draw,
                self._sum_toks,
                self._sum_draw / self._sum_toks))

class InsertionTransformModel(TransformModel):
    def __init__(self, data_config):
        super().__init__(data_config)
        self.vocab = None

    def warm_up(self, vocabs):
        # assumes shared vocab
        self.vocab = list(vocabs['shared'].keys())

    def get_transform(self, transform, task):
        return InsertionTransform(self.data_config, self.vocab)

class InsertionTransform(SwitchOutTransform):
    def set_train_opts(self, data_config):
        self.temperature = data_config['meta']['train'].get(
            'insertion_temperature', 2.0) * -1

    def _switchout(self, tokens):
        n_tokens = len(tokens)
        indices = np.arange(n_tokens)
        logits = indices * self.temperature
        logits = softmax(logits)
        n_switchouts = np.random.choice(indices, p=logits)
        self._sum_toks += n_tokens

        switchout_indices = set(np.random.choice(indices,
                                                 size=n_switchouts,
                                                 replace=False))
        out = []
        for (i, tok) in enumerate(tokens):
            if i in switchout_indices:
                # first the insertion
                self._sum_draw += 1
                out.append(self._replace(tok))
            # then the old token
            out.append(tok)
        return tuple(out)

    def stats(self):
        if self._sum_toks == 0:
            yield('no insertion')
        else:
            yield('insertion  {} / {} = {}'.format(
                self._sum_draw,
                self._sum_toks,
                self._sum_draw / self._sum_toks))

class FilterTooLongTransform(SimpleTransform):
    def __init__(self, data_config):
        super().__init__(data_config)
        self._n_dropped = 0
        self._n_kept = 0

    def set_train_opts(self, data_config):
        self.max_len = data_config['meta']['train']['max_len']

    def apply(self, tpl, task, is_train=True):
        if not is_train:
            return tpl
        if any(len(side) > self.max_len for side in tpl):
            self._n_dropped += 1
            return None
        self._n_kept += 1
        return tpl

    def stats(self):
        tot = self._n_dropped + self._n_kept
        if tot <= 0:
            yield('nothing filtered')
        else:
            yield('length filtered {} / {} = {}'.format(
                self._n_dropped, tot, self._n_dropped / tot))

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.max_len)

class PrefixTransformModel(TransformModel):
    def __init__(self, data_config):
        super().__init__(data_config)

    def get_transform(self, transform, task):
        # FIXME can NOT override in set_train_opts
        src_lang = self.data_config['tasks'][task]['meta']['src_lang']
        tgt_lang = self.data_config['tasks'][task]['meta']['tgt_lang']
        prefix = ('<FROM_{}>'.format(src_lang),
                  '<TO_{}>'.format(tgt_lang))
        try:
            extra_prefix = \
                self.data_config['tasks'][task]['meta']['extra_prefix']
            prefix = prefix + (extra_prefix,)
        except KeyError:
            pass
        return PrefixTransform(prefix)

class PrefixTransform(Transform):
    def __init__(self, prefix):
        self.prefix = tuple(prefix)

    def get_specials(self):
        return set(self.prefix)

    def apply(self, tpl, task, is_train=True):
        # src is required, tgt is optional
        src, tail = tpl[0], tpl[1:]
        src = self.prefix + tuple(src)
        return (src,) + tail

    def __repr__(self):
        return '{}({})'.format(self.__class__.__name__, self.prefix)

def mix_half(xs, ys):
    assert len(xs) == len(ys)
    len_seq = len(xs)
    indices = list(range(len_seq))
    random.shuffle(indices)
    chosen = indices[:(len_seq // 2)]
    out_a = [x if i in chosen else y for (i, (x, y)) in enumerate(zip(xs, ys))]
    out_b = [y if i in chosen else x for (i, (x, y)) in enumerate(zip(xs, ys))]
    return out_a, out_b

def flatten(deep):
    out = []
    for token in deep:
        out.extend(token)
    return tuple(out)

def load_word_counts(data_config):
    vocab = Vocabulary(data_config)
    path = vocab.path('shared', segmentation='words')
    return vocab.load(path)

class MorfessorEmStdTransform(Transform):
    def __init__(self, data_config, seg_model, task):
        super().__init__(data_config)
        self.data_config = data_config
        self.seg_model = seg_model
        self.task = task
        # can NOT override in set_train_opts
        self.n_samples = data_config['meta']['train'].get(
            'seg_n_samples', 5)
        self.theta = data_config['meta']['train'].get(
            'seg_theta', 0.5)
        print('n_samples {} theta {}'.format(self.n_samples, self.theta))
        self._cache = StdSampleCache(
            self.seg_model,
            n_samples=self.n_samples,
            addcount=0,
            theta=self.theta,
            maxlen=30)

    def warm_up(self, vocabs):
        # load the word (not subword) counts
        counts = load_word_counts(self.data_config)
        # populate cache for most frequent words
        for w, c in counts.most_common(WARM_UP):
            self._cache._populate_cache(w)

    def apply(self, tpl, task, is_train=True):
        out = []
        for tokens in tpl:
            morphs = []
            for token in tokens:
                morphs.extend(self._cache.segment(token))
            out.append(tuple(morphs))
        return tuple(out)

    def stats(self):
        yield('hits {} vs misses {}'.format(self._cache.hits,
                                            self._cache.misses))

    def __repr__(self):
        try:
            return '{}({}, {})'.format(self.__class__.__name__,
                                       self.n_samples,
                                       self.theta)
        except AttributeError:
            return '{}[{}]'.format(self.__class__.__name__, 'old')

class MorfessorEmTabooTransform(Transform):
    def __init__(self, data_config, seg_model, task):
        super().__init__(data_config)
        self.data_config = data_config
        self.seg_model = seg_model
        self.task = task
        self.lang = self.data_config['tasks'][task]['meta']['src_lang']
        self.prefix = ('<TABOO_AE_{}>'.format(self.lang),)
        # can NOT override in set_train_opts
        self.n_samples = data_config['meta']['train'].get(
            'seg_n_samples', 5)
        self.theta = data_config['meta']['train'].get(
            'seg_theta', 0.5)
        print('n_samples {} theta {}'.format(self.n_samples, self.theta))
        self._cache = TabooSampleCache(
            self.seg_model,
            n_samples=self.n_samples,
            addcount=0,
            theta=self.theta,
            maxlen=30)

    def warm_up(self, vocabs):
        # load the word (not subword) counts
        counts = load_word_counts(self.data_config)
        # populate cache for most frequent words
        for w, c in counts.most_common(WARM_UP):
            self._cache._populate_cache(w)

    def get_specials(self):
        return set(self.prefix)

    def apply(self, tpl, task, is_train=True):
        # src -> xs, ys
        src, = tpl
        xs, ys = zip(*[self._cache.segment(token) for token in src])
        # xs, ys -mix-> src, tgt
        src, tgt = mix_half(xs, ys)
        src = flatten(src)
        tgt = flatten(tgt)
        # prepend '<TABOO_AE_{lang}>'
        src = self.prefix + src
        return src, tgt

    def stats(self):
        yield('hits {} vs misses {}'.format(self._cache.hits,
                                            self._cache.misses))

    def __repr__(self):
        try:
            return '{}({}, {})'.format(self.__class__.__name__,
                                       self.n_samples,
                                       self.theta)
        except AttributeError:
            return '{}[{}]'.format(self.__class__.__name__, 'old')

class MorfessorEmTabooReorderedTransform(MorfessorEmTabooTransform):
    """ taboo segmentation, after reordering src """
    def apply(self, tpl, task, is_train=True):
        # src -> xs, ys
        src, tgt = tpl
        all_morphs = sorted(set(src).union(tgt))
        xs, ys = zip(*[self._cache.segment(token) for token in all_morphs])
        # mix sides (per type)
        src_segs, tgt_segs = mix_half(xs, ys)
        # apply side-specific segmentation maps
        src_map = {morph: src_seg
                   for morph, src_seg
                   in zip(all_morphs, src_segs)}
        tgt_map = {morph: tgt_seg
                   for morph, tgt_seg
                   in zip(all_morphs, tgt_segs)}
        src = [src_map[token] for token in src]
        tgt = [tgt_map[token] for token in tgt]
        src = flatten(src)
        tgt = flatten(tgt)
        # prepend '<TABOO_AE_{lang}>'
        src = self.prefix + src
        return src, tgt

class MorfessorEmTransformModel():
    transform_classes = {
        'morfessor_em': MorfessorEmStdTransform,
        'morfessor_em_taboo': MorfessorEmTabooTransform,
        'morfessor_em_taboo_reordered': MorfessorEmTabooReorderedTransform,
    }

    def __init__(self, data_config):
        self.data_config = data_config

    def warm_up(self, vocabs):
        import morfessor
        # load the segmentation model
        io = morfessor.MorfessorIO()
        expected = io.read_expected_file(
            self.data_config['meta']['train']['segmentation_model'])
        self.seg_model = morfessor.BaselineModel(
            use_em=True,
            em_substr=expected)

    def get_transform(self, transform, task):
        try:
            transform_cls = self.transform_classes[transform]
        except KeyError:
            raise Exception('Unknown transform {}'.format(transform))
        return transform_cls(self.data_config, self.seg_model, task)

class SentencepieceTransform(SimpleTransform):
    def __init__(self, data_config):
        super().__init__(data_config)
        self.data_config = data_config
        self.model_path = \
            self.data_config['meta']['train']['segmentation_model']
        # can NOT override in set_train_opts
        self.n_samples = data_config['meta']['train'].get(
            'seg_n_samples', -1)
        self.theta = data_config['meta']['train'].get(
            'seg_theta', 0.5)
        if data_config['meta']['shard'].get('pretokenize', False):
            raise Exception(
                'SentencepieceTransform should not be used with "pretokenize"')
        if not data_config['meta']['shard'].get('predetokenize', False):
            logger.warn(
                'SentencepieceTransform used without "predetokenize". '
                'Make sure that the input is not tokenized.')

    def warm_up(self, vocabs=None):
        # load the segmentation model
        import sentencepiece as spm
        self.seg_model = spm.SentencePieceProcessor()
        self.seg_model.Load(self.model_path)

    def apply(self, tpl, task, is_train=True):
        out = []
        for tokens in tpl:
            out.append(tuple(self.seg_model.SampleEncodeAsPieces(
                ' '.join(tokens), self.n_samples, self.theta)))
        return tuple(out)

    def __getstate__(self):
        return {'data_config': self.data_config,
                'model_path': self.model_path,
                'n_samples': self.n_samples,
                'theta': self.theta}

    def __setstate__(self, d):
        self.data_config = d['data_config']
        self.model_path = d['model_path']
        self.n_samples = d['n_samples']
        self.theta = d['theta']
        self.warm_up()

class DeterministicSegmentationTransform(SimpleTransform):
    def __init__(self, data_config):
        super().__init__(data_config)
        self.data_config = data_config
        self.mapping = None
        self.mapping_path = \
            self.data_config['meta']['train']['segmentation_model']
        n_samples = data_config['meta']['train'].get(
            'seg_n_samples', 1)
        if n_samples != 1:
            raise Exception(
                'DeterministicSegmentationTransform requires seg_n_samples=1,'
                ' not {}'.format(n_samples))

    def warm_up(self, vocabs=None):
        # load the word to segmentation mapping
        # format: "<word> <tab> <morph1> <space> ... <morphN>"
        self.mapping = {}
        with open(self.mapping_path, 'r') as lines:
            for line in lines:
                line = line.strip()
                if len(line) == 0:
                    continue
                word, seg = line.split(None, 1)
                seg = tuple(seg.split(' '))
                self.mapping[word] = seg

    def apply(self, tpl, task, is_train=True):
        out = []
        for tokens in tpl:
            mapped = []
            for token in tokens:
                # intentionally die if token not in mapping
                mapped.extend(self.mapping[token])
            out.append(tuple(mapped))
        return tuple(out)

class SampleCache(object):
    def __init__(self, model, n_samples=5,
                 addcount=0, theta=0.5, maxlen=30):
        import morfessor
        self.model = model
        self.n_samples = n_samples
        self.addcount = addcount
        self.theta = theta
        self.maxlen = maxlen
        self.cache = {}
        self.hits = 0
        self.misses = 0
        self._categorical = morfessor.utils.categorical

    def segment(self, compound):
        if compound not in self.cache:
            self.misses += 1
            self._populate_cache(compound)
        else:
            self.hits += 1
        analyses, distr = self.cache[compound]
        if len(analyses) == 1:
            return analyses[0]
        _, sample = self._categorical(analyses, distr)
        return sample

class StdSampleCache(SampleCache):
    def _populate_cache(self, compound):
        n_best = self.model.viterbi_nbest(
            compound,
            self.n_samples,
            addcount=self.addcount,
            theta=self.theta,
            maxlen=self.maxlen)
        analyses, logps = zip(*n_best)
        distr = [math.exp(-x) for x in logps]
        divisor = sum(distr)
        if divisor == 0:
            distr = [1]
            analyses = [analyses[0]]
        else:
            distr = [x / divisor for x in distr]
        self.cache[compound] = (analyses, distr)

class TabooSampleCache(SampleCache):
    def _populate_cache(self, compound):
        samples = []
        logps = []
        # first side
        n_best = self.model.viterbi_nbest(
            compound,
            self.n_samples,
            addcount=self.addcount,
            theta=self.theta,
            maxlen=self.maxlen)

        for first, first_logp in n_best:
            taboo = [morph for morph in first if len(morph) > 1]
            second, second_logp = self.model.sample_segment(
                compound, theta=self.theta, maxlen=self.maxlen, taboo=taboo)
            logp = first_logp + second_logp
            samples.append((first, second))
            logps.append(logp)

        distr = [math.exp(-x) for x in logps]
        divisor = sum(distr)
        if divisor == 0:
            distr = [1]
            samples = [samples[0]]
        else:
            distr = [x / divisor for x in distr]
        self.cache[compound] = (samples, distr)

DEFAULT_TRANSFORMS = {
    'duplicate_mono': DuplicateMonoTransform,
    'reorder': ReorderTransform,
    'drop': DropTransform,
    'switchout': SwitchOutTransformModel,
    'wb_noise': WbNoiseTransformModel,
    'insertion': InsertionTransformModel,
    'lang_prefix_both': PrefixTransformModel,
    'morfessor_em': MorfessorEmTransformModel,
    'morfessor_em_taboo': MorfessorEmTransformModel,
    'morfessor_em_taboo_reordered': MorfessorEmTransformModel,
    'sentencepiece': SentencepieceTransform,
    'deterministic_segmentation': DeterministicSegmentationTransform,
    'filter_too_long': FilterTooLongTransform,
}

def make_transform_models(data_config, vocabs, custom_transforms=None):
    transform_model_classes = dict(DEFAULT_TRANSFORMS)
    if custom_transforms is not None:
        transform_model_classes.update(custom_transforms)
    transform_models = {}
    for key in data_config['_transforms']:
        transform_model = transform_model_classes[key](data_config)
        transform_model.warm_up(vocabs)
        transform_models[key] = transform_model
    return transform_models

def make_transforms(transform_models, data_config, vocabs):
    transforms = {}
    for task in data_config['tasks']:
        keys = data_config['tasks'][task]['transforms']
        transforms[task] = [transform_models[key].get_transform(key, task)
                            for key in keys]
        for transform in transforms[task]:
            transform.warm_up(vocabs)
    return transforms

def get_specials(transforms):
    # assumes shared vocabulary
    all_specials = set()
    for task_transforms in transforms.values():
        for transform in task_transforms:
            all_specials.update(transform.get_specials())
    return all_specials

def set_train_opts(data_config, transforms):
    for task_transforms in transforms.values():
        for transform in task_transforms:
            transform.set_train_opts(data_config)

ext: py

max_stars_repo_name: Waino/OpenNMT-py

max_issues_repo_name: MarcelRaschke/OpenNMT-py

max_forks_repo_name: MarcelRaschke/OpenNMT-py

reason: import numpy,from scipy


text:
#!/usr/bin/env python3

import numpy as np
import sys

def bump_neighbors(state, r, c):
    if r-1 >= 0:
        for ci in range(max(c-1, 0), min(c+2, state.shape[1])):
            state[ r-1, ci ] += 1

    for ci in range(max(c-1, 0), min(c+2, state.shape[1])):
        if ci == c: continue
        state[ r, ci ] += 1

    if r+1 < state.shape[1]:
        for ci in range(max(c-1, 0), min(c+2, state.shape[1])):
            state[ r+1, ci ] += 1

def evolve(state):
    #print("---- PRE  ----")
    #print(state)

    flashed = np.zeros_like(state, dtype=bool)

    # Energy of all octopi += 1
    state = state + np.ones_like(state)

    # Anybody with energy > 9 flashes; have to repeat this until it
    # converges (i.e., no new flashes)
    last_flashed = None
    num_flashed = 0
    rounds = 0
    while num_flashed != last_flashed:
        rounds += 1
        last_flashed = num_flashed
        for (r,c), val in np.ndenumerate(state):
            if val > 9:
                #print(f'FLASHED: {r}, {c} {flashed[r,c]}')
                if flashed[r, c] == False:
                    flashed[r, c] = True
                    num_flashed += 1
                    bump_neighbors(state, r, c)
    print(f'Converged with {num_flashed} flash(es) after {rounds} rounds')

    # Anybody that flashed gets their energy reset to 0
    state[flashed] = 0
    #print("---- POST ----")
    #print(state)

    return num_flashed, state

if __name__ == '__main__':
    if not len(sys.argv) == 3:
        print('USAGE: day11.py input_file iterations')
        sys.exit(1)

    steps = int(sys.argv[2])

    with open(sys.argv[1], 'r') as infile:
        lines = infile.readlines()
        lines = [ [int(x) for x in s.strip()] for s in lines ]
        state = np.array(lines)

        flashes = 0
        # Guess how much time I spent figuring out that my part two
        # answer was wrong because I was zero-indexing the steps?
        for i in range(1, steps+1):
            print(f'{i:4d}: ', end='')
            num_flashes, state = evolve(state)
            flashes += num_flashes
            if np.count_nonzero(state) == 0:
                print(f'Simultaneous megaflash after {i} steps')
                break

        print('Total flashes after {} step(s): {}'.format(i, flashes))
        #print('Final state:')
        #print(state)

reason: import numpy


text:
import numpy as np
import pandas as pd
import sys
import os
import pdb
import yaml
import tensorflow as tf
from classifier.DenseNet import pretrained_classifier
from utils import read_data_file, load_images_and_labels
import argparse
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore")

def train():
    parser = argparse.ArgumentParser()
    parser.add_argument(
    '--config', '-c', default='configs/celebA_YSBBB_Classifier.yaml'
)
    args = parser.parse_args()
    # ============= Load config =============
    config_path = args.config
    config = yaml.load(open(config_path))
    print(config)
    # ============= Experiment Folder=============
    output_dir = os.path.join(config['log_dir'], config['name'])
    try: os.makedirs(output_dir)
    except: pass
    try: os.makedirs(os.path.join(output_dir, 'logs'))
    except: pass
    # ============= Experiment Parameters =============
    BATCH_SIZE = config['batch_size']
    EPOCHS = config['epochs']
    channels = config['num_channel']
    input_size = config['input_size']
    N_CLASSES = config['num_class']
    ckpt_dir_continue = config['ckpt_dir_continue']
    if ckpt_dir_continue == '':
        continue_train = False
    else:
        continue_train = True
    # ============= Data =============
    try:
        categories, file_names_dict = read_data_file(config['image_label_dict'])
    except:
        print("Problem in reading input data file : ", config['image_label_dict'])
        sys.exit()
    data_train = np.load(config['train'])
    data_test = np.load(config['test'])
    print("The classification categories are: ")
    print(categories)
    print('The size of the training set: ', data_train.shape[0])
    print('The size of the testing set: ', data_test.shape[0])
    fp = open(os.path.join(output_dir, 'setting.txt'), 'w')
    fp.write('config_file:'+str(config_path)+'\n')
    fp.close()
    # ============= placeholder =============
    with tf.name_scope('input'):
        x_ = tf.placeholder(tf.float32, [None, input_size, input_size, channels], name='x-input')
        y_ = tf.placeholder(tf.int64, [None, N_CLASSES], name='y-input')
        isTrain = tf.placeholder(tf.bool)
    # ============= Model =============
    if N_CLASSES == 1:
        y = tf.reshape(y_, [-1])
        y = tf.one_hot(y,2,on_value=1.0,off_value=0.0,axis=-1)
        logit,prediction = pretrained_classifier(x_, n_label=2, reuse=False,  name='classifier', isTrain =isTrain)
    else:
        logit,prediction = pretrained_classifier(x_, n_label=N_CLASSES, reuse=False,  name='classifier', isTrain =isTrain)
        y = y_
    classif_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=y,logits=logit)
    loss = tf.losses.get_total_loss()
    # ============= Optimization functions =============
    train_step = tf.train.AdamOptimizer(0.0001).minimize(loss)
    # ============= summary =============
    cls_loss = tf.summary.scalar('classif_loss', classif_loss)
    total_loss = tf.summary.scalar('total_loss', loss)
    sum_train = tf.summary.merge([cls_loss, total_loss])
    # ============= Variables =============
    # Note that this list of variables only include the weights and biases in the model.
    lst_vars = []
    for v in tf.global_variables():
        lst_vars.append(v)
    # ============= Session =============
    sess=tf.InteractiveSession()
    saver = tf.train.Saver(var_list=lst_vars)
    tf.global_variables_initializer().run()
    writer = tf.summary.FileWriter(output_dir + '/train', sess.graph)
    writer_test = tf.summary.FileWriter(output_dir + '/test', sess.graph)
    # ============= Checkpoints =============
    if continue_train:
        print("Before training, Load checkpoint ")
        print("Reading checkpoint...")
        ckpt = tf.train.get_checkpoint_state(ckpt_dir_continue)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            saver.restore(sess, os.path.join(ckpt_dir_continue, ckpt_name))
            print(ckpt_name)
            print("Successful checkpoint upload")
        else:
            print("Failed checkpoint load")
            sys.exit()
    # ============= Training =============
    train_loss = []
    test_loss = []
    itr_train = 0
    itr_test = 0
    for epoch in range(EPOCHS):
        total_loss = 0.0
        perm = np.arange(data_train.shape[0])
        np.random.shuffle(perm)
        data_train = data_train[perm]
        num_batch = int(data_train.shape[0]/BATCH_SIZE)
        for i in range(0, num_batch):
            start = i*BATCH_SIZE
            ns = data_train[start:start+BATCH_SIZE]
            xs, ys = load_images_and_labels(ns, config['image_dir'],N_CLASSES, file_names_dict, input_size, channels, do_center_crop=True)
            [_, _loss,summary_str] = sess.run([train_step, loss, sum_train], feed_dict={x_:xs, isTrain:True, y_: ys})
            writer.add_summary(summary_str, itr_train)
            itr_train+=1
            total_loss += _loss
        total_loss /= i
        print("Epoch: " + str( epoch) + " loss: " + str(total_loss) + '\n')
        train_loss.append(total_loss)

        total_loss = 0.0
        perm = np.arange(data_test.shape[0])
        np.random.shuffle(perm)
        data_test = data_test[perm]
        num_batch = int(data_test.shape[0]/BATCH_SIZE)
        for i in range(0, num_batch):
            start = i*BATCH_SIZE
            ns = data_test[start:start+BATCH_SIZE]
            xs, ys = load_images_and_labels(ns, config['image_dir'],N_CLASSES, file_names_dict, input_size, channels, do_center_crop=True)
            [_loss, summary_str] = sess.run([loss, sum_train], feed_dict={x_:xs, isTrain:False, y_: ys})
            writer_test.add_summary(summary_str, itr_test)
            itr_test+=1
            total_loss += _loss
        total_loss /= i
        print("Epoch: "+ str(epoch) + " Test loss: "+ str(total_loss) + '\n')
        test_loss.append(total_loss)

        if epoch %2 == 0:
            checkpoint_name = os.path.join( output_dir, 'cp1_epoch'+str(epoch)+'.ckpt')
            save_path = saver.save(sess, checkpoint_name)
            np.save( os.path.join( output_dir,'logs','train_loss.npy'), np.asarray(train_loss))
            np.save( os.path.join( output_dir,'logs','test_loss.npy'), np.asarray(test_loss))
    checkpoint_name = os.path.join( output_dir, 'cp1_epoch'+str(epoch)+'.ckpt')
    save_path = saver.save(sess, checkpoint_name)
    np.save( os.path.join( output_dir,'logs','train_loss.npy'), np.asarray(train_loss))
    np.save( os.path.join( output_dir,'logs','test_loss.npy'), np.asarray(test_loss))

if __name__ == "__main__":
    train()