"""Define report loader class."""
import warnings
import re
import bioc
import pandas as pd
from negbio.pipeline import text2bioc, ssplit, section_split

from constants import *


class Loader(object):
    """Report impression loader."""
    def __init__(self, reports_path, sections_to_extract, extract_strict):
        self.reports_path = reports_path
        self.sections_to_extract = sections_to_extract
        self.extract_strict = extract_strict
        self.punctuation_spacer = str.maketrans({key: f"{key} "
                                                 for key in ".,"})
        self.splitter = ssplit.NegBioSSplitter(newline=False)

    def load(self):
        """Load and clean the reports."""
        collection = bioc.BioCCollection()
        reports = pd.read_csv(self.reports_path,
                              header=None,
                              names=[REPORTS])[REPORTS].tolist()

        for i, report in enumerate(reports):
            clean_report = self.clean(report)
            document = text2bioc.text2document(str(i), clean_report)

            if self.sections_to_extract:
                document = self.extract_sections(document)

            split_document = self.splitter.split_doc(document)

            assert len(split_document.passages) == 1,\
                ('Each document must be given as a single passage.')

            collection.add_document(split_document)

        self.reports = reports
        self.collection = collection

    def extract_sections(self, document):
        """Extract the Impression section from a Bioc Document."""
        split_document = section_split.split_document(document)
        passages = []
        for i, passage in enumerate(split_document.passages):
            if 'title' in passage.infons:
                if (passage.infons['title'] in self.sections_to_extract and
                    len(split_document.passages) > i+1):
                    next_passage = split_document.passages[i+1]
                    if 'title' not in next_passage.infons:
                        passages.append(next_passage)
        
        if passages or self.extract_strict:
            extracted_passages = bioc.BioCPassage()
            if passages:
                extracted_passages.offset = passages[0].offset
                extracted_passages.text = ' '.join(map(lambda x: x.text, passages))
            else:
                extracted_passages.offset = 0
                extracted_passages.text = ''
            split_document.passages = [extracted_passages]
            return split_document
        else:
            warnings.warn('Loader found document containing none of the ' + 
                          'provided sections to extract. Returning original ' + 
                          'document.')
            return document

    def clean(self, report):
        """Clean the report text."""
        lower_report = report.lower()
        # Change `and/or` to `or`.
        corrected_report = re.sub('and/or',
                                  'or',
                                  lower_report)
        # Change any `XXX/YYY` to `XXX or YYY`.
        corrected_report = re.sub('(?<=[a-zA-Z])/(?=[a-zA-Z])',
                                  ' or ',
                                  corrected_report)
        # Clean double periods
        clean_report = corrected_report.replace("..", ".")
        # Insert space after commas and periods.
        clean_report = clean_report.translate(self.punctuation_spacer)
        # Convert any multi white spaces to single white spaces.
        clean_report = ' '.join(clean_report.split())
        # Remove empty sentences
        clean_report = re.sub(r'\.\s+\.', '.', clean_report)

        return clean_report