# Copyright 2022 Cerebras Systems.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Script to format PubMed Fulltext commercial, PubMed Baseline and Update file Abstracts

Reference: https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow/LanguageModeling/BERT

"""

import csv
import glob
import os

import pubmed_parser as pmp


class TextFormatting:
    def __init__(
        self,
        pubmed_path,
        output_filename,
        filesize_limit=5 * (10 ** 9),
        recursive=False,
    ):
        """
        :param str pubmed_path: Path to folder containing PubMed files
        :param str output_folder : Path to where the txt file to be written
        :param Optional[int] filesize_limit: Max size of each text file
        :param Optional[bool] recursive: Flag if true, 
        searches for nxml/xml files recursively within subfolders
        """

        self.pubmed_path = pubmed_path

        print(f"self.pubmed_path:{pubmed_path}")
        self.recursive = recursive
        self.filesize = int(filesize_limit)

        self.output_folder = os.path.dirname(output_filename)
        if not os.path.exists(self.output_folder):
            os.makedirs(self.output_folder)

        self.filename = output_filename

    def merge_abstracts(self):

        file_num = 0
        num_articles = 0
        total_articles = 0

        output_filename = (
            self.filename.split('.')[0] + f"_{int(file_num)}" + ".txt"
        )

        csv_file = output_filename.split('.')[0] + "_stats.csv"
        csv_fh = open(csv_file, 'w')
        fieldnames = ['fname', 'num_articles']
        csv_writer = csv.DictWriter(csv_fh, fieldnames=fieldnames)
        csv_writer.writeheader()

        ofile = open(output_filename, mode='w', newline='\n')
        it = glob.iglob(self.pubmed_path + '/*.xml', recursive=self.recursive)

        for filename in it:
            print(f"Processing: {filename}")

            dicts_out = pmp.parse_medline_xml(filename)
            for dict_out in dicts_out:

                if not dict_out['abstract']:
                    # Some articles have no abstract : https://pubmed.ncbi.nlm.nih.gov/13787/
                    continue
                try:
                    for line in dict_out['abstract'].splitlines():
                        if len(line) < 30:
                            # Refer to https://pubmed.ncbi.nlm.nih.gov/4969/
                            # Multiple paragraphs in abstract with subtitles such as  "Result".
                            # Removing these subtitles ONLY
                            continue
                        ofile.write(line.strip() + " ")
                    ofile.write("\n\n")

                    num_articles += 1

                except:
                    ofile.write("\n\n")
                    continue

                if int(ofile.tell()) > self.filesize:
                    ofile.close()

                    # Write to csv stats:
                    csv_writer.writerow(
                        {
                            'fname': output_filename,
                            'num_articles': num_articles,
                        }
                    )
                    total_articles += num_articles

                    # Open another file
                    file_num += 1
                    output_filename = (
                        self.filename.split('.')[0]
                        + f"_{int(file_num)}"
                        + ".txt"
                    )
                    print(f" -- Creating new file: {output_filename}")
                    ofile = open(output_filename, mode='w', newline='\n')

                    # Reset abstracts count per file
                    num_articles = 0

        total_articles += num_articles
        csv_writer.writerow(
            {'fname': output_filename, 'num_articles': num_articles}
        )
        csv_writer.writerow(
            {'fname': 'Total abstracts', 'num_articles': total_articles}
        )
        csv_fh.close()
        ofile.close()
        print(f"**** Total number of abstracts = {total_articles}")

    def merge_fulltext(self):
        # This puts one article per line

        file_num = 0
        num_articles = 0
        total_articles = 0

        output_filename = (
            self.filename.split('.')[0] + f"_{int(file_num)}" + ".txt"
        )

        csv_file = output_filename.split('.')[0] + "_stats.csv"
        csv_fh = open(csv_file, 'w')
        fieldnames = ['fname', 'num_articles']
        csv_writer = csv.DictWriter(csv_fh, fieldnames=fieldnames)
        csv_writer.writeheader()

        top_level_folders = [
            os.path.join(self.pubmed_path, x)
            for x in os.listdir(self.pubmed_path)
        ]
        top_level_folders = [x for x in top_level_folders if os.path.isdir(x)]
        print(top_level_folders)

        not_written = os.path.join(self.output_folder, "exceptions.txt")

        with open(not_written, mode='w', newline='\n') as ex_fh:

            ofile = open(output_filename, mode='w', newline='\n')

            for folder in top_level_folders:

                it = glob.iglob(folder + '/**/*.nxml', recursive=self.recursive)

                for filename in it:
                    print(f"Processing: {filename}")
                    header_dict = pmp.parse_pubmed_xml(filename)
                    body_list = pmp.parse_pubmed_paragraph(
                        filename, all_paragraph=True
                    )

                    if not header_dict and not body_list:
                        ex_fh.write(filename)
                        ex_fh.write('\n')
                        continue

                    try:
                        if header_dict:
                            ofile.write(
                                header_dict['full_title'].strip() + ". "
                            )

                        if header_dict.get('abstract', None):
                            for line in header_dict['abstract'].splitlines():
                                if len(line) < 30:
                                    continue

                                ofile.write(line.strip() + " ")

                        if body_list:
                            for dict_entry in body_list:
                                section = dict_entry['section']

                                if len(section) > 30:
                                    ofile.write(section.strip() + ". ")

                                for line in dict_entry['text'].splitlines():
                                    ofile.write(line.strip() + " ")

                            ofile.write("\n\n")

                        num_articles += 1

                    except:
                        ofile.write("\n\n")
                        continue

                    if int(ofile.tell()) > self.filesize:
                        ofile.close()

                        # Write to csv stats:
                        csv_writer.writerow(
                            {
                                'fname': output_filename,
                                'num_articles': num_articles,
                            }
                        )
                        total_articles += num_articles

                        # Open another file
                        file_num += 1
                        output_filename = (
                            self.filename.split('.')[0]
                            + f"_{int(file_num)}"
                            + ".txt"
                        )

                        print(f" -- Creating new file: {output_filename}")
                        ofile = open(output_filename, mode='w', newline='\n')

                        # Reset articles count
                        num_articles = 0

                total_articles += num_articles
                csv_writer.writerow(
                    {'fname': output_filename, 'num_articles': num_articles,}
                )
                csv_writer.writerow(
                    {
                        'fname': 'Total num articles',
                        'num_articles': total_articles,
                    }
                )
                csv_fh.close()
                ofile.close()
                print(
                    f"**** Total number of full text articles = {total_articles}"
                )

    def merge(self, dataset_name):

        if (
            dataset_name == "pubmed_baseline"
            or dataset_name == "pubmed_daily_update"
        ):
            self.merge_abstracts()

        elif (
            dataset_name == "pubmed_fulltext"
            or dataset_name == "pubmed_open_access"
        ):
            self.merge_fulltext()

        else:
            raise ValueError(f"Incorrect dataset_name: {dataset_name} passed")
