from utils import read_local
import preprocess
import preprocessing_funcs
import keywords
import argparse
import os
from keywords import KEYWORD_MAPPINGS
from collections import defaultdict

# import tqdm


def parse_arguments():
    parser = argparse.ArgumentParser(description="Smart Chunking Tester")
    parser.add_argument(
        "--input_dir", help="Input directory (local path or S3 URI)", default="src/ip2"
    )
    parser.add_argument(
        "--output_dir", help="Output directory (local)", default="src/output_chunks"
    )
    parser.add_argument("--keyword", help="Keyword to be used for chunking")
    parser.add_argument(
        "--case_sensitive",
        help="case sensitivity",
        action=argparse.BooleanOptionalAction,
    )
    return parser.parse_args()


def main():
    args = parse_arguments()
    input_dir = args.input_dir
    output_dir = args.output_dir  # src/output_chunks
    case_sensitive = args.case_sensitive

    if case_sensitive is None:  # This shouldn't have to be assigned
        case_sensitive = False

    keyword = args.keyword
    keyword = list(map(str, keyword.split(",")))
    keyword = [kw.strip() for kw in keyword]
    print("Generating Chunks for keywords --> ", keyword)
    print("Case sensitive -->", case_sensitive)
    print(type(case_sensitive))

    kws = KEYWORD_MAPPINGS

    retained_rates = []

    for file in os.listdir(input_dir):
        full_path = os.path.join(input_dir, file)
        contract_text = read_local(full_path)
        contract_text = preprocessing_funcs.clean_newlines(contract_text)
        contract_text = preprocessing_funcs.clean_law_symbols(contract_text)
        text_dict = preprocessing_funcs.split_text(
            contract_text
        )  # return a dictionary with keys - page_num (str), values as the page_text

        test_chunck = preprocessing_funcs.smart_chunk_ac(
            text_dict=text_dict,
            keyword_mappings={
                "place_holder": {
                    "methodology": "or",
                    "keywords": keyword,
                    "case_sensitive": case_sensitive,
                }
            },
        )

        retained_rate = len(test_chunck["place_holder"]) / len(contract_text)
        retained_rates.append(retained_rate)

        print(
            f"\nfrom file {file} {len(contract_text)} characters were retrieved;\n{len(test_chunck['place_holder'])} were retained by smart chunking on keywords:\n{keyword}"
        )
        print(f"Smart chunking reduced the document by {100*(1-retained_rate):0.2f}%")

        with open(f"{output_dir}/{file}_CHUNKED.txt", "w") as f:
            f.write(test_chunck["place_holder"])

    print("avg reduction:", 1 - (sum(retained_rates) / len(retained_rates)))
    print("max retention:", max(retained_rates))
    print("min retention:", min(retained_rates))


if __name__ == "__main__":
    main()