#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import logging
import math
import re
import os
from dataclasses import dataclass, field
from typing import Dict, Optional, Sequence

import torch
import torch.distributed as dist
import transformers
from torch.nn import functional as F
import json
import yaml
import argparse

from peft import PeftModel, LoraConfig, TaskType, get_peft_model
from peft import PeftModel
from datasets import load_dataset, concatenate_datasets
from accelerate.utils import set_seed
from src.model import (
    CODI_CT,
    ModelArguments,
    DataArguments,
    TrainingArguments,
    load_safetensors_state_dict,
)

import numpy as np

do_print = True


def setup_distributed():
    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))

    if torch.cuda.is_available():
        device_count = torch.cuda.device_count()
        local_rank_env = os.environ.get("LOCAL_RANK")
        local_rank = int(local_rank_env) if local_rank_env is not None else (0 if device_count == 1 else rank)
        if local_rank < 0 or local_rank >= device_count:
            raise RuntimeError(
                f"Invalid LOCAL_RANK={local_rank} for {device_count} visible CUDA devices"
            )
        torch.cuda.set_device(local_rank)
        device = torch.device("cuda", local_rank)
        backend = "nccl"
    else:
        device = torch.device("cpu")
        backend = "gloo"

    if world_size > 1 and not dist.is_initialized():
        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)

    return rank, world_size, device


def gather_results(local_results, rank, world_size):
    if world_size == 1:
        return local_results

    gathered_results = [None] * world_size if rank == 0 else None
    dist.gather_object(local_results, gathered_results, dst=0)
    if rank != 0:
        return None
    return [result for worker_results in gathered_results for result in worker_results]


def read_json(file_path):
    """
    从指定路径读取JSON文件并返回对应的Python对象。
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
            return data
    except Exception as e:
        print(f"读取JSON文件时出错: {e}")
        return None

def read_svamp(file_path):
    """
    从指定路径读取JSON文件并返回对应的Python对象。
    """
    new_data = []
    sample = {}
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        for i in data:
            sample = {
                "question": i['Body'] + ' ' + i['Question'],
                "answer": i["Answer"]
            }
            new_data.append(sample.copy())
        return new_data
    except Exception as e:
        print(f"读取JSON文件时出错: {e}")
        return None


def write_json(data, file_path):
    """
    将Python对象写入指定路径的JSON文件中。
    """
    try:
        with open(file_path, 'w', encoding='utf-8') as file:
            json.dump(data, file, ensure_ascii=False, indent=4)
    except Exception as e:
        print(f"写入JSON文件时出错: {e}")


def evaluation(model_args, data_args, training_args, rank, world_size, device):
    if model_args.lora_init:
        task_type = TaskType.CAUSAL_LM
        if any(name in model_args.model_name_or_path.lower() for name in ["llama", "mistral", "falcon", "qwen"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["phi"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["gpt2"]):
            target_modules = ["c_attn", "c_proj", 'c_fc']
        else:
            raise ValueError(f"Only support LLAMA, Mistral, Falcon, Phi-2, but got {model_args.model_name_or_path}.")
        lora_config = LoraConfig(
            task_type=task_type,
            inference_mode=False,
            r=model_args.lora_r,
            lora_alpha=model_args.lora_alpha,
            lora_dropout=0.1,
            target_modules=target_modules,
            init_lora_weights=True,
        )
    else:
        raise NotImplementedError
    
    model = CODI_CT(model_args, training_args, lora_config)
    #if "llama" in model_args.model_name_or_path:
    #    model.codi.resize_token_embeddings(128261)
    state_dict = load_safetensors_state_dict(model_args.ckpt_dir)
    
    # new_state_dict = { k.replace("coconut", "codi"): v for k, v in state_dict.items() }
    # torch.save(new_state_dict, "/scratch/prj/inf_multimodal_qa/scratch_tmp/transfer/pytorch_model.bin")
    model.load_state_dict(state_dict, strict=False)
    model.codi.tie_weights()
    
    tokenizer_path = model_args.model_name_or_path 
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        tokenizer_path,
        token=model_args.token,
        model_max_length=training_args.model_max_length,
        padding_side="left",
        use_fast=False,
    )

    if tokenizer.pad_token_id is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        tokenizer.pad_token_id = model.pad_token_id
        if tokenizer.pad_token_id is None: # error handling
            tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids('[PAD]')

    model = model.to(device)
    model.to(torch.bfloat16)

    ######################
    #      dataset       #
    ######################
    logging.warning("Downloading Data")
    question_name = "question"
    answer_name = "answer"
    if "gsm-hard" == data_args.data_name:
        # dataset = load_dataset("juyoung-trl/gsm-hard")
        # test_set = dataset['train']
        # question_name = "instruction"
        # answer_name = "response"
        test_set = read_json('/data/oss_bucket_0/liqi/dataset/gsm-hard/gsm-hard.json')
        question_name = "input"
        answer_name = "target"
    elif "multi-arith" == data_args.data_name:
        # test_set = read_json("/data/oss_bucket_0/liqi/dataset/MultiArith/test_all.json")
        test_set = read_json("/data/oss_bucket_0/liqi/dataset/MultiArith/test.json")
        answer_name = "final_ans"
    elif "svamp" == data_args.data_name:
        # dataset = load_dataset("ChilleD/SVAMP")
        # test_set = concatenate_datasets([dataset["train"], dataset["test"]])
        test_set = read_svamp('/data/oss_bucket_0/liqi/dataset/SVAMP/test.json')
        # test_set = read_svamp('/data/oss_bucket_0/liqi/dataset/SVAMP/test_all.json')
    elif "commonsense" == data_args.data_name:
        dataset = load_dataset("zen-E/CommonsenseQA-GPT4omini")
        test_set = dataset['validation']
    elif "gsm8k" == data_args.data_name:
        test_set = read_json('/data/oss_bucket_0/liqi/gsk8k-aug/gsm_test.json')
    else:
        raise NotImplementedError

    total_examples = len(test_set)
    sample_indices = list(range(rank, total_examples, world_size))
    test_set = [test_set[index] for index in sample_indices]

    logging.warning(
        f"Rank {rank}/{world_size} formatting {len(test_set)} of {total_examples} inputs..."
    )
    question = [f"{example[question_name].strip().replace('  ', ' ')}" for example in test_set]
    answer = []

    # get numerical answer
    for example in test_set:
        example = str(example[answer_name])
        if isinstance(example, bool):
            answer.append(example)
            continue
        if example in ["True", "False"]:
            if example == "True":
                ans = True
            else:
                ans = False
            answer.append(ans)
            continue
        if example in "ABCDE":
            answer.append(example)
            continue
        if "####" in example:
            ans = example.split('####')[-1]
        else:
            ans = example
        ans = ans.replace(',', '')  # handle numbers like 2,000
        try:
            ans = float(ans)
        except ValueError:
            ans = float("inf")
        answer.append(ans)

    logging.warning("Tokenizing inputs...")
    eval_step = math.ceil(len(question)/data_args.batch_size)
    logging.warning(f"Total example: {len(question)} | eval batch size: {data_args.batch_size}"
                    f"eval steps: {eval_step}")
    
    question_data = []
    for i in range(eval_step):
        if i < eval_step - 1:
            batch = tokenizer(
                question[i*data_args.batch_size: (i+1)*data_args.batch_size],
                return_tensors="pt",
                padding="longest",
            )
        else:
            batch = tokenizer(
                question[i*data_args.batch_size:],
                return_tensors="pt",
                padding="longest",
            )
        
        if training_args.remove_eos:
            bot_tensor = torch.tensor([model.bot_id], dtype=torch.long).expand(batch["input_ids"].size(0), 1)
        else:
            bot_tensor = torch.tensor([tokenizer.eos_token_id, model.bot_id], dtype=torch.long).expand(batch["input_ids"].size(0), 2)
        batch["input_ids"] = torch.cat((batch["input_ids"], bot_tensor), dim=1)
        batch["attention_mask"] = torch.cat((batch["attention_mask"], torch.ones_like(bot_tensor)), dim=1)
        batch['input_len'] = len(batch['input_ids'][0])
        question_data.append(batch.to(device))

    model.eval()
    gen_kwargs = {
        "max_new_tokens": 256,
        "temperature":0.1,
        "top_k": 40,
        "top_p": 0.95,
        "do_sample": True,
    }

    ans_pred_list = []
    ans_pred_list_accu_at_n_passes = []
    attention_map_weights = []
    attention_to_latents_against_len_sum = []
    attention_to_latents_against_len_count = []
    set_seed(42)
    gating_probs_sums = None
    len_cot = []
    full_out = []
    model.eval()
    attn_to_latent_list = []
    embed_layer = model.get_embd(model.codi, model.model_name)
    
    for step, batch in enumerate(question_data):
        batch_size = batch["input_ids"].size(0)
        with torch.no_grad():
            # encode the question
            past_key_values = None
            outputs = model.codi(input_ids=batch["input_ids"], use_cache=True, output_hidden_states=True, past_key_values=past_key_values, attention_mask=batch["attention_mask"])
            past_key_values = outputs.past_key_values

            if training_args.enable_top_k:
                logits_last = outputs.logits[:, -1, :]  # [batch, vocab_size]
                latent_embd = model._compute_topk_weighted_embedding(logits_last, embed_layer)  # [batch, 1, hidden_size]
            else:
                latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1)

            if training_args.use_prj:
                latent_embd = model.prj(latent_embd)
            
            inf_latent_iterations = training_args.inf_latent_iterations

            for i in range(inf_latent_iterations):
                # decode the latent embeddings
                outputs = model.codi(inputs_embeds=latent_embd, use_cache=True, output_hidden_states=True, past_key_values=past_key_values)
                past_key_values = outputs.past_key_values

                if training_args.enable_top_k:
                    logits_last = outputs.logits[:, -1, :]  # [batch, vocab_size]
                    latent_embd = model._compute_topk_weighted_embedding(logits_last, embed_layer)  # [batch, 1, hidden_size]
                else:
                    latent_embd = outputs.hidden_states[-1][:, -1, :].unsqueeze(1)
                
                if training_args.use_prj:
                    latent_embd = model.prj(latent_embd)
            

            # batch_size = latent_embd.size(0)
            # hidden_size = latent_embd.size(2)
            # device_ = latent_embd.device
            # dtype_ = latent_embd.dtype
            # with torch.no_grad():
            #     ref_std = latent_embd.detach().std()
            # random_latents = torch.randn(
            #     batch_size, inf_latent_iterations, hidden_size,
            #     device=device_, dtype=dtype_,
            # ) * ref_std  # [batch, 6, hidden_size]
            # outputs = model.codi(inputs_embeds=random_latents, use_cache=True, output_hidden_states=True, past_key_values=past_key_values)
            # past_key_values = outputs.past_key_values


            if training_args.remove_eos:
                eot_emb = model.get_embd(model.codi, model.model_name)(torch.tensor([model.eot_id], dtype=torch.long, device=device)).unsqueeze(0)
            else:
                eot_emb = model.get_embd(model.codi, model.model_name)(torch.tensor([model.eot_id, tokenizer.eos_token_id], dtype=torch.long, device=device)).unsqueeze(0)
            
            eot_emb = eot_emb.expand(batch["input_ids"].size(0), -1, -1)

            output = eot_emb
            
            seq_len = 0
            finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
            pred_tokens = [[] for _ in range(batch_size)]
            for i in range(gen_kwargs["max_new_tokens"]):
                seq_len += 1

                out = model.codi(
                        inputs_embeds=output,
                        output_hidden_states=False,
                        attention_mask=None,
                        use_cache=True,
                        output_attentions=False,
                        past_key_values=past_key_values
                    )
                past_key_values = out.past_key_values
                logits = out.logits[:, -1, :model.codi.config.vocab_size-1]

                # implement the sampling process
                if training_args.greedy:
                    next_token_ids = torch.argmax(logits, dim=-1)
                else:
                    logits /= gen_kwargs["temperature"]
                    if gen_kwargs["top_k"] > 1:
                        top_k_values, _ = torch.topk(logits, gen_kwargs["top_k"], dim=-1)
                        min_top_k_value = top_k_values[:, -1].unsqueeze(-1)
                        logits[logits < min_top_k_value] = -float("inf")

                    if gen_kwargs["top_p"] < 1.0:
                        sorted_logit, sorted_indices = torch.sort(logits, descending=True, dim=-1)
                        cumulative_probs = torch.cumsum(F.softmax(sorted_logit, dim=-1), dim=-1)

                        sorted_indices_to_remove = cumulative_probs > gen_kwargs["top_p"]
                        if sorted_indices_to_remove.any():
                            sorted_indices_to_remove = sorted_indices_to_remove.roll(1, dims=-1)
                            sorted_indices_to_remove[:, 0] = False

                        for b in range(logits.size(0)):
                            logits[b, sorted_indices[b, sorted_indices_to_remove[b]]] = -float("inf")
                    
                    probs = F.softmax(logits, dim=-1)
                    next_token_ids = torch.multinomial(probs, num_samples=1).squeeze(-1)

                # Handle EOS for each sequence
                for b in range(batch_size):
                    if not finished[b]:
                        pred_tokens[b].append(next_token_ids[b].item())
                        if next_token_ids[b] == tokenizer.eos_token_id:
                            finished[b] = True

                # Break if all sequences have finished
                if finished.all():
                    break

                #output = model.codi.get_base_model().transformer.wte(next_token_ids).unsqueeze(1).to(device)
                output = model.get_embd(model.codi, model.model_name)(next_token_ids).unsqueeze(1).to(device)

            for mini_step, pred_token in enumerate(pred_tokens):
                local_index = step * data_args.batch_size + mini_step
                global_index = sample_indices[local_index]
                decoded_pred = tokenizer.decode(pred_token, skip_special_tokens=True)
                prediction = extract_answer_number(decoded_pred)
                len_cot.append(len(pred_token))
                full_out.append(decoded_pred)
                if do_print:
                    print(f"Rank {rank} | Question {global_index} Starts...")
                    print(f"Q: {question[local_index]}")
                    print(decoded_pred)
                    print(f"Rank {rank} | Question {global_index} Ends")
                    print(f"Prediction={prediction}; Groundtruth={answer[local_index]}")
                    print("")
                ans_pred_list.append(prediction)

    local_results = []
    for index, gt, target, pred, out, cot_length in zip(
        sample_indices, test_set, answer, ans_pred_list, full_out, len_cot
    ):
        local_results.append({
            "index": index,
            "question": gt[question_name],
            "ground truth": gt[answer_name],
            "target": target,
            "full output": out,
            "prediction": pred,
            "cot_length": cot_length,
        })

    all_results = gather_results(local_results, rank, world_size)
    metrics = [None, None]
    if rank == 0:
        all_results.sort(key=lambda result: result["index"])
        metrics[0] = compute_accuracy(
            [result["target"] for result in all_results],
            [result["prediction"] for result in all_results],
        )
        metrics[1] = sum(result["cot_length"] for result in all_results) / len(all_results)
    if world_size > 1:
        dist.broadcast_object_list(metrics, src=0)
    accuracy, average_cot_length = metrics

    if rank == 0:
        save_dir = os.path.join(training_args.output_dir, os.path.basename(model_args.ckpt_dir))
        os.makedirs(save_dir, exist_ok=True)
        save_list = [
            {
                "question": result["question"],
                "ground truth": result["ground truth"],
                "full output": result["full output"],
                "prediction": result["prediction"],
            }
            for result in all_results
        ]
        write_json(
            {
                "accuracy": 100 * accuracy,
                "average length of COT": average_cot_length,
                "output": save_list,
            },
            os.path.join(save_dir, "answer.json"),
        )
        print(f"adapter: {model_args.adapter_name_or_path} | GSM8K test accuracy: {100*accuracy:.2f}% | ")
        print(f"average length of COT: {average_cot_length}")

    # Explicitly release model and GPU memory before loading next checkpoint
    del model
    del question_data
    import gc
    gc.collect()
    torch.cuda.empty_cache()

    return 100*accuracy

def extract_answer_number(sentence: str) -> float:
    sentence = sentence.replace(',', '')
    pred = [s for s in re.findall(r'-?\d+\.?\d*', sentence)]
    if not pred:
        if "commonsense" in data_args.data_name:
            pred = sentence.split("The answer is:")[-1].strip()
            if pred[0] not in "ABCDE":
                return "C" 
            return pred[0]
        elif "strategy" in data_args.data_name or "prontoqa" in data_args.data_name.lower():
            if "True" in sentence:
                return True
            elif "False" in sentence:
                return False
            else:
                raise ValueError
        return float('inf')

    # use the last number as the answer
    pred_answer = float(pred[-1])

    return pred_answer


def compute_accuracy(gold: list, pred: list):
    acc = 0.0
    for p, g in zip(pred, gold):
        if isinstance(p, list):
            if g in p:
                acc += 1
        else:
            if p == g:
                acc += 1

    return acc / len(gold)

if __name__ == "__main__":
    rank, world_size, device = setup_distributed()
    if rank == 0:
        print(f"Distributed evaluation: world_size={world_size}, device={device}")

    arg_parser = argparse.ArgumentParser(description="CODI Testing with YAML config support")
    arg_parser.add_argument("--config_file", type=str, default=None, help="Path to YAML config file")
    args, remaining_args = arg_parser.parse_known_args()
    
    # If config_file is provided, load from YAML
    if args.config_file:
        if rank == 0:
            print(f"Loading configuration from YAML file: {args.config_file}")
        with open(args.config_file, 'r') as f:
            config = yaml.safe_load(f)
        
        # Convert YAML config to command line arguments
        cmd_args = []
        for key, value in config.items():
            if value is None:
                continue
            if isinstance(value, bool):
                if value:
                    cmd_args.extend([f"--{key}", "True"])
                else:
                    cmd_args.extend([f"--{key}", "False"])
            else:
                cmd_args.extend([f"--{key}", str(value)])
        
        # Parse with HfArgumentParser
        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args)
    else:
        # Original behavior: parse from command line
        if rank == 0:
            print("Loading configuration from command line arguments")
        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # ---------------------------------------------------------------
    # Checkpoint list: fill in the checkpoint directories to evaluate.
    # Each entry will override model_args.ckpt_dir in turn.
    # Leave as an empty list [] to fall back to the single-ckpt mode
    # (i.e. whatever ckpt_dir is set in the config / command line).
    # ---------------------------------------------------------------
    # ckpt_list = ["9000","9500","10000","10500","11000","11500","12000","12500","13000","13500","14000","14500","15000"]
    # ckpt_list = ["30000","29500","29000","28500","28000","27500","27000","26500","26000","25500","25000","24500","24000","23500","23000","22000","21000","20000","19000","18000","17000","16000","15000","14000","13000","12000","11000","10000"]
    # ckpt_list = ["60000","59500","59000","58500","58000","57500","57000","56500","56000","55500","55000","54500","54000","53500","53000","52000","51000","50000","9000"]
    # ckpt_list = ["7500","7000","6500","6000","5500","5000","4500","4000","3000","2000","1000"]
    # ckpt_list = ["15000","14500","14000","13500","13000","12500","12000","11500","11000","10500","10000","9500","9000","8500","8000","7500","7000","6500","6000","5500","5000","4000"]
    # ckpt_list = ["19000","18500","18000","17500","17000","16500","16000","15500","15000","14500","14000","13500","13000","12500","12000","11500","11000","10000"]
    # ckpt_list = ["99500","99000","98500","98000","97500","97000","96500","96000","95500","95000","94500","94000","93500","93000","92500","92000","91500","91000","90000","88000"]
    # ckpt_list = ['89970', '86971', '83972', '80973', '77974', '74975', '71976', '68977', '65978', '62979', '59980', '56981', '53982', '50983', '47984', '44985', '41986', '38987', '35988', '32989']
    # ckpt_list = ["45000","44500","44000","43500","43000","42500","42000","41500","41000","40500","40000","39500","39000","38000","37000","36000","35000","34000","33000","9500"]
    # ckpt_list = ['60000', '58500', '57000', '55500', '54000', '52500', '51000', '49500', '48000', '46500', '45000', '43500', '42000', '40500', '39000']
    # ckpt_list = ['45000', '43500', '42000', '40500', '39000', '37500', '36000', '34500', '33000', '31500', '30000', '28500', '27000', '25500', '24000']
    # ckpt_list = ['45000', '42000', '39000', '36000', '33000', '30000', '27000', '24000', '21000', '18000', '15000', '12000', '9000', '6000', '3000']
    # ckpt_list = ['30000', '27000', '24000', '21000', '18000', '15000', '12000', '9000', '6000', '3000']
    # ckpt_list = ['22500', '21750', '21000', '20250', '19500', '18750', '18000', '17250', '16500', '15750', '15000', '13500', '12000', '10500']


    # ckpt_list = ["15000","14250","13500","12750","12000","11250","10500","9750","9000","8250","7500","6750","6000"]
    ckpt_list = ['45090', '43587', '42084', '40581', '39078', '37575', '36072', '34569', '33066', '31563', '30060', '28557', '27054', '25551', '24048', '22545', '21042', '19539', '18036', '16533']
    # ckpt_list = ['75000', '73500', '72000', '70500', '69000', '67500', '66000', '64500', '63000', '61500', '60000', '58500', '57000', '55500', '54000']
    # ckpt_list = ['10000', '9500', '9000', '8500', '8000', '7500', '7000', '6500', '6000']
    # ckpt_list = ['30000', '28500', '27000', '25500', '24000', '22500', '21000', '19500', '18000', '16500', '15000', '13500', '12000', '10500', '9000']
    # ckpt_list = ['19474', '17976', '16478', '14980', '13482', '11984', '10486', '8988', '7490', '5992', '4494']
    # ckpt_list = ['19474', '17976', '16478', '14980', '13482', '11984', '10486', '8988', '7490', '5992']
    # ckpt_list = ['12560', '11932', '11304', '10676', '10048', '9420', '8792', '8164', '7536', '6908', '6280', '5652', '5024', '4396', '3768']
    # ckpt_list = ['15000', '13500', '12000', '10500', '9000', '7500', '6000', '4500', '3000', '1500']
    # ckpt_list = ['22470', '20972', '19474', '17976', '16478', '14980', '13482', '11984', '10486', '8988', '7490', '5992', '4494', '2996', '1498']
    # ckpt_list = ['22500', '21000', '19500', '18000', '16500', '15000', '13500', '12000', '10500', '9000', '7500', '6000', '4500', '3000', '1500']
    # ckpt_list = ['14700', '13965', '13230', '12495', '11760', '11025', '10290', '9555', '8820', '8085', '7350', '6615', '5880', '5145', '4410']
    # ckpt_list = ['12600', '11970', '11340', '10710', '10080', '9450', '8820', '8190', '7560', '6930', '6300', '5670', '5040', '4410', '3780']
    # ckpt_list = ['12560', '11932', '11304', '10676', '10048', '9420', '8792', '8164', '7536', '6908', '6280', '5652', '5024', '4396', '3768']
    # ckpt_list = ['44985', '41986', '38987', '35988', '32989', '29990', '26991', '23992', '20993', '17994', '14995', '11996', '8997', '5998', '2999']
    # ckpt_list = ['37635', '35126', '32617', '30108', '27599', '25090', '22581', '20072', '17563', '15054', '12545', '10036', '7527', '5018', '2509']
    # ckpt_list = ['17140', '16283', '15426', '14569', '13712', '12855', '11998', '11141', '10284', '9427', '8570', '7713', '6856', '5999', '5142']
    # ckpt_list = ['45075', '42070', '39065', '36060', '33055', '30050', '27045', '24040', '21035', '18030', '15025', '12020', '9015', '6010', '3005']
    # ckpt_list = ['1']
    # Multi-checkpoint mode: evaluate every checkpoint in ckpt_list
    if model_args.ckpt_dir is None:
        raise ValueError(
            "ckpt_dir must be specified in the YAML config or command line "
            "when evaluating multiple checkpoints. "
            "Please add `ckpt_dir: <path_to_checkpoints_root>` to your config."
        )
    all_ckpt_results = {}
    base_dir = model_args.ckpt_dir
    for ckpt_id in ckpt_list:
        if rank == 0:
            print(f"\n{'='*60}")
            print(f"Evaluating checkpoint: {ckpt_id}")
            print(f"{'='*60}")
        
        accu_list = []
        for i in range(training_args.inf_num_iterations):
            model_args.ckpt_dir = os.path.join(base_dir, f"checkpoint-{ckpt_id}")
            # model_args.ckpt_dir = base_dir
            accu = evaluation(model_args, data_args, training_args, rank, world_size, device)
            accu_list.append(accu)
            torch.cuda.empty_cache()
            if world_size > 1:
                dist.barrier()

        avg_accu = sum(accu_list) / len(accu_list)
        all_ckpt_results[ckpt_id] = avg_accu
        if rank == 0:
            print(f"Checkpoint {ckpt_id} | Average accuracy over {training_args.inf_num_iterations} sampling: {avg_accu:.2f}%")

    if rank == 0:
        write_json(all_ckpt_results, f"{training_args.output_dir}/all_ckpt_results.json")

        print(f"\n{'='*60}")
        print("Summary of all checkpoints:")
        for ckpt_dir, avg_accu in all_ckpt_results.items():
            print(f"  {ckpt_dir}: {avg_accu:.2f}%")

    if dist.is_initialized():
        dist.destroy_process_group()