# @eval_vlmeval_shiyi_10_20_2025.sh 把python file抽出来单独放一个文件
"""Generate VLMEvalKit configuration from environment variables."""

from __future__ import annotations

import json
import os


FALSEY_VALUES = {"0", "false", "no", ""}


# Dataset name to class mapping for auto-detection
DATASET_CLASS_MAP = {
    # VQA datasets (ImageVQADataset)
    "GQA_TestDev_Balanced": "ImageVQADataset",
    "TextVQA_VAL": "ImageVQADataset",
    "OCRVQA_TEST": "ImageVQADataset",
    "OCRVQA_TESTCORE": "ImageVQADataset",
    "DocVQA_VAL": "ImageVQADataset",
    "DocVQA_TEST": "ImageVQADataset",
    "InfoVQA_VAL": "ImageVQADataset",
    "InfoVQA_TEST": "ImageVQADataset",
    "ChartQA_TEST": "ImageVQADataset",
    # Y/N datasets (ImageYORNDataset)
    "POPE": "ImageYORNDataset",
    "MME": "ImageYORNDataset",
    "HallusionBench": "ImageYORNDataset",
    # MCQ datasets (ImageMCQDataset)
    "MMBench_DEV_EN": "ImageMCQDataset",
    "MMBench_DEV_EN_V11": "ImageMCQDataset",
    "MMBench_TEST_EN": "ImageMCQDataset",
    "MMBench_TEST_EN_V11": "ImageMCQDataset",
    "SEEDBench_IMG": "ImageMCQDataset",
    "ScienceQA_VAL": "ImageMCQDataset",
    "ScienceQA_TEST": "ImageMCQDataset",
    "AI2D_TEST": "ImageMCQDataset",
    "RealWorldQA": "ImageMCQDataset",
    # MMMU datasets (MMMUDataset)
    "MMMU_DEV_VAL": "MMMUDataset",
    "MMMU_TEST": "MMMUDataset",
    "MMMU_Pro_V": "MMMUProDataset",
    "MMMU_Pro_V_COT": "MMMUProDataset",
    # MathVista (MathVista class)
    "MathVista_MINI": "MathVista",
    "MathVista": "MathVista",
    # MathVision (MathVision class) - for CoT math reasoning
    "MathVision": "MathVision",
    "MathVision_MINI": "MathVision",
    # Image2Struct (Image2Struct class) - structure extraction benchmark
    "Image2Struct": "Image2Struct",
    "Image2Struct_Latex": "Image2Struct",
    "Image2Struct_Webpage": "Image2Struct",
    "Image2Struct_Music": "Image2Struct",
    # ChartMimic (ChartMimic class) - chart-to-code generation benchmark
    "ChartMimic_v2_direct": "ChartMimic",
    "ChartMimic_v2_direct_600": "ChartMimic",
    "ChartMimic_v2_direct_1800": "ChartMimic",
    "ChartMimic_v2_customized": "ChartMimic",
    "ChartMimic_v2_customized_600": "ChartMimic",
    "ChartMimic_v2_customized_1800": "ChartMimic",
    "ChartMimic_v1_direct": "ChartMimic",
    "ChartMimic_v1_customized": "ChartMimic",
}


def get_dataset_class(dataset_name: str, default_class: str = "ImageMCQDataset") -> str:
    """Auto-detect the dataset class based on dataset name."""
    # Check exact match first
    if dataset_name in DATASET_CLASS_MAP:
        return DATASET_CLASS_MAP[dataset_name]
    
    # Check prefix patterns
    for key, cls in DATASET_CLASS_MAP.items():
        if dataset_name.startswith(key.split("_")[0]):
            return cls
    
    return default_class


def main() -> None:
    # Determine model class to use
    # Use MDM model class if MODEL_CLASS is set to "Qwen2VLChatMDM" or if USE_MDM is set
    use_mdm = os.environ.get("USE_MDM", "").lower() not in FALSEY_VALUES
    model_class_env = os.environ.get("MODEL_CLASS", "")
    model_path_abs = os.environ.get("MODEL_PATH_ABS", "")
    model_path_lower = model_path_abs.lower()

    if model_class_env:
        model_class = model_class_env
    elif use_mdm:
        # Auto-route Nemotron-Ministral checkpoints to dedicated wrapper.
        if "nemotron" in model_path_lower or "ministral" in model_path_lower:
            model_class = "NemotronMinistralChatMDM"
        else:
            model_class = "Qwen2VLChatMDM"
    else:
        model_class = "Qwen2VLChat"  # Default VLMEvalKit class

    # Build model configuration
    model_cfg = {
        "class": model_class,
        "model_path": os.environ["MODEL_PATH_ABS"],
        "processor_path": os.environ.get(
            "PROCESSOR_PATH", "Qwen/Qwen2.5-VL-3B-Instruct"
        ),
        "torch_dtype": os.environ.get("TORCH_DTYPE", "bfloat16"),
    }

    # Add MDM-specific parameters if using MDM model
    if model_class in {"Qwen2VLChatMDM", "NemotronMinistralChatMDM"}:
        block_size_env = os.environ.get("BLOCK_SIZE", "")
        threshold_env = os.environ.get("THRESHOLD", "")
        if block_size_env:
            model_cfg["block_size"] = int(block_size_env)
        if threshold_env:
            model_cfg["threshold"] = float(threshold_env)
        model_cfg["max_tokens"] = int(os.environ.get("MAX_TOKENS", "128"))
        model_cfg["mask_token"] = os.environ.get("MASK_TOKEN", "|<MASK>|")
        model_cfg["device"] = os.environ.get("DEVICE", "cuda")

        if model_class == "Qwen2VLChatMDM":
            # A/B test: speculative decoding toggle
            # USE_SPECULATIVE=1 enables speculative decoding, USE_SPECULATIVE=0 (default) uses standard decoding
            use_speculative_env = os.environ.get("USE_SPECULATIVE", "0").lower()
            model_cfg["use_speculative"] = use_speculative_env not in FALSEY_VALUES

            # A/B test: force left-to-right decoding toggle
            # FORCE_LEFT_TO_RIGHT=1 uses sequential L2R decoding, FORCE_LEFT_TO_RIGHT=0 uses confidence-based parallel
            force_l2r_env = os.environ.get("FORCE_LEFT_TO_RIGHT", "1").lower()
            model_cfg["force_left_to_right"] = force_l2r_env not in FALSEY_VALUES

            # Track tokens_per_step statistics (default: enabled)
            track_tps_env = os.environ.get("TRACK_TOKENS_PER_STEP", "1").lower()
            model_cfg["track_tokens_per_step"] = track_tps_env not in FALSEY_VALUES
        else:
            # Nemotron-Ministral specific optional knobs.
            dlm_paradigm = os.environ.get("DLM_PARADIGM", "bidirectional")
            if dlm_paradigm:
                model_cfg["dlm_paradigm"] = dlm_paradigm

            steps_env = os.environ.get("STEPS", "")
            if steps_env:
                model_cfg["steps"] = int(steps_env)

            shift_logits_env = os.environ.get("SHIFT_LOGITS", "")
            if shift_logits_env:
                model_cfg["shift_logits"] = shift_logits_env.lower() not in FALSEY_VALUES

            causal_context_env = os.environ.get("CAUSAL_CONTEXT", "")
            if causal_context_env:
                model_cfg["causal_context"] = causal_context_env.lower() not in FALSEY_VALUES

            track_tps_env = os.environ.get("TRACK_TOKENS_PER_STEP", "1").lower()
            model_cfg["track_nfe_stats"] = track_tps_env not in FALSEY_VALUES

    cfg = {
        "model": {
            os.environ.get("MODEL_KEY", "Qwen2.5-VL-FT"): model_cfg
        },
        "data": {},
    }

    datasets = os.environ.get("DATASETS", "MMBench_DEV_EN_V11").split()
    # Get default dataset class from env, or None for auto-detection
    default_dataset_class = os.environ.get("DATASET_CLASS", "")
    auto_detect = os.environ.get("AUTO_DETECT_CLASS", "1").lower() not in FALSEY_VALUES
    model_verbose = os.environ.get("MODEL_VERBOSE", "")

    for dataset in datasets:
        if auto_detect and not default_dataset_class:
            dataset_class = get_dataset_class(dataset)
        else:
            dataset_class = default_dataset_class or "ImageMCQDataset"
        cfg["data"][dataset] = {"class": dataset_class, "dataset": dataset}
        print(f"[config] Dataset: {dataset} -> Class: {dataset_class}")

    if model_verbose and model_verbose.lower() not in FALSEY_VALUES:
        for model_cfg_item in cfg["model"].values():
            model_cfg_item["verbose"] = True

    cfg_path = os.environ["CFG_PATH"]
    with open(cfg_path, "w", encoding="utf-8") as file:
        json.dump(cfg, file, indent=2)

    print(f"[config] Saved to: {cfg_path}")


if __name__ == "__main__":
    main()

