import re


def remove_punc(rawtext, ignore_case=False, remove_numbers=False, sub_numbers=True,
                names_list=set(), locations_list=set(), is_mhddata=False):
    """

    Parameters
    ----------
    rawtext : str
    ignore_case: bool
    remove_numbers: bool
        Removes all numbers if True
    sub_numbers: bool
        Substitutes all numbers to a token -num-.
        if 'remove_numbers' is True, it will have no effect.
    is_mhddata: bool
        True if the data is MHD
        False if it's from some other dataset

    Returns
    -------
    str

    """
    # from NLTK

    #ending quotes
    ENDING_QUOTES = [
        (re.compile(r'"'), " '' "),
        (re.compile(r'(\S)(\'\')'), r'\1 \2 '),
        (re.compile(r"([^' ])('[sS]|'[mM]|'[dD]|') "), r"\1 \2 "),
        (re.compile(r"([^' ])('ll|'LL|'re|'RE|'ve|'VE|n't|N'T) "), r"\1 \2 "),
    ]

    # List of contractions adapted from Robert MacIntyre's tokenizer.
    CONTRACTIONS2 = [re.compile(r"(?i)\b(can)(not)\b"),
                     re.compile(r"(?i)\b(d)('ye)\b"),
                     re.compile(r"(?i)\b(gim)(me)\b"),
                     re.compile(r"(?i)\b(gon)(na)\b"),
                     re.compile(r"(?i)\b(got)(ta)\b"),
                     re.compile(r"(?i)\b(lem)(me)\b"),
                     re.compile(r"(?i)\b(mor)('n)\b"),
                     re.compile(r"(?i)\b(wan)(na) ")]
    CONTRACTIONS3 = [re.compile(r"(?i) ('t)(is)\b"),
                     re.compile(r"(?i) ('t)(was)\b")]

    txt = re.sub(r'[\[\{\(\<]patient name[\]\}\)\>]', ' -name- ', rawtext)
    # For proper names and location.
    # Couldn't use string match since it might be part of other words therefore using regex
    # Takes forever!! we use dict at the end
    # for lc in locations_list:
    #     txt = re.sub(r"\b" + re.escape(lc) + r"\b", "-location-", txt)
    # for n in names_list:
    #     txt = re.sub(r"\b" + re.escape(n) + r"\b", "-name-", txt)

    # Laugh
    txt = re.sub(r'[\[\{\(\<]laugh[ter]*[\]\}\)\>]', ' -laugh- ', txt)

    # Delete brackets (also contents inside the brackets)
    txt = re.sub(r'[\[\{\<][^\[\{\<]*[\]\}\>]', ' ', txt)

    # Remove parentheses, while leaving the contents inside.
    txt = re.sub(r'[\(\)]', ' ', txt)

    if remove_numbers:
        # remove everything except for alpha characters and ', -, ?, !, .
        txt = re.sub(r'[^\.,A-Za-z\'\-\?\! ]', ' ', txt)
    else:
        # remove everything except for alphanumeric characters and ', -, ?, !, .
        txt = re.sub(r'[^\.,A-Za-z0-9\'\-\?\! ]', ' ', txt)
        # split numbers
        txt = re.sub(r"([0-9]+)", r" \1 ", txt)

        if sub_numbers:
            txt = re.sub(r'[0-9]+', '-num-', txt)

    if is_mhddata:
        # Replace non-ascii characters (for now just replace one)
        txt = re.sub("\x92", "'", txt)
        txt = re.sub("í", "'", txt)
        # Remove all NAs
        txt = re.sub(r"\bNA\b", "", txt)

    if ignore_case:
        txt = txt.lower()

    # space out the periods and commas
    txt = re.sub(r"\.", " . ", txt)
    txt = re.sub(r",", " , ", txt)
    # split ! and ?
    txt = re.sub(r'\!', ' ! ', txt)
    txt = re.sub(r'\?', ' ? ', txt)

    #add extra space to make things easier
    txt = " " + txt + " "

    # remove dashes that are used alone -- but this might be useful later
    # for now we're just removing them
    txt = re.sub(r' \-\-* ', ' ', txt)
    # remove -- these (consecutive dashes)
    txt = re.sub(r'\-\-+', ' ', txt)
    # remove dashes in the words that start with a dash or end with a dash
    # (Note: words that start AND end with single dashes are meaningful tokens.)
    txt = re.sub(r"\-([A-Za-z\']+)\s", r"\1 ", txt)
    txt = re.sub(r"\s([A-Za-z\']+)\-", r" \1", txt)

    #add extra space to make things easier
    txt = " " + txt + " "

    # Added these two to find a match with the pre-trained words
    # txt = re.sub(r'[A-Za-z]\-$', '<PARTIALWORD>', txt) # add this? probably not at the moment.
    # txt = re.sub(r'\-', ' ', txt)
    txt = re.sub(r"''", "'", txt)
    txt = re.sub(r"\s+'\s+", " ", txt)

    for regexp, substitution in ENDING_QUOTES:
        txt = regexp.sub(substitution, txt)

    for regexp in CONTRACTIONS2:
        txt = regexp.sub(r' \1 \2 ', txt)
    for regexp in CONTRACTIONS3:
        txt = regexp.sub(r' \1 \2 ', txt)

    txt = re.sub(r'\s+', ' ', txt) # make multiple white spaces into 1

    tokenized = txt.strip().split()  # tokenized temporarily to check proper nouns
    for i, w in enumerate(tokenized):
        if w in names_list:
            tokenized[i] = '-name-'
        # considers up to bigrams. First looks up if there is a match in bigrams. if not, checks unigram
        if i < len(tokenized)-1 and ' '.join(tokenized[i:i+2]) in locations_list:
            tokenized[i] = '-location-'
            del(tokenized[i+1])
        elif w in locations_list:
            tokenized[i] = '-location-'
    return ' '.join(tokenized)