from django.core.management.base import BaseCommand
from ngram_viewer.models import *
from itertools import islice
from datetime import datetime
from tqdm import tqdm
import csv
import fnmatch
import os


class Command(BaseCommand):
    help = ("Adds n-grams to the database using the django models"
            " syntax. N-grams will be added from csv files with three columns."
            " First column is the n-gram string, second column is the key "
            " (e.g. year or speaker) and the third column is the counter."
            " Input (input_path) is a path pointing to one folder containing all"
            " 37 alphabetical sorted n-gram csv-files for one kind of n-gram."
            " Thus the user must specify with the parameter n_grams if the"
            " csv-files in the folder are 1-grams, 2-grams etc."
            " parameter 'n_grams'. The user also need to specifiy the corpus_type.")

    def add_arguments(self, parser):
        parser.add_argument("n_grams",
                            type=int,
                            choices=[1, 2, 3, 4, 5],
                            help="Tells the script to either import given input\
                            csv-files as 1-grams 2-grams etc.")
        parser.add_argument("input_folder",
                            type=str,
                            help="File path to the csv-files containing one    \
                            kind of ngrams.")
        parser.add_argument("corpus_type",
                            choices=["lm_ns_year", "tk_ws_year", "lm_ns_speaker",
                                     "tk_ws_speaker"],
                            help="User has to choose what kind of ngrams will  \
                            be imported. lm_ns_year: Lemmatized without        \
                            stopwords per year, tk_ws_year: not lemmatized     \
                            with stopwords per year etc.",
                            type=str)
        parser.add_argument(
                            "--batch_size",
                            "-bs",
                            type=int,
                            default=1000000,
                            required=False,
                            help="Int to set how many rows(entries) should be  \
                            inserted via bulk at once. Default is 1 million.   \
                            Optional parameter.")

    def handle(self, *args, **options):
        start_time = datetime.now()
        self.stdout.write("Start time of script is: " + str(start_time))
        folder_path = options["input_folder"]
        n_grams = options["n_grams"]
        corpus_type = options["corpus_type"]
        batch_size = options["batch_size"]

        list_of_files = []
        for path, subdirs, files in os.walk(folder_path):
            for name in files:
                if fnmatch.fnmatch(name, "*.csv"):
                    list_of_files.append(os.path.join(path, name))
        list_of_files = sorted(list_of_files)

        for file in tqdm(list_of_files, desc="File status"):
            with open(file, newline="") as csvfile:
                n_gram_reader = csv.reader(csvfile, delimiter="\t")

                row_count = sum(1 for row in n_gram_reader)  # closes csvfile
                iterations = int(row_count/batch_size) + 1
                self.stdout.write("Number of rows in csv is: " + str(row_count))
                self.stdout.write("Batch size is " + str(batch_size))
                self.stdout.write((str(iterations)
                                   + " iterations are needed to import the"
                                     " data into the database."))

            with open(file, newline="") as csvfile:  # reopens csvfile
                sort_key = os.path.basename(file)[0:1]
                if(sort_key == "_"):
                    sort_key = "_Non_ASCII"
                n_gram_reader = csv.reader(csvfile, delimiter="\t")

                if(n_grams == 1):
                    main_class = "One"
                elif(n_grams == 2):
                    main_class = "Two"
                elif(n_grams == 3):
                    main_class = "Three"
                elif(n_grams == 4):
                    main_class = "Four"
                elif(n_grams == 5):
                    main_class = "Five"

                model = "Key{}_{}Gram_{}".format(sort_key, main_class, corpus_type)
                print(model)
                while True:
                    batch = [globals()[model](ngram=row[0],
                                              key=row[1],
                                              count=row[2])
                             for row in tqdm(islice(n_gram_reader, batch_size),
                                             desc="Creating batch from row")]
                    if not batch:
                        break
                    self.stdout.write("Starting bulk insert.")
                    globals()[model].objects.bulk_create(batch, batch_size)
                    self.stdout.write("---------------------------------------")

        end_time = datetime.now()
        self.stdout.write("End time of script is: " + str(end_time))
        duration = end_time - start_time
        self.stdout.write("Duration of script is: " + str(duration))
