export SQUAD_DIR=outputs
DATASET=${1:-none}
METHOD=${2:-none}
ACTION=${3:-run}
SPLIT=${4:-dev}
DEVICES=${5:-0}

MAX_SEQ_LENGTH=512
if [ "$DATASET" = "simple" ]; then
    MAX_SEQ_LENGTH=30
elif [ "$DATASET" = "comp" ]; then
    MAX_SEQ_LENGTH=30
elif [ "$DATASET" = "synth" ]; then
    MAX_SEQ_LENGTH=35
elif [ "$DATASET" = "hpqa" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "squad" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "bioasq" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "newsqa" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "natq" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "trivia" ]; then
    MAX_SEQ_LENGTH=512
elif [ "$DATASET" = "hotpot" ]; then
    MAX_SEQ_LENGTH=512
else
  echo "Invalid dataset ${DATASET}"
  exit 1
fi

MODEL_TYPE="roberta-base"
if [ "$ACTION" = "run" ]; then
  if [ "$METHOD" = "ig" ]; then
    echo "Run intergrated gradient"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_int_grad.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --do_eval \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --per_gpu_eval_batch_size 10 \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "tokig" ]; then
    echo "Run token level intergrated gradient"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_tok_ig.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --do_eval \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --per_gpu_eval_batch_size 10 \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "probe" ]; then
    echo "Run probing method"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_probe.py \
      --model_type roberta \
      --model_name_or_path checkpoints/squad_roberta-base \
      --dataset ${DATASET} \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "finter" ]; then
    echo "Run feature interaction method"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_inter.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "arch" ]; then
    echo "Run archipelago method"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_arch.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "lime" -o "$METHOD" = "shap" ]; then
    echo "Run "${METHOD}
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_${METHOD}.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --overwrite_output_dir \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --output_dir pred_output \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  else
    echo "No such method" $METHOD
  fi
elif [ "$ACTION" = "vis" ]; then
  if [ "$METHOD" = "ig" ]; then
    echo "Vis intergrated gradient"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_int_grad.py \
      --model_type roberta \
      --tokenizer_name $MODEL_TYPE \
      --model_name_or_path $MODEL_TYPE \
      --dataset ${DATASET} \
      --do_vis \
      --output_dir pred_output \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} \
      --visual_dir visualizations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "tokig" ]; then
    echo "Vis token level intergrated gradient"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_tok_ig.py \
      --model_type roberta \
      --tokenizer_name $MODEL_TYPE \
      --model_name_or_path $MODEL_TYPE \
      --dataset ${DATASET} \
      --do_vis \
      --output_dir pred_output \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} \
      --visual_dir visualizations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "lime" -o "$METHOD" = "shap" ]; then
    echo "Vis ${METHOD}"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_${METHOD}.py \
      --model_type roberta \
      --tokenizer_name $MODEL_TYPE \
      --model_name_or_path $MODEL_TYPE \
      --dataset ${DATASET} \
      --do_vis \
      --output_dir pred_output \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} \
      --visual_dir visualizations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  elif [ "$METHOD" = "probe" ]; then
    echo "Vis intergrated gradient"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_probe.py \
      --model_type roberta \
      --tokenizer_name $MODEL_TYPE \
      --model_name_or_path $MODEL_TYPE \
      --dataset $DATASET \
      --do_vis \
      --output_dir pred_output \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} \
      --visual_dir visualizations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  else
    echo "No such method"
  fi
elif [ "$ACTION" = "prune" ]; then
  if [ "$METHOD" = "probe" ]; then
    echo "run prune link"
    CUDA_VISIBLE_DEVICES=$DEVICES \
    python -u run_prune.py \
      --model_type roberta \
      --model_name_or_path checkpoints/${DATASET}_roberta-base \
      --dataset ${DATASET} \
      --output_dir pred_output \
      --max_seq_length ${MAX_SEQ_LENGTH} \
      --predict_file $SQUAD_DIR/${SPLIT}_${DATASET}.json \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE} \
      --visual_dir visualizations/${METHOD}_prune/${DATASET}_${SPLIT}_${MODEL_TYPE} 2>&1
  else
    echo "No such method"
  fi
elif [ "$ACTION" = "eval" ]; then
  if [ "$METHOD" = "ig" ] || [ "$METHOD" = "arch" ] || [ "$METHOD" = "probe" ] || [ "$METHOD" = "tokig" ]; then
    echo "Eval ${METHOD}"
    python -u auto_eval.py \
      --method ${METHOD}   \
      --dataset ${DATASET} \
      --interp_dir interpretations/${METHOD}/${DATASET}_${SPLIT}_${MODEL_TYPE}      
  else
    echo "No such method"
  fi
else
  echo "run or vis"
fi
