# -*- coding: utf-8 -*-
"""
Created on Tue Oct 27 15:58:42 2020

@author: DrLC
"""

from token_level import UIDStruct, WSStruct, UIDStruct_Java
from exp_level import CMPStruct
from utils import is_uid, is_special_id, normalize, denormalize, is_java_uid, is_java_special_id
from torch.nn import CrossEntropyLoss, Softmax
from codebert_attack import CodeBERT_Attack_UID, CodeBERT_Attack_CMP, CodeBERT_Attack_WS
from mhm import MHM_Baseline

# Import codebert & rnn later

import torch    
import argparse
import os
import pickle
import json
import time
import numpy
import random
import copy

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', type=str, default='-1', help="Gpu selection")
    parser.add_argument('--mlm_path', type=str,
                        default="/var/data/lushuai/bertvsbert/save/only_mlm/poj/checkpoint-9000-1.0555",
                        help="Path to the masked language model")
    parser.add_argument('--cls_path', type=str,
                        default="/var/data/lushuai/bertvsbert/save/only_mlm/poj-classifier/checkpoint-51000-0.986",
                        help="Path to the downstream CodeBERT classifier")
    parser.add_argument('--trainset', type=str, default="../data/train.pkl",
                        help="Path to the train set")
    parser.add_argument('--testset', type=str, default="../data/test.pkl",
                        help="Path to the test set")
    parser.add_argument('--uid_path', type=str, default="../data/all_uids.pkl",
                        help="Path to the uid file")
    parser.add_argument('--attack', type=str, default='cba',
                        help="Attack approach")
    parser.add_argument('--max_perturb_iter', type=int, default=20,
                        help="Maximal iteration of perturbtions")
    parser.add_argument('--max_vulnerable', type=int, default=5,
                        help="Maximal vulnerable number")
    parser.add_argument('--max_mask_ws', type=int, default=30,
                        help="Maximal mask number for white space attack")
    parser.add_argument('--max_candidate', type=int, default=10,
                        help="Maximal candidate number")
    parser.add_argument('--smooth_factor', type=float, default=0.1,
                        help="Smoothing factor during merging")
    parser.add_argument('--init_temperature', type=float, default=1,
                        help="Temperature initialization for SA")
    parser.add_argument('--cooling_factor', type=float, default=0.8,
                        help="Temperature decreasement for SA")
    parser.add_argument('--model', type=str, default="CB",
                        help="Target model / victim model")
    parser.add_argument('--so_path', type=str, default='../data/java-language.so',
                        help="Path to the java parser library")
    parser.add_argument('--word2vec', type=str, default="../data/w2v.model",
                        help="Path to the word2vec matrix")
    parser.add_argument('--rnn_path', type=str, default='../model/oj_lstm/model.pt',
                        help="Path to the downstream RNN classifier")
    
    parser.add_argument('--data_name', type=str, default="OJ",
                        help='Name of the dataset')
    parser.add_argument('--rand_idx', type=bool, default=False,
                        help="Random indexing without vulnerablility")
    parser.add_argument('--rand_cand', type=bool, default=True,
                        help="Random candidate without codebert")
    
    opt = parser.parse_args()
    
    print ("CONFIGS")
    args = vars(opt)
    for k in args.keys():
        print ("  " + k + " = " + str(args[k]))
    
    _attack = opt.attack.upper()
    # CodeBERT-Attack, MHM, Simulated Annealling CBA, Compare-swapping CBA, Compare-swapping SACBA， White-space-inserting CBA, White-space-inserting SACBA
    assert _attack in ['CBA', 'MHM', 'SACBA', 'CBA-CMP', 'SACBA-CMP', 'CBA-WS', 'SACBA-WS']
    _victim_model = opt.model.upper()
    assert _victim_model in ['CB', 'LSTM']   # CodeBERT, LSTM
    
    n_perturb = opt.max_perturb_iter
    n_vulnerable = opt.max_vulnerable
    n_mask_ws = opt.max_mask_ws
    n_candidate = opt.max_candidate
    smoothing = opt.smooth_factor

    # Load the models (MLM and downstream CLS)
    if int(opt.gpu) < 0:
        device = torch.device("cpu")
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu
        device = torch.device("cuda")
        
    # Cannot select GPU after importing transformers...
    from codebert import codebert_mlm, codebert_cls
    from rnn import RNNClassifier
        
    data_name = opt.data_name.upper()
    n_class = {"OJ": 104, "JDP": 2}[data_name]
    class_op = {"OJ": -1, "JDP": 0}[data_name]
    if data_name == "JDP":
        _is_uid = is_java_uid
        _is_special_id = is_java_special_id
    elif data_name == "OJ":
        _is_uid = is_uid
        _is_special_id = is_special_id
    _lang = {"OJ": "C", "JDP": "JAVA"}[data_name]
    
    mlm_model = codebert_mlm(opt.mlm_path, device)
    if _victim_model == "CB":
        cls_model = codebert_cls(opt.cls_path, device, n_class=n_class)
    else:
        hidden_size = 600
        n_layers = 2
        max_len = 512
        attn = True
        bidirection = True
        cls_model = RNNClassifier(num_class=n_class,
                                  hidden_dim=hidden_size,
                                  n_layers=n_layers,
                                  tokenizer_path=opt.cls_path,
                                  w2v_path=opt.word2vec,
                                  max_len=max_len,
                                  drop_prob=0.,
                                  model=_victim_model,
                                  brnn=bidirection,
                                  attn=attn,
                                  device=device).to(device)
        cls_model.load_state_dict(torch.load(opt.rnn_path))
        cls_model.eval()
    # Load the vocabulary of MLM
    vocab_path = mlm_model.tokenizer.vocab_files_names["vocab_file"]
    with open(os.path.join(opt.mlm_path, vocab_path), "r") as f:
        txt2idx = json.load(f)
        tmp = sorted(txt2idx.items(), key=lambda it: it[1])
        idx2txt = [it[0] for it in tmp]
        assert txt2idx[idx2txt[-1]] == len(idx2txt) - 1, \
            "\n"+idx2txt[-1]+"\n"+str(txt2idx[idx2txt[-1]])+"\n"+str(len(idx2txt)-1)
    # Load the test set
    with open(opt.testset, "rb") as f:
        d = pickle.load(f)
        
    if opt.rand_cand:
        with open(opt.uid_path, "rb") as f:
            all_uids = pickle.load(f)
        print (len(all_uids))
            
    # Renaming attack
    if _attack == 'CBA':
        atk = CodeBERT_Attack_UID(lang=_lang)
    elif _attack == 'MHM':
        # When using MHM, collect all uids from the training set
        with open(opt.uid_path, "rb") as f:
            all_uids = pickle.load(f)
        print (len(all_uids))
        atk = MHM_Baseline(all_uids, lang=_lang)
    elif _attack == 'SACBA':
        atk = CodeBERT_Attack_UID(lang=_lang)
        temperature = lambda n: opt.init_temperature * (opt.cooling_factor ** n)
    elif _attack == 'CBA-CMP':
        atk = CodeBERT_Attack_CMP()
    elif _attack == 'SACBA-CMP':
        atk = CodeBERT_Attack_CMP()
        temperature = lambda n: opt.init_temperature * (opt.cooling_factor ** n)
    elif _attack == 'CBA-WS':
        atk = CodeBERT_Attack_WS()
    elif _attack == 'SACBA-WS':
        atk = CodeBERT_Attack_WS()
        temperature = lambda n: opt.init_temperature * (opt.cooling_factor ** n)
    else:
        assert False
    
    n_total_including_originally_wrong = 0
    n_total = 0
    n_succ = 0
    time_total = 0
    
    ce = CrossEntropyLoss(reduction="none")
    softmax = Softmax(dim=-1)
    
    for i in range(len(d['norm'])):
        print ("Attack %d / %d. Class %d" % \
               (i+1, len(d['norm']), d['label'][i]))
        succ = False
        n_total_including_originally_wrong += 1
        time_st = time.time()
        s = [t.strip() for t in d['norm'][i]]
        s_norm = normalize(s)
        s = denormalize(s)
        y = d['label'][i] + class_op
        logits = cls_model.run([" ".join(s_norm)])[0]
        # Skip those original erroneously predicted examples
        if logits.argmax().item() != y:
            print ("  WRONG. SKIP!")
            continue
        # Start adversarial attack
        old_prob = softmax(logits)[y].item()
        if _attack in ['MHM', 'CBA', 'SACBA']:
            if _lang == 'C':
                uid = UIDStruct(s, mask=mlm_model.tokenizer.unk_token)
            elif _lang == 'JAVA':
                uid = UIDStruct_Java(s, mask=mlm_model.tokenizer.unk_token, so_path=opt.so_path)
            print ("  UIDs: ", end="")
            for i in uid.sym2pos.keys():
                print (i, end=" ")
            print ()
        elif _attack in ['CBA-CMP', 'SACBA-CMP']:
            cmp = CMPStruct(s, mask=mlm_model.tokenizer.unk_token)
            print ("  CMPs: ", end="")
            for s, _, t in cmp.cmptab:
                for w in cmp.code[s: t]:
                    print (w, end=" ")
                print (",", end=" ")
            print ()
        elif _attack in ['CBA-WS', 'SACBA-WS']:
            ws = WSStruct(s, mask=mlm_model.tokenizer.unk_token, max_len=512)
    
        # MHM
        if _attack == 'MHM':
            
            res = atk.mcmc(uid=uid, label=y,
                           classifier=cls_model,
                           n_candi=n_candidate,
                           max_iter=n_perturb)
            if res['succ']:
                succ = True
            
        # CodeBert-Attack (UID Renaming)
        elif _attack in ['CBA', 'SACBA']:
            
            for it in range(n_perturb):
                # Find vulnerable identifiers
                if opt.rand_idx:
                    vulnerables = random.sample(list(uid.sym2pos.keys()), 1)
                    vulnerables = {i: None for i in vulnerables}
                else:
                    vulnerables = atk.find_vulnerable_uids(cls=cls_model,
                                                           uid=uid,
                                                           ground_truth_label=y,
                                                           n_vul=n_vulnerable)
                candidate_uids, candidate_seqs, candidate_old_uids = [], [], []
                for v in vulnerables.keys():
                    # Generate possible candidates for each vulnerable uids
                    if opt.rand_cand:
                        c = []
                        s = []
                        for _c in random.sample(all_uids, n_candidate):
                            if _c not in uid.sym2pos.keys():
                                c.append(_c)
                                s.append(copy.deepcopy(uid.code))
                                for jjj in uid.sym2pos[v]:
                                    s[-1][jjj] = _c
                                s[-1] = " ".join(s[-1])
                    else:
                        c, s = atk.generate_candidates(uid=uid,
                                                       mlm=mlm_model,
                                                       vulnerable=v,
                                                       idx2txt=idx2txt,
                                                       bpe_indicator='Ġ',
                                                       n_candidate=n_candidate,
                                                       smoothing=smoothing,
                                                       batch_size=n_candidate*2,
                                                       criterion=ce,
                                                       max_computational=1e6,
                                                       len_threshold=510)
                    if c is None or s is None:
                        continue
                    candidate_old_uids += [v for _ in c]
                    candidate_uids += c
                    candidate_seqs += s
                if len(candidate_seqs) <= 0:
                    break
                # Probe the target model
                probs = softmax(cls_model.run(candidate_seqs))
                preds = probs.argmax(dim=-1)
                for pi in range(len(preds)):
                    # Find an adversarial example
                    if preds[pi].item() != y:
                        print ("  %s => %s, %d (%.5f%%) => %d %d (%.5f%% %.5f%%)" % \
                               (candidate_old_uids[pi], candidate_uids[pi], y, old_prob*100,
                                y, preds[pi], probs[pi][y].item()*100, probs[pi][preds[pi]].item()*100))
                        succ = True
                        assert uid.update_sym(candidate_old_uids[pi], candidate_uids[pi]), \
                            "\n"+str(uid.sym2pos.keys())+"\n"+candidate_old_uids[pi]+"\n"+candidate_uids[pi]
                        assert _is_uid(candidate_uids[pi])
                        break
                if succ:
                    break
                
                next_i = probs[:, y].argmin().item()
                # CBA - Test if the ground truth probability decreases
                # SACBA - Accept / reject to jump to the candidate
                accept = (probs[next_i, y] < old_prob)
                if _attack == 'SACBA' and (not accept):
                    acc_prob = torch.exp(-(probs[next_i, y] - old_prob) / temperature(it+1))
                    accept = (numpy.random.uniform(0,1) < acc_prob)
                if accept:
                    print ("  %s => %s, %d (%.5f%%) => %d (%.5f%%)" % \
                           (candidate_old_uids[next_i], candidate_uids[next_i], y, old_prob*100,
                            y, probs[next_i][y].item()*100))
                    old_prob = probs[next_i][y].item()
                    assert uid.update_sym(candidate_old_uids[next_i], candidate_uids[next_i]), \
                        "\n"+str(uid.sym2pos.keys())+"\n"+candidate_old_uids[next_i]+"\n"+candidate_uids[next_i]
                    assert _is_uid(candidate_uids[next_i]), candidate_uids[next_i]
                else:
                    break
                
        # CodeBert-Attack (white space insertion)
        elif _attack in ['CBA-WS', 'SACBA-WS']:
            
            for it in range(n_perturb):
                # Find vulnerable identifiers
                vulnerables = atk.find_vulnerable_tokens(cls=cls_model,
                                                         ws=ws,
                                                         ground_truth_label=y,
                                                         n_vul=n_vulnerable,
                                                         n_mask=n_mask_ws)
                if len(vulnerables) <= 0:
                    break
                candidate_wss, candidate_seqs = atk.generate_candidates(ws=ws,
                                                                        mlm=mlm_model,
                                                                        vulnerable=vulnerables, 
                                                                        n_candidate=n_candidate,
                                                                        batch_size=n_candidate*3,
                                                                        criterion=ce,
                                                                        len_threshold=510)
                if len(candidate_seqs) <= 0:
                    break
                # Probe the target model
                probs = softmax(cls_model.run(candidate_seqs))
                preds = probs.argmax(dim=-1)
                for pi in range(len(preds)):
                    # Find an adversarial example
                    if preds[pi].item() != y:
                        print ("  idx %d <%s>, %d (%.5f%%) => %d %d (%.5f%% %.5f%%)" % \
                               (candidate_wss[pi][0],
                                {" ":"SPC", "\t":"TAB", "\n":"NL"}[candidate_wss[pi][1]],
                                y, old_prob*100, y, preds[pi], probs[pi][y].item()*100,
                                probs[pi][preds[pi]].item()*100))
                        succ = True
                        assert ws.update_ws(candidate_wss[pi][0], candidate_wss[pi][1])
                        break
                if succ:
                    break
                
                next_i = probs[:, y].argmin().item()
                # CBA-WS - Test if the ground truth probability decreases
                # SACBA-WS - Accept / reject to jump to the candidate
                accept = (probs[next_i, y] < old_prob)
                if _attack == 'SACBA' and (not accept):
                    acc_prob = torch.exp(-(probs[next_i, y] - old_prob) / temperature(it+1))
                    accept = (numpy.random.uniform(0,1) < acc_prob)
                if accept:
                    print ("  idx %d <%s>, %d (%.5f%%) => %d (%.5f%%)" % \
                           (candidate_wss[next_i][0],
                            {" ":"SPC", "\t":"TAB", "\n":"NL"}[candidate_wss[next_i][1]],
                            y, old_prob*100, y, probs[next_i][y].item()*100))
                    old_prob = probs[next_i][y].item()
                    assert ws.update_ws(candidate_wss[next_i][0], candidate_wss[next_i][1])
                else:
                    break
                
        # CodeBert-Attack (CMP Swapping)
        elif _attack in ['CBA-CMP', 'SACBA-CMP']:
            
            for it in range(n_perturb):
                # Find vulnerable identifiers
                vulnerables = atk.find_vulnerable_cmps(cls=cls_model,
                                                       cmp=cmp,
                                                       ground_truth_label=y,
                                                       n_vul=n_vulnerable)
                if len(vulnerables) <= 0:
                    break
                candidate_cmps, candidate_new_cmps, candidate_seqs = atk.generate_candidates(cmp=cmp,
                                                                                             mlm=mlm_model,
                                                                                             vulnerable=vulnerables,
                                                                                             batch_size=n_candidate,
                                                                                             n_candidate=n_candidate)
                candidate_old_cmps = [" ".join(cmp.code[cmp.cmptab[i][0]: cmp.cmptab[i][2]]) \
                                      for i in candidate_cmps]
                if len(candidate_seqs) <= 0:
                    break
                # Probe the target model
                probs = softmax(cls_model.run(candidate_seqs))
                preds = probs.argmax(dim=-1)
                for pi in range(len(preds)):
                    # Find an adversarial example
                    if preds[pi].item() != y:
                        print ("  %s => %s, %d (%.5f%%) => %d %d (%.5f%% %.5f%%)" % \
                               (candidate_old_cmps[pi], candidate_new_cmps[pi], y, old_prob*100,
                                y, preds[pi], probs[pi][y].item()*100, probs[pi][preds[pi]].item()*100))
                        succ = True
                        assert cmp.update_cmp(candidate_cmps[pi])
                        break
                if succ:
                    break
                
                next_i = probs[:, y].argmin().item()
                # CBA-CMP - Test if the ground truth probability decreases
                # SACBA-CMP - Accept / reject to jump to the candidate
                accept = (probs[next_i, y] < old_prob)
                if _attack == 'SACBA-CMP' and (not accept):
                    acc_prob = torch.exp(-(probs[next_i, y] - old_prob) / temperature(it+1))
                    accept = (numpy.random.uniform(0,1) < acc_prob)
                if accept:
                    print ("  %s => %s, %d (%.5f%%) => %d (%.5f%%)" % \
                           (candidate_old_cmps[next_i], candidate_new_cmps[next_i], y, old_prob*100,
                            y, probs[next_i][y].item()*100))
                    old_prob = probs[next_i][y].item()
                    assert cmp.update_cmp(candidate_cmps[next_i])
                else:
                    break
            
        if succ:
            n_succ += 1
            n_total += 1
            time_total += time.time() - time_st
            print ("  SUCC!")
        else:
            n_total += 1
            print ("  FAIL!")
            
        if n_total > 0:
            succ_rate = n_succ/n_total
        else:
            succ_rate = 0
        acc_rate = (n_total-n_succ)/n_total_including_originally_wrong
        if n_succ > 0:
            avg_time = time_total/n_succ
        else:
            avg_time = float("NaN")
        
        print ("  Succ %% = %.5f%%, Acc %% = %.5f%%, Avg time = %.5f sec" % \
               (succ_rate*100, acc_rate*100, avg_time))