
import os 
import glob 
import json 
import os
import shutil

def read_txt(txt):
    with open(txt, 'r', encoding='utf8') as f:
        data = [line.strip() for line in f]
    return data 

def write_txt(txt, data):
    with open(txt, 'w', encoding='utf8') as f:
        for line in data:
            f.write(line + "\n")

def write_json(json_path, data):
    with open(json_path, 'w', encoding='utf8') as f:
        json.dump(data, f, ensure_ascii=False)


def read_json(json_path):
    with open(json_path, 'r', encoding='utf8') as f:
        data = json.load(f)

    return data 

def create_template_info(data_dir, json_out):
    outputs ={}
    txt_paths = sorted(glob.glob(data_dir + "/*.txt"))
    
    for txt_path in txt_paths:
        txt_name = os.path.basename(txt_path)
        txt_data = read_txt(txt_path)
        wordgroups = [item for item in txt_data if "____kie_wordgroup	seller_company_name_value" in item]
        num_line_company = len(wordgroups)
        outputs[txt_name] = num_line_company
    write_json(json_out, outputs)
    

def filter_data(template_file, data_file, img_dir, txt_dir,  output):
    template_data = read_json(template_file)
    data = read_json(data_file)

    new_data = []
    for txt_name, num_wordgroup_line in template_data.items():
        id = txt_name.split("_type99_")[0]
        for txt_name_target, num_wordgroup_target_line in data.items():
            id_target = txt_name_target.split("_type99_")[0]
            # print(id, id_target, id_target != id)
            if id_target != id:
                # print(id_target, id)
                continue
            if num_wordgroup_line != num_wordgroup_target_line:
                continue
                
            new_data.append(txt_name_target)

    new_data = sorted(list(set(new_data)))
    print(new_data[:5])
    # return new_data

    if not os.path.exists(output):
        os.makedirs(output, exist_ok=True)

    for txt_name in new_data:
        img_path = os.path.join(img_dir, txt_name.replace(".txt", ".jpg"))
        shutil.copy(img_path, output)
        shutil.copy(os.path.join(txt_dir, txt_name), output)


            

if __name__ == "__main__":
    target_dir = "/mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_one_line_r20_txt"
    template_dir = "/mnt/ssd1T/hoanglv/Projects/KIE/sdsvkie/workdirs/visualize/vnpt_one_line_txt"
    output_dir = "/mnt/ssd1T/hoanglv/Projects/KIE/craw_data/output/synth_vnpt_r20/one_line_filtered"
    img_target_dir = "/mnt/ssd1T/hoanglv/Projects/KIE/craw_data/output/synth_vnpt_r20/one_line"
    out_template_json = "vnpt_template.json"
    out_target_json = "vnpt_r20.json"
    create_template_info(template_dir, out_template_json)
    create_template_info(target_dir, out_target_json)

    new_data = filter_data(out_template_json, out_target_json,img_target_dir, target_dir,  output_dir)
    print("Total after filter: ", len(new_data))

