import torch
from PIL import Image
import argparse
import os
import json
import yaml
from tqdm import tqdm
from datetime import datetime
import random
from pprint import pprint
from copy import deepcopy
from env import ROOT_DIR, DATA_DIR # /badVLM, /lavis


def update_config(args_dict):
    attack_type = args_dict.pop('attack_type', 'blended')
    
    with open(f'{ROOT_DIR}/backdoors/config/{attack_type}/default.yaml') as f:
        default_cfg = yaml.safe_load(f)
    
    default_cfg.update(args_dict)

    with open(f'{ROOT_DIR}/backdoors/config/{attack_type}/run_cfg.yaml', "w") as f:
        yaml.dump(default_cfg, f, default_flow_style=False)

 
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description="Evaluating Backdoor")

    parser.add_argument("--attack-type", help="Attack type", required=True)
    parser.add_argument("--target-label", help="Target label", type=str, default=None)
    parser.add_argument("--pattern", help="Trigger type", type=str, default=None)
    parser.add_argument("--dual", help="Using dual key", type=bool, default=False)
    parser.add_argument("--patch-location", help="Trigger location", type=str, default=None)
    parser.add_argument("--pattern-size", help="Trigger size", type=int, default=None)
    parser.add_argument("--trigger-path", help="Path to trigger pattern", type=str, default=None)
    parser.add_argument("--poison-size", help="Num of poisoned samples", type=int, default=None)
    parser.add_argument("--dataset-size", help="The total dataset size", type=int, default=None)
    parser.add_argument("--dataset", help="The dataset name", type=str,choices=['coco', 'cc3m'], default='coco')
    parser.add_argument("--eps", help="Blended ratio for BadVLM", type=float, default=None)
    parser.add_argument("--blended-ratio", help="Blended ratio for Blended", type=float, default=None)
    
    args = parser.parse_args()

    ## Delete items that have None values
    args_dict = {k: v for k, v in args.__dict__.items() if v is not None}

    update_config(args_dict)
