import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from Pdf_malware.Scripts.LEMNA import xai_rnn
from Pdf_malware.Scripts.utils import *
import matplotlib.pyplot as plt
import innvestigate




class FidelityMetric():
    def __init__(self, data, model, important, maxselNum, neg_x, step = 1):
        self.data = data.copy()
        self.important = important
        self.maxselNum = maxselNum
        self.neg_x  = neg_x
        self.step = step
        self.iter = int(maxselNum / step)
        self.model = model


    def AugmentTest(self):
        AugmentRes = []
        for i in range(self.iter):
            testdata = self.neg_x.copy()
            xpos = np.arange(0, len(self.data))
            for j in range(i * self.step):
                pos = np.int32(self.important[:, j])
                testdata[xpos, pos] = self.data[xpos, pos]
            AugmentRes.append(np.sum(self.model.predict(testdata) > 0.5) / len(self.data))
        return AugmentRes


    def DeductionTest(self):
        DeductionRes = []
        for i in range(self.iter):
            testdata = self.data.copy()
            xpos = np.arange(0, len(self.data))
            for j in range(i * self.step):
                pos = np.int32(self.important[:, j])
                testdata[xpos, pos] = 1 - self.data[xpos, pos]
            DeductionRes.append(np.sum(self.model.predict(testdata) > 0.5)/ len(self.data))
        return DeductionRes





class FidelityTest():
    def __init__(self, model, puppetmodel, x, y, neg_x, testNum = 100, selNum = 10, step = 1):
        self.model = model
        self.puppetmodel = puppetmodel
        self.x = x[0 : testNum]
        self.y = y[0 : testNum]
        self.testNum = testNum
        self.selNum = selNum
        self.neg_x = neg_x[0 : testNum]
        self.step = step
        self.baseline = [
            "gradient", #gradient
            "integrated_gradients",
            "deep_taylor",
        ]

    def Denasfidelity(self, ):
        importantInx = []
        for i in range(self.testNum):
            x = self.x[i: i + 1]
            rule = []
            importantpos = []
            while len(rule) <= self.selNum:
                activationState = calAcStateFromRule(rule, self.model)
                contributionVec = calContributionVec(self.puppetmodel, activationState)
                for pt in rule:
                    contributionVec[int(pt/2)] = 0
                add = 0
                while add <= self.step:
                    newpt = np.argmax(np.abs(contributionVec))
                    if contributionVec[newpt] > 0:
                        newpt = newpt * 2 + 1
                    else:
                        newpt = newpt * 2
                    pos = int(newpt / BIT)
                    val = newpt % BIT
                    if x[0, pos] == val:
                        importantpos.append(pos)
                        rule.append(newpt)
                        add += 1
                    contributionVec[pos] = 0
            importantInx.append(np.array(importantpos))
        importantInx = np.array(importantInx)
        print("denas explain finished")

        RuleSet = []
        for i in range(len(self.x)):
            rule = [[j, self.x[i, j]] for j in importantInx[i]]
            RuleSet.append(rule)
        f = open('../RuleSet/denas.pkl', 'wb')
        pickle.dump(RuleSet, f)
        f.close()



        metric = FidelityMetric(self.x, self.model, importantInx, self.selNum, self.neg_x, step = self.step)
        a = metric.AugmentTest()
        b = metric.DeductionTest()
        return a, b


    def Lemnafidelity(self,):
        importantInx = np.zeros_like(self.x)
        for i in range(self.testNum):
            x_test = self.x[i: (i + 1)]
            xai_test = xai_rnn(self.model, x_test)
            importantInx[i] = np.array(xai_test.xai_feature(500))
        print("lemna finish extract explanation")

        RuleSet = []
        for i in range(len(self.x)):
            rule = [[j, self.x[i, j]] for j in importantInx[i]]
            RuleSet.append(rule)
        f = open('../RuleSet/lemna.pkl', 'wb')
        pickle.dump(RuleSet, f)
        f.close()

        metric = FidelityMetric(self.x, self.model, importantInx, self.selNum, self.neg_x, step = self.step)
        a = metric.AugmentTest()
        b = metric.DeductionTest()
        return a, b



    def Baselinefidelity(self, i_num, num = 2):
        analysis = np.zeros_like(self.x, dtype=np.float32)
        step = int(self.testNum / num)
        analyzer= innvestigate.create_analyzer(self.baseline[i_num], self.model)
        for i in range(num):
            st = int((i) * step)
            ed = int((i + 1) * step)
            analysis[st : ed] = analyzer.analyze(self.x[st : ed])

        importIndex = np.argsort(analysis * -1, axis= 1)
        print(self.baseline[i_num], "finish explanation")



        RuleSet = []
        for i in range(len(self.x)):
            rule = [[j, self.x[i, j]] for j in importIndex[i]]
            RuleSet.append(rule)
        f = open('../RuleSet/' + self.baseline[i_num] + '.pkl', 'wb')
        pickle.dump(RuleSet, f)
        f.close()


        metric = FidelityMetric(self.x, self.model, importIndex, self.selNum, self.neg_x)
        a = metric.AugmentTest()
        b = metric.DeductionTest()
        return a, b





def getFidelityRes(fid ):
    baselineNum = 3
    AugplotArr = np.zeros([10, baselineNum + 2])
    DecplotArr = np.zeros([10, baselineNum + 2])
    x_axis = np.arange(0, 10, 1)

    for i in range(baselineNum):
        AugplotArr[:, i], DecplotArr[:, i] = fid.Baselinefidelity(i)
    AugplotArr[:, baselineNum], DecplotArr[:, baselineNum] = fid.Lemnafidelity()
    AugplotArr[:, baselineNum + 1], DecplotArr[:, baselineNum + 1] = fid.Denasfidelity()

    print('AugplotArr')
    print(AugplotArr)
    print('DecplotArr')
    print(DecplotArr)
    np.savetxt("../Results/AugplotArr.csv", AugplotArr, delimiter=',')
    np.savetxt("../Results/DecplotArr.csv", DecplotArr, delimiter=',')

    name = ["gradient", "ig", "deeptaylor", 'lemna', 'denas']
    L = []
    for i in range(baselineNum + 2):
        l,  = plt.plot(x_axis, AugplotArr[:, i], label=str(i))
        L.append(l)
    plt.legend(handles=L, labels=name)
    plt.savefig('../Results/AugplotArr.png')
    plt.cla()
    #plt.show()

    name = ["gradient", "ig","deeptaylor", 'lemna', 'denas']
    L = []
    for i in range(baselineNum + 2):
        l, = plt.plot(x_axis, DecplotArr[:, i], label=str(i))
        L.append(l)
    plt.legend(handles=L, labels=name)
    plt.savefig('../Results/DecplotArr.png')
    #plt.show()








def main():
    model = load_model("../model/MLP_model.h5")
    x, y = loaddata("../data/train_model_data.pkl")

    index = np.random.choice(np.arange(len(x)), 3000)
    x = x[index]
    y = y[index]
    np.save('../Results/x', x)
    np.save('../Results/y', y)
    pred_y = model.predict(x, batch_size=5000) > 0.5
    pos_x = x[np.where(pred_y > 0)[0]]
    neg_x = x[np.where(pred_y == 0)[0]]
    puppetModel = getPuppetModel("../model/MLP_model.h5")
    fid = FidelityTest(model, puppetModel, pos_x, y, neg_x, testNum= 100,  step = 1)
    getFidelityRes(fid)
    print("finish accuracy experiment")

if __name__ == "__main__":
    main()









