import argparse
import os
import pandas as pd
import numpy as np
from pathlib import Path
from augmenter import *
import yaml

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', '-cfg', type=str, default='./configs/ucf101/ucf_zero_shot.yaml') # './configs/ucf101/ucf_zero_shot.yaml' './configs/hmdb51/hmdb_zero_shot.yaml' './configs/k700/k700_zero_shot.yaml'
    parser.add_argument('--log_time', type=str, default='20240105_172223')
    parser.add_argument('--name', type=str, default='ucf') # 'kinetics_700' 'hmdb51' 'ucf'
    parser.add_argument('--action', type=str, choices=['delete', 'substitute', 'insert', 'unified','space_del', 'space_insert'], default='space_del')
    parser.add_argument('--aug_p', type=float, default=0.2)
    parser.add_argument('--num_sim', type=int, default=1)
    args = parser.parse_args()

    with open(args.config, 'r') as f:
        config = yaml.full_load(f)
    working_dir = os.path.join('./exp', config['network']['type'], config['network']['arch'], config['data']['dataset'],
                               args.log_time)
    
    Path(working_dir).mkdir(parents=True, exist_ok=True)

    if args.action=="space_del":
        aug = MyRandomCharAug_spacedel(aug_char_p=args.aug_p)
    elif args.action=="space_insert":
        aug = MyRandomCharAug_spaceinsert(aug_char_p=args.aug_p)
    elif args.action=="unified":
        aug = MyRandomCharAugg(aug_char_p=args.aug_p)
    else:
        aug = MyRandomCharAug(args.action, aug_char_p=args.aug_p, spec_char='')

    for i in range(1,args.num_sim+1):
        if args.name == 'ucf':
            df = pd.read_csv(f"./lists/{args.name}_labels_sep.csv", index_col=0)
        else:
            df = pd.read_csv(f"./lists/{args.name}_labels.csv", index_col=0)
        if not np.isclose(args.aug_p, 0):
            df.name = df.name.apply(lambda x: aug.augment(x)[0])
        df.to_csv(os.path.join(working_dir, f"{args.name}_labels_{i}.csv"))
    

if __name__ == '__main__':
    main()
