"""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

def word_count(args: tg.Tuple[tg.Iterable[str, str], tg.Dict[str, tg.List]]) -> tg.List:
    """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, keyin = args

    result = []
    id_, string = row
    for type, keys in keyin.items():
        is_match = False
        if is_match:
            break

        if type in ("abbreviation", "regular_expression"):
            for key in keys:
                if re.search(key, string):
                    is_match = True
                    break
        elif type in ("full_name"):
            for key in keys:
                if re.search(key, string, re.IGNORECASE):
                    is_match = True
                    break

        result.append((id_, is_match))

    return result
    
DEBUG = True

if __name__ == "__main__":
    if DEBUG:
        textin = "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:
            textin = sys.argv[1]
            with open(sys.argv[2], 'r', encoding='utf_8') as keyin:
                keys = json.load(keyin)

    with open(textin, '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])
@@ -198,22 +198,13 @@
 %5Btg.
-Iterable%5Bstr, 
+List%5B
 str%5D
@@ -460,21 +460,23 @@
    row, 
+raw_
 key
-in
  = args%0A
@@ -526,27 +526,37 @@
 for 
-type, 
+entity, entity_
 key
-s
  in 
+raw_
 key
-in
 .ite
@@ -586,16 +586,74 @@
 = False%0A
+        for key_type, key_list in entity_key.items():%0A    
         
@@ -669,39 +669,51 @@
 ch:%0A            
+    
 break%0A
-%0A
+    %0A    
         if type 
@@ -707,16 +707,20 @@
      if 
+key_
 type in 
@@ -763,32 +763,36 @@
 %22):%0A            
+    
 for key in keys:
@@ -781,33 +781,37 @@
   for key in key
-s
+_list
 :%0A              
@@ -804,32 +804,36 @@
                 
+    
 if re.search(key
@@ -855,32 +855,36 @@
                 
+    
 is_match = True%0A
@@ -899,30 +899,38 @@
             
+ 
+   
 break%0A
+    
         elif
@@ -950,16 +950,17 @@
 ll_name%22
+,
 ):%0A     
@@ -958,32 +958,36 @@
 ,):%0A            
+ 
+   
 for key in keys:
@@ -988,11 +988,17 @@
  key
-s:%0A
+_list:%0A  
     
@@ -1001,32 +1001,34 @@
                 
+  
 if re.search(key
@@ -1065,32 +1065,36 @@
                 
+    
 is_match = True%0A
@@ -1109,24 +1109,28 @@
             
+   
+ 
 break%0A%0A     
@@ -1151,16 +1151,24 @@
 nd((id_,
+ entity,
  is_matc
@@ -1262,18 +1262,21 @@
     text
-in
+_file
  = %22TEST
@@ -1510,18 +1510,21 @@
     text
-in
+_file
  = sys.a
@@ -1657,18 +1657,21 @@
 pen(text
-in
+_file
 , 'r', e

change confused varieble names