import re

filepath = 'mlsysim/mlsysim/core/constants.py'
with open(filepath, 'r') as f:
    lines = f.readlines()

mapping_keys = [
    'GPT2_PARAMS', 'GPT2_LAYERS', 'GPT2_HIDDEN_DIM', 'GPT2_HEADS',
    'GPT3_PARAMS', 'GPT3_LAYERS', 'GPT3_HIDDEN_DIM', 'GPT3_HEADS', 'GPT3_TRAINING_OPS',
    'GPT4_EST_PARAMS', 'GPT4_LAYERS', 'GPT4_HIDDEN_DIM', 'GPT4_HEADS',
    'BERT_BASE_PARAMS', 'BERT_BASE_LAYERS', 'BERT_BASE_HIDDEN_DIM', 'BERT_BASE_HEADS', 'BERT_BASE_FLOPs',
    'BERT_LARGE_PARAMS', 'BERT_LARGE_LAYERS', 'BERT_LARGE_HIDDEN_DIM', 'BERT_LARGE_HEADS', 'BERT_LARGE_FLOPs',
    'LLAMA2_70B_PARAMS', 'LLAMA2_70B_LAYERS', 'LLAMA2_70B_HIDDEN_DIM', 'LLAMA2_70B_HEADS',
    'LLAMA3_8B_PARAMS', 'LLAMA3_8B_LAYERS', 'LLAMA3_8B_HIDDEN_DIM', 'LLAMA3_8B_HEADS', 'LLAMA3_8B_KV_HEADS',
    'LLAMA3_70B_PARAMS', 'LLAMA3_70B_LAYERS', 'LLAMA3_70B_HIDDEN_DIM', 'LLAMA3_70B_HEADS', 'LLAMA3_70B_KV_HEADS',
    'RESNET50_PARAMS', 'RESNET50_FLOPs',
    'MOBILENETV2_PARAMS', 'MOBILENETV2_FLOPs',
    'ALEXNET_PARAMS', 'ALEXNET_FLOPs',
    'YOLOV8_NANO_PARAMS', 'YOLOV8_NANO_FLOPs', 'YOLOV8_NANO_LAYERS',
    'KWS_DSCNN_PARAMS', 'KWS_DSCNN_FLOPs',
    'WAKEVISION_PARAMS', 'WAKEVISION_FLOPs',
    'ANOMALY_MODEL_PARAMS',
    'DLRM_MODEL_SIZE_FP32',
    'MAMBA_130M_PARAMS', 'MAMBA_130M_LAYERS', 'MAMBA_130M_HIDDEN_DIM', 'MAMBA_130M_STATE_SIZE',
    'MAMBA_2_8B_PARAMS', 'MAMBA_2_8B_LAYERS', 'MAMBA_2_8B_HIDDEN_DIM', 'MAMBA_2_8B_STATE_SIZE',
    'STABLE_DIFFUSION_V1_5_PARAMS', 'STABLE_DIFFUSION_V1_5_RESOLUTION', 'STABLE_DIFFUSION_V1_5_STEPS', 'STABLE_DIFFUSION_V1_5_FLOPs_PER_STEP',
    'V100_FLOPS_FP16_TENSOR', 'V100_FLOPS_FP32', 'V100_MEM_BW', 'V100_MEM_CAPACITY', 'V100_TDP', 'V100_UNIT_COST',
    'A100_FLOPS_FP16_TENSOR', 'A100_FLOPS_TF32', 'A100_FLOPS_FP32', 'A100_TOPS_INT8', 'A100_MEM_BW', 'A100_MEM_CAPACITY', 'A100_TDP', 'A100_UNIT_COST',
    'H100_FLOPS_FP16_TENSOR', 'H100_FLOPS_FP8_TENSOR', 'H100_FLOPS_TF32', 'H100_TOPS_INT8', 'H100_MEM_BW', 'H100_MEM_CAPACITY', 'H100_TDP', 'H100_UNIT_COST',
    'H200_MEM_BW', 'H200_MEM_CAPACITY', 'H200_TDP', 'H200_UNIT_COST',
    'B200_FLOPS_FP16_TENSOR', 'B200_FLOPS_FP8_TENSOR', 'B200_FLOPS_FP4_TENSOR', 'B200_TOPS_INT4', 'B200_MEM_BW', 'B200_MEM_CAPACITY', 'B200_TDP', 'B200_UNIT_COST',
    'NVL72_GPUs', 'NVL72_FLOPS_FP16_TENSOR', 'NVL72_FLOPS_FP8_TENSOR', 'NVL72_FLOPS_FP4_TENSOR', 'NVL72_MEM_CAPACITY', 'NVL72_MEM_BW', 'NVL72_NVLINK_BW', 'NVL72_TDP', 'NVL72_UNIT_COST',
    'MI300X_FLOPS_FP16_TENSOR', 'MI300X_FLOPS_FP8', 'MI300X_TOPS_INT8', 'MI300X_FLOPS_FP32', 'MI300X_MEM_BW', 'MI300X_MEM_CAPACITY', 'MI300X_TDP', 'MI300X_UNIT_COST',
    'MI250X_FLOPS_FP16_TENSOR', 'MI250X_FLOPS_FP32', 'MI250X_TOPS_INT8', 'MI250X_MEM_BW', 'MI250X_MEM_CAPACITY', 'MI250X_TDP',
    'GAUDI2_FLOPS_BF16', 'GAUDI2_FLOPS_FP8', 'GAUDI2_MEM_BW', 'GAUDI2_MEM_CAPACITY', 'GAUDI2_TDP',
    'GAUDI3_FLOPS_BF16', 'GAUDI3_FLOPS_FP8', 'GAUDI3_MEM_BW', 'GAUDI3_MEM_CAPACITY', 'GAUDI3_TDP',
    'TRAINIUM2_FLOPS_BF16', 'TRAINIUM2_FLOPS_FP8', 'TRAINIUM2_MEM_BW', 'TRAINIUM2_MEM_CAPACITY', 'TRAINIUM2_TDP',
    'TPUV4_FLOPS_BF16', 'TPUV4_MEM_BW',
    'TPUV5P_FLOPS_BF16', 'TPUV5P_TOPS_INT8', 'TPUV5P_MEM_BW', 'TPUV5P_MEM_CAPACITY', 'TPUV5P_TDP',
    'TPUV6_FLOPS_BF16', 'TPUV6_MEM_BW', 'TPUV6_MEM_CAPACITY',
    'T4_FLOPS_FP16_TENSOR', 'T4_TOPS_INT8', 'T4_MEM_BW', 'T4_TDP', 'T4_UNIT_COST',
    'WSE3_FLOPS_FP16', 'WSE3_MEM_CAPACITY', 'WSE3_MEM_BW', 'WSE3_TDP', 'CEREBRAS_CS3_UNIT_COST',
    'MOBILE_NPU_TOPS_INT8', 'MOBILE_NPU_MEM_BW',
]

new_lines = []
for line in lines:
    match = re.match(r'^([A-Z0-9_]+)\s*=', line)
    if match and match.group(1) in mapping_keys:
        continue # Skip this line
    new_lines.append(line)

with open(filepath, 'w') as f:
    f.writelines(new_lines)
print("Removed legacy constants from constants.py")