import sys, os

import glob

import itertools

import shutil

import time

from pathlib import Path

from nltk import sent_tokenize

from summarizers import Lexrank

import argparse

parser = argparse.ArgumentParser()

parser.add_argument('--train', action='store', metavar='path', type=str, required=True,  help='Directory path containing training documents. If test path not specified, also treated as testing documents.')

parser.add_argument('--saveto', action='store', metavar='path', type=str, required=True, help='directory path for saving summaries produced')

parser.add_argument('--test', action='store', metavar='path', type=str, help='directory path containing testing documents')

parser.add_argument('--ntrain', action='store', type=int, metavar='n', help='First n number of documents to train on')

parser.add_argument('--ntest', action='store', type=int, metavar='n', help='First n number of documents to produce summaries for')

parser.add_argument('--threshold', action='store', type=float, help="default 0.03")

parser.add_argument('--size', action='store', type=int, help='summary size. default 1')

args = parser.parse_args()

train_dir_path = args.train

test_dir_path = args.test

saveto_dir_path = args.saveto

threshold = args.threshold or 0.03

summary_size = args.size or 1

if not os.path.isdir(train_dir_path):

    print('The train path specified does not exist')

    sys.exit()

if test_dir_path and not os.path.isdir(test_dir_path):

    print('The test path specified does not exist')

    sys.exit()

if not os.path.isdir(saveto_dir_path):

    print('The save to path specified does not exist')

    sys.exit()

if args.ntest and args.ntest < 1:

    print('ntest should be greater than 0')

if args.ntrain and args.ntrain < 1:

    print('ntrain should be greater than 0')

start = time.time()

train_documents = {}

test_documents = {}

for i, filepath in enumerate(glob.glob(os.path.join(train_dir_path, '*'))):

    with open(filepath) as fp:

        fname = Path(filepath).stem

        sentences = [] 

        for line in fp.readlines():

            sentences.extend(sent_tokenize(line))

        train_documents[fname] = sentences

        if args.ntrain and i == args.ntrain - 1:

            break

if test_dir_path:

    for i, filepath in enumerate(glob.glob(os.path.join(test_dir_path, '*'))):

        with open(filepath) as fp:

            fname = Path(filepath).stem

            sentences = [] 

            for line in fp.readlines():

                sentences.extend(sent_tokenize(line))

            test_documents[fname] = sentences

        if args.ntest and i == args.ntest - 1:

            break

lxr = Lexrank(train_documents.values(), threshold=threshold)

if test_dir_path:

    documents = test_documents

else:

    documents = train_documents

for i, fname in enumerate(documents):

    summary = lxr.get_summary(documents[fname], summary_size=summary_size)#, threshold=.1)

    joined_summary = " ".join(summary)

    summary_path = os.path.join(saveto_dir_path, fname + ".sum")

    with open(summary_path, 'w') as sum:

        sum.write(joined_summary)

    if args.ntest and i == args.ntest - 1:

        break

end = time.time()

#print("----Summary----")

print("Runtime " + str(end - start))