import time
import torch
import torch.nn as nn
import rnn_tools
import rnn_model
from torch.autograd import Variable
from copy import deepcopy
import pickle
import numpy as np
from math import pow
from itertools import combinations

TAU = 0.99999999
INDICE_K = 6
MODEL_TYPE = 'nonsub'
SECONDS = 180

TITLE = '=== ' + MODEL_TYPE + ' target prob = ' + str(TAU) + ' k = ' \
        + str(INDICE_K) + ' time = ' + str(SECONDS) + ' ==='

NUM_ATTACK_SAMPLES = 1000

log_f = open('./Logs/GB_%s_k=%d_t=%s_s=%d.bak' % (MODEL_TYPE, INDICE_K, str(TAU), SECONDS), 'w+')


class Attacker(object):
    def __init__(self, options, emb_weights):
        print("Loading pre-trained classifier...", file=log_f, flush=True)

        self.model = rnn_model.LSTM(options, emb_weights)
        self.model2 = rnn_model.LSTM(options, emb_weights)

        if MODEL_TYPE == 'sub':
            self.model.load_state_dict(torch.load('./Classifiers/Submodular_lstm.7')) # abs
            self.model2.load_state_dict(torch.load('./Classifiers/Submodular_lstm.7'))
        elif MODEL_TYPE == 'nonsub':
            self.model.load_state_dict(torch.load('./Classifiers/Nonsubmodular_lstm.7')) # positive and negative
            self.model2.load_state_dict(torch.load('./Classifiers/Nonsubmodular_lstm.7'))

        self.model.eval()
        self.model2.eval()

        self.criterion = nn.CrossEntropyLoss()

    def classify(self, funccall):

        model_input, weight_of_embed_codes = self.input_handle(funccall)

        logit = self.model2(model_input, weight_of_embed_codes)

        pred = torch.max(logit, 1)[1].view((1,)).data.numpy()

        prob = logit[0][pred]

        return pred, prob

    def input_handle(self, funccall):
        funccall = [funccall]
        t_diagnosis_codes = rnn_tools.pad_matrix(funccall)
        model_input = deepcopy(t_diagnosis_codes)
        for i in range(len(model_input)):
            for j in range(len(model_input[i])):
                idx = 0
                for k in range(len(model_input[i][j])):
                    model_input[i][j][k] = idx
                    idx += 1

        model_input = Variable(torch.LongTensor(model_input))
        return model_input, torch.tensor(t_diagnosis_codes)

    def forward_lstm(self, weighted_embed_codes, model):
        x = model.relu(weighted_embed_codes)
        x = torch.mean(x, dim=2)
        h0 = Variable(torch.FloatTensor(torch.randn(1, x.size()[1], x.size()[2])))
        c0 = Variable(torch.FloatTensor(torch.randn(1, x.size()[1], x.size()[2])))
        output, h_n = model.lstm(x, (h0, c0))
        embedding, attn_weights = model.attention(output.transpose(0, 1))
        x = model.dropout(embedding)  # (n_samples, hidden_size)

        logit = model.fc(x)  # (n_samples, n_labels)

        logit = model.softmax(logit)
        return logit

    def attack(self, funccall, y):
        st = time.time()

        success_flag = 1

        one_hot_codes = np.zeros((len(funccall), 1104))
        for i in range(len(funccall)):
            one_hot_codes[i][funccall[i]] = 1

        orig_pred, orig_prob = self.classify(funccall)
        print(orig_pred, orig_prob, file=log_f, flush=True)

        pred, pred_prob = orig_pred, orig_prob

        if not (pred == y or pred_prob < TAU):
            success_flag = 0
            return pred_prob, {}, {}, success_flag, 0

        maxCy = 1 - pred_prob

        S_t = set()
        S_t_timestep = set()
        time_step = -1

        new_code = (-1,-1)
        Code_S_t = []

        iteration = 0

        first_time_flag = 1

        while (pred == y or maxCy < TAU):

            iteration += 1

            persons = []

            flip_num = int(pow(2, len(S_t)))

            selected_codes = []

            if first_time_flag:
                flip_num = 1
                persons.append(funccall)
                first_time_flag = 0
            else:
                for i in range(len(S_t) + 1):
                    comb_set = combinations(list(S_t), i)
                    for s in comb_set:
                        selected_codes.append(list(s))
                if selected_codes:
                    for c in selected_codes:
                        c.append(new_code)
                else:
                    selected_codes.append([new_code])

                for c in selected_codes: # [[C], [A, C], [B, C], [A, B, C]]
                    temp_person = deepcopy(funccall)
                    for code in c:
                        temp_person[code[0]] = code[1]
                    persons.append(temp_person)

                S_t.add(new_code)
                S_t_timestep.add(time_step)

            temp_vs = []
            temp_Cys = []
            map_temp_v_subset = {}

            for i in range(flip_num):
                model_input, weight_of_embed_codes = self.input_handle(persons[i])
                embed_codes = self.model.embed(torch.LongTensor(model_input))

                weight_of_embed_codes = Variable(weight_of_embed_codes.data, requires_grad=True)

                weighted_embed_codes = embed_codes * torch.unsqueeze(weight_of_embed_codes, dim=3)

                output = self.forward_lstm(weighted_embed_codes, self.model)

                loss = self.criterion(output, Variable(torch.LongTensor([y])))

                loss.backward()

                score_time_step = np.zeros(20)

                for j in range(20):
                    if j in S_t_timestep: continue
                    a = weight_of_embed_codes.grad.data[j][0][funccall[j]]
                    score_time_step[j] = abs(a.data)

                indices = np.argsort(-np.reshape(score_time_step, (-1)))[:INDICE_K]
                indices_pair = []

                for t in indices:
                    score_features = np.zeros(1104)
                    for feature_idx in range(1104):
                        if feature_idx == funccall[t] or feature_idx == 0:
                            score_features[feature_idx] = 0
                            continue
                        a = weight_of_embed_codes.grad.data[t][0][feature_idx]
                        score_features[feature_idx] = abs(a.data)
                    f = np.argmax(score_features)
                    if f == 0:
                        f += 1
                    indices_pair.append((t, f))

                temp_v, temp_Cy = self.code_paraphrase(persons[i], indices_pair, y)

                temp_vs.append(temp_v)
                temp_Cys.append(temp_Cy)
                map_temp_v_subset[temp_v] = i

            best_temp_Cy_idx = np.argmax(temp_Cys)
            best_temp_Cy = np.max(temp_Cys)
            new_code = temp_vs[best_temp_Cy_idx]
            time_step = new_code[0]

            if best_temp_Cy > maxCy:
                maxCy = best_temp_Cy
                if selected_codes:
                    Code_S_t = selected_codes[map_temp_v_subset[new_code]]
                    Code_S_t.append(new_code)
                else:
                    Code_S_t = [new_code]

            if best_temp_Cy > 0.5:
                pred = 1-y

            if (time.time() - st) > SECONDS and (pred == y or pred_prob < TAU):
                print("===== Time out! Attack Fail =====", file=log_f, flush=True)
                success_flag = -1
                break

        S_t.add(new_code)

        if len(Code_S_t) == 0:
            success_flag = -1

        return maxCy, S_t, Code_S_t, success_flag, iteration

    def code_paraphrase(self, funccall, indices_pair, y):
        temp_v_index = -1
        temp_Cy = -1
        for pos in indices_pair:
            candidate = deepcopy(funccall)

            candidate[pos[0]] = pos[1]

            model_input, weight_of_embed_codes = self.input_handle(candidate)
            pred_prob = self.model(model_input, weight_of_embed_codes)
            if (1-pred_prob[0][y]) > temp_Cy:
                temp_Cy = (1-pred_prob[0][y])
                temp_v_index = pos
        return temp_v_index, temp_Cy


def main(emb_weights, training_file, validation_file,
                                       n_diagnosis_codes, n_labels,
                                       batch_size, dropout_rate,
                                       L2_reg, n_epoch, log_eps, n_claims, visit_size, hidden_size,
                                       model_name):
    options = locals().copy()
    print("Loading dataset...", file=log_f, flush=True)
    test = rnn_tools.load_data(validation_file)

    n_samples = NUM_ATTACK_SAMPLES

    attacker = Attacker(options, emb_weights)

    n_success = 0
    n_fail = 0

    total_node_change = 0

    n_iteration = 0

    saving_time = {}

    attack_code_dict = {}

    for i in range(n_samples):
        print("-------- %d ---------" % (i), file=log_f, flush=True)

        sample = test[0][i]

        y = test[1][i]

        print('* Processing:%d/%d samples' % (i, n_samples), file=log_f, flush=True)

        print("* Original: "+str(sample), file=log_f, flush=True)

        print("  Original label: %d" % (y), file=log_f, flush=True)

        st = time.time()

        maxCy, S_t, Code_S_t, success_flag, iteration = attacker.attack(sample, y)

        et = time.time()
        all_t = et-st

        if success_flag == 1:
            n_success += 1
            n_iteration += iteration
            total_node_change += len(Code_S_t)
            saving_time[i] = all_t
            attack_code_dict[i] = Code_S_t
        elif success_flag == -1:
            n_fail += 1

        print("* Result: ", file=log_f, flush=True)

        print(Code_S_t, file=log_f, flush=True)
        print(S_t, file=log_f, flush=True)
        print(maxCy, file=log_f, flush=True)

        print("  Nnumber of changed codes: %d" % (len(Code_S_t)), file=log_f, flush=True)

        print("  Number of iterations for this: " + str(iteration), file=log_f, flush=True)

        print(" Time: "+str(all_t), file=log_f, flush=True)

        print("* SUCCESS Number NOW: %d " % (n_success), file=log_f, flush=True)
        print("* Failure Number NOW: %d " % (n_fail), file=log_f, flush=True)

        if n_success:
            print("  Average Number of success changed codes: " + str(float(total_node_change)/float(n_success)), file=log_f, flush=True)
            print("  Average Number of success iterations: " + str(float(n_iteration) / float(n_success)), file=log_f, flush=True)

        if i % 20 == 0:
            pickle.dump(attack_code_dict, open('./AttackCodes/codes_GB_%s_k=%d_t=%s_s=%d.pickle' % (MODEL_TYPE, INDICE_K, str(TAU), SECONDS), 'wb'))
            pickle.dump(saving_time, open('Each_TimeNumber/time_GB_%s_k=%d_t=%s_s=%d.pickle' % (MODEL_TYPE, INDICE_K, str(TAU), SECONDS), 'wb'))

    pickle.dump(attack_code_dict,
                open('./AttackCodes/codes_GB_%s_k=%d_t=%s_s=%d.pickle' % (MODEL_TYPE, INDICE_K, str(TAU), SECONDS), 'wb'))
    pickle.dump(saving_time,
                open('Each_TimeNumber/time_GB_%s_k=%d_t=%s_s=%d.pickle' % (MODEL_TYPE, INDICE_K, str(TAU), SECONDS), 'wb'))

    print("--- Total Success Number: " + str(n_success) + " ---", file=log_f, flush=True)
    print("--- Total Attack Success Rate: " + str(float(n_success) / float(n_success + n_fail)), file=log_f, flush=True)

    print(TITLE, file=log_f, flush=True)
    print(TITLE)


if __name__ == '__main__':
    print(TITLE)
    print(TITLE, file=log_f, flush=True)
    # parameters
    batch_size = 5
    dropout_rate = 0.2
    L2_reg = 0.001  # 0.001
    log_eps = 1e-8
    n_epoch = 50
    n_labels = 3  # binary classification
    visit_size = 70
    hidden_size = 70
    n_diagnosis_codes = 1104
    n_claims = 504

    # use_gpu = True
    model_name = 'lstm'

    training_file = './SourceData/tri_sub0.01_training_malware_index_data.pickle'
    validation_file = './SourceData/tri_sub0.003_validation_malware_index_data.pickle'

    emb = torch.load("./SourceData/PretrainedEmbedding1104.0", map_location=torch.device('cpu'))['embeddings.weight']

    emb_weights = torch.tensor(emb, dtype=torch.float)

    # ma = torch.max(emb_weights[:1104])
    # mi = torch.min(emb_weights[:1104])

    main(emb_weights, training_file, validation_file,
                n_diagnosis_codes, n_labels,
                batch_size, dropout_rate,
                L2_reg, n_epoch, log_eps, n_claims, visit_size, hidden_size,
                model_name)

