"""Word Co-occurence finder tool
"""
import os
import sys
import mmap
import json
import typing as tg
from multiprocessing import Pool
import csv
import pdb
import re
from collections import defaultdict

def word_count(args: tg.Tuple[tg.List[str], tg.Dict[str, tg.List]]) -> tg.Tuple[str, tg.Dict[str, str]]:
    """function that count the number of defined keys in given text
    using tuple to store args due to Pool's map method only support one parameter
    """
    if DEBUG:
        #print(args)
        pass
    
    row, raw_key = args

    id_, string = row
    
    result = defaultdict()
    
    for entity, entity_key in raw_key.items():
        is_match = False
        for key_type, key_list in entity_key.items():
            if key_type in ("abbreviation", "regular_expression"):
                for key in key_list:
                    match = re.search(key, string, re.UNICODE)
                    if match:
                        is_match = True
                        break
            elif type in ("full_name",):
                for key in key_list:
                    match = re.search(key, string, re.UNICODE + re.IGNORECASE)
                    if match:
                        is_match = True
                        break
        if is_match:
            result[entity] = match.group()

    return id_, result # tg.Dict[str, str]
    
DEBUG = True

if __name__ == "__main__":
    if DEBUG:
        text_file = "TEST.tsv"
        with open("TEST.json", 'r', encoding='utf_8') as keyin:
            keys = json.load(keyin)
        pass
    else:
        if len(sys.argv) < 3:
            sys.exit("Not enough parameters")
        else:
            text_file = sys.argv[1]
            with open(sys.argv[2], 'r', encoding='utf_8') as keyin:
                keys = json.load(keyin)

    with open(text_file, 'r', encoding='utf_8') as tsvin:
        tsv_reader = csv.reader(tsvin, delimiter='\t')
        with Pool(8) as p:
            result = p.imap_unordered(word_count, ((row, keys) for row in tsv_reader), 100)
            print([item for item in result if len(item[1]) > 1])
@@ -277,24 +277,29 @@
 g.Tuple%5Bstr,
+ str,
  tg.Dict%5Bstr
@@ -1332,16 +1332,24 @@
 urn id_,
+ string,
  result 
@@ -1381,20 +1381,80 @@
 DEBUG = 
-True
+False%0ANUM_PROCESS = 8%0ACHUNK_SIZE = 100%0ACOOCCURENCE_THRESHOLD = 1
 %0A%0Aif __n
@@ -1877,16 +1877,45 @@
 keyin)%0A%0A
+    print(%22Config loaded.%22)%0A%0A
     with
@@ -1960,24 +1960,66 @@
 ) as tsvin:%0A
+        print(%22Text opened, processing.%22)%0A
         tsv_
@@ -2079,17 +2079,27 @@
 th Pool(
-8
+NUM_PROCESS
 ) as p:%0A
@@ -2185,19 +2185,26 @@
 eader), 
-100
+CHUNK_SIZE
 )%0A      
@@ -2213,56 +2213,706 @@
     
-print(%5Bitem for item in result if len(item%5B1%5D) %3E 1%5D
+with open(%22out_1.tsv%22, 'w', encoding='utf_8', newline='') as outfile1, open(%22out_2.tsv%22, 'w', encoding='utf_8', newline='') as outfile2:%0A                tsv_writer1 = csv.writer(outfile1, delimiter='%5Ct')%0A                tsv_writer2 = csv.writer(outfile2, delimiter='%5Ct')%0A                for id_, original_string, result_dict in result:%0A                    if len(result_dict) %3E COOCCURENCE_THRESHOLD:%0A                        tsv_writer1.writerows((%5Bid_, entity, raw_mention%5D for entity, raw_mention in result_dict.items()))%0A                        tsv_writer2.writerow(%5Bid_, original_string%5D + %5Bentity for entity in result_dict%5D)%0A        print(%22Finished. Check out_1.tsv and out_2.tsv file for result%22
 )

add status notification