# # w.o. adapter
# python3 reeval_zero.py suffix_xgiga_generate_zero.sh mbartV2_adapterV1_rouge_en zh

# # w. adapter
# python3 reeval_zero.py suffix_xgiga_generate_zero_w_adapter_proj.sh mbartV2_iadapterfix_2stage_adapterV2_encProjFix_LAv2 zh hdfs://haruna/home/byte_arnold_hl_mlnlc/user/wuxianze.0/Workspace/Multilingual/xgiga/ZhEnFr_Unsup_infillNoise_r0.35_maxlen256_right_fix/checkpoints/mbartV2_cc100_adapterV2_encProjFix_trainLA_ls0

import os
import sys
import json

def generate_suffixs(interval=4000, end=112000, epoch=15625, start=4000):
    current = start
    results = []
    while current <= end:
        epoch_idx = current // epoch + 1
        result = "{}_{}".format(epoch_idx, current)
        results.append(result)
        current += interval
    return results

if __name__ == "__main__":
    script = sys.argv[1]
    prefix = sys.argv[2]
    lg = sys.argv[3] # en, zh, fr
    lg_code = {
        "en": "en_XX",
        "fr": "fr_XX",
        "zh": "zh_CN"
    }

    # start, end, interval
    start, end, interval = 4000, 112000, 4000
    if len(sys.argv) > 5:
        start = int(sys.argv[5])
    if len(sys.argv) > 6:
        end = int(sys.argv[6])
    if len(sys.argv) > 7:
        interval = int(sys.argv[7])
    
    suffixs = generate_suffixs(
        start=start,
        end=end,
        interval=interval
    )
    for (si, suffix) in enumerate(suffixs):
        command = ""
        if script == "suffix_xgiga_generate_zero.sh":
            command = "bash {script} mspm4_xgiga en {lg} {prefix} None {suffix} --prefix-tokens {lg_code}".format(
                script=script,
                lg=lg,
                prefix=prefix,
                suffix=suffix,
                lg_code=lg_code[lg]
            )
        else:
            hdfs_adapter = sys.argv[4]
            adapter_args = "{} enzhfr_final_ln.json".format(hdfs_adapter)
            command = "bash {script} mspm4_xgiga en {lg} {prefix} None {adapter_args} {suffix} --prefix-tokens {lg_code}".format(
                script=script,
                lg=lg,
                prefix=prefix,
                suffix=suffix,
                adapter_args=adapter_args,
                lg_code=lg_code[lg]
            )
        if si == 0:
            print("The first command:\n {}".format(command))
        pipe = os.popen(command)
        last_line = pipe.readlines()[-1].strip()
        info = {
            "prefix": prefix,
            "suffix": suffix,
            "score": last_line
        }
        print(json.dumps(info), flush=True)
