import argparse
from sdsvkie.utils import read_json, yaml_load, write_json
# from sdsvkie.utils.post_processing.invoice_post_processing import *
# from sdsvkie.utils.post_processing.common_post_processing import normalize_number
from tqdm import tqdm

INVOICE_KEYS =  [
 'no_key',
 'no_value',
 'form_key',
 'form_value',
 'serial_key',
 'serial_value',
 
 'date_value',
 
 'seller_company_name_key',
 'seller_company_name_value',
 'seller_address_value',
 'seller_address_key',
 'seller_tel_key',
 'seller_tel_value',
 'seller_tax_code_key',
 'seller_tax_code_value',

 'buyer_personal_name_key',
 'buyer_personal_name_value',
 'buyer_company_name_value', 
 'buyer_company_name_key',

 'buyer_address_key',
 'buyer_address_value',
  'buyer_tax_code_key',
 'buyer_tax_code_value',
 'buyer_tel_key',
 'buyer_tel_value',
 
 'tax_amount_key',
 'tax_amount_value',
 'total_key',
 'total_value',
 'total_in_words_key',
 'total_in_words_value',
 'other'
]

def sort_invoice(data):
    
    sorted_data = {}
    for img_key, img_data in tqdm(data.items()):
        sorted_img_data = {}
        for field_key in INVOICE_KEYS:
            if "_key" in field_key or "other" in field_key:
                continue
            sorted_img_data[field_key] = img_data.get(field_key, "")
            
        sorted_data[img_key] = sorted_img_data
        
    return sorted_data


def sort_receipt(data):
    return data



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, help="e2e label file path")
    parser.add_argument("--out", type=str, help='postprocess e2e label')
    parser.add_argument("--doc_type", default="invoice")
    args = parser.parse_args()

    data = read_json(args.input)
    if args.doc_type:
        
        sorted_data = sort_invoice(data)
    else:
        sorted_data = sort_receipt(data)
        
    write_json(args.out, sorted_data, sort_keys=False)


