from sklearn.linear_model import LinearRegression
import scipy.optimize as opt
from scipy import stats
from Android_malware.Scripts.utils import *

np.random.seed(1234)


FEANUM = 282515
n_fea_select = 300

TestNum = 100

def obj_fun(theta, x, y_):
    pre_dis = np.dot(x, theta)
    loss = np.sum((pre_dis - y_) ** 2)
    return loss



def perf_measure(y_true, y_pred):
    TP_FN = np.count_nonzero(y_true)
    FP_TN = y_true.shape[0] * y_true.shape[1] - TP_FN
    FN = np.where((y_true - y_pred) == 1)[0].shape[0]
    TP = TP_FN - FN
    FP = np.count_nonzero(y_pred) - TP
    TN = FP_TN - FP
    Precision = float(float(TP) / float(TP + FP + 1e-9))
    Recall = float(float(TP) / float((TP + FN + 1e-9)))
    accuracy = float(float(TP + TN) / float((TP_FN + FP_TN + 1e-9)))
    F1 =  2*((Precision * Recall) / (Precision + Recall))
    return Precision, Recall, accuracy



class xai_rnn(object):
    """class for explaining the rnn prediction"""
    def __init__(self, model, data,):
        """
        Args:
            model: target rnn model.
            data: data sample needed to be explained.
            label: label of the data sample.
            start: value of function start.
        """
        self.model = model
        self.data = data
        self.seq_len = data.shape[1]
        self.seq_len = data[0].shape[0]



    def xai_feature(self, samp_num, ):
        """extract the important features from the input data
        Arg:
            fea_num: number of features that needed by the user
            samp_num: number of data used for explanation
        return:
            fea: extracted features
        """
        # cen = int(self.seq_len/2)
        # half_tl = int(self.tl/2)
        sample = np.random.randint(0, FEANUM, samp_num)
        # # features_range = range(self.tl+1)
        # data_explain = np.copy(self.data).reshape(1, self.data.shape[1])
        # # data_sampled = np.copy(self.trunc_data_test)
        data_explain = np.random.randint(0, 2, [samp_num, FEANUM], dtype= np.bool)
        for i, size in enumerate(sample, start=0):
            inactive = np.random.choice(FEANUM, size, replace=False)
            tmp_sampled = np.copy(self.data)
            tmp_sampled[0,inactive] = np.random.choice(range(2), size, replace = True)     #1 - tmp_sampled[0,inactive]
            data_explain[i] = tmp_sampled
        #     #tmp_sampled = tmp_sampled.reshape(1, self.data.shape[1])
        #     data_explain = np.concatenate((data_explain, tmp_sampled), axis=0)
        #     data_sampled_mutate = np.copy(self.data)
        #
        #     data_sampled_mutate[0, 0:FEANUM] = tmp_sampled
        #
        #     data_sampled = np.concatenate((data_sampled, data_sampled_mutate),axis=0)
        # #
        # # if option == "Fixed":
        # #     print("Fix start points")
        # #     data_sampled[:, self.real_sp] = self.start

        # label_sampled = label_sampled.reshape(label_sampled.shape[0], 1)


        label_sampled = self.model.predict(data_explain, batch_size = 500, verbose = 0)



        # x = cp.Variable(n)
        # objective = cp.Minimize(cp.sum_squares(A * x - b))
        # constraints = [0 <= x, x <= 1]
        # prob = cp.Problem(objective, constraints)
        # result = prob.solve()
        # print(x.value)
        # print(constraints[0].dual_value)
        #


        linreg = LinearRegression(fit_intercept=False)
        linreg.fit(data_explain, label_sampled)
        result = np.argsort(np.abs(linreg.coef_)).reshape([FEANUM])
        importance_score = result[::-1]
        self.fea = importance_score
        # self.fea = self.fea[np.where(self.fea < 200)]
        # self.fea = self.fea[np.where(self.fea >= 0)]
        #print 'self.fea ...',self.fea
        return self.fea

    def gauss_mix_model(self, X, y, model_num, min_err=0.1):
        [data_num, fea_num] = np.shape(X)
        mean = np.random.random([model_num, 1]) * 100
        sigma = np.random.random([model_num, 1]) * 100
        pi = np.ones([model_num, 1]) / model_num

        new_mean = np.zeros_like(mean)
        new_sigma = np.zeros_like(sigma)
        new_pi = np.zeros_like(pi)

        new_z = np.zeros([data_num, model_num])
        err = np.linalg.norm(mean - new_mean)
        nk = np.zeros((model_num, 1))
        while (err > min_err):
            for index in range(data_num):
                for kindex in range(model_num):
                    new_z[index, kindex] = pi[kindex] * stats.norm(mean[kindex, 0], sigma[kindex, 0]).pdf(y[index])
                new_z[index, :] = new_z[index, :] / np.sum(new_z[index, :])
            for kindex in range(model_num):
                nk[kindex] = np.sum(new_z[:, kindex])
                new_pi[kindex, 0] = np.sum(new_z[:, kindex]) / data_num
                sum_mean = 0
                sum_sigma = 0
                for nindex in range(data_num):
                    tmp = new_z[nindex, kindex] * y[nindex]
                    sum_mean = sum_mean + tmp
                new_mean[kindex, :] = sum_mean / nk[kindex]

                for nindex in range(data_num):
                    vec = (y[nindex] - new_mean[kindex, 0])
                    tmp = new_z[nindex, kindex] * (vec * vec)
                    sum_sigma = sum_sigma + tmp
                new_sigma[kindex, 0] = sum_sigma / nk[kindex]
            err = np.linalg.norm(new_mean - mean)
            sigma = new_sigma
            pi = new_pi
            mean = new_mean

        weight = np.zeros(fea_num)
        s = 1e4
        init_k = np.ones([fea_num, 1])
        cons = ({'type': 'ineq',
                 'fun': lambda x: np.array([s - (x[0] - x[1]) ** 2 - (x[1] - x[2]) ** 2])})
        # for iindex in range(fea_num):
        #     for kindex in range(model_num):
        #         result = opt.minimize(obj_fun, init_k, args=(np.array(X), np.array(mean[kindex, :])),
        #                               constraints=cons)
        #         weight[iindex] = result.x[iindex] * pi[kindex] + weight[iindex]
        result = opt.minimize(obj_fun, init_k, args=(np.array(X), np.array(mean[kindex, :])),
                              constraints=cons)
        weight = result.x
######################################################################################
        # X = r.matrix(X, nrow=data_explain.shape[0], ncol=data_explain.shape[1])
        # Y = r.matrix(label_sampled, nrow=label_sampled.shape[0], ncol=label_sampled.shape[1])
        #
        # n = r.nrow(X)
        # p = r.ncol(X)
        # results = r.fusedlasso1d(y=Y, X=X)  # eps = 1e4
        #
        # result = np.array(r.coef(results, np.sqrt(n * np.log(p)))[0])[:, -1]

        return weight


def getRandomSample(x,y,num):
    #####################################################
    batchSize = int(len(x) / num)
    batindex = np.random.randint(0, num, size=[batchSize])
    for i in range(batchSize):
        batindex[i] = num * i + batindex[i]
    x = x[batindex]
    y = y[batindex]
    #####################################################
    return x, y


def LemnaRule(x_test):
    RuleSet = []
    for i in range(TestNum):
        tic = datetime.datetime.now()
        x = x_test[i: (i + 1)]
        xai_test = xai_rnn(model, x, )
        xai_fea = xai_test.xai_feature(500)
        if len(xai_fea) <= n_fea_select:
            selected_fea = np.array(xai_test.fea, dtype=np.int32)
        else:
            selected_fea = np.array(xai_test.fea[0:n_fea_select], dtype=np.int32)

        rule = []
        for j in range(len(selected_fea)):
            rule.append([selected_fea[j], xai_test.data[0, selected_fea[j]]])
        RuleSet.append(rule)
        toc = datetime.datetime.now()
        print("this is ", i, "rule", "cost time", toc - tic)
    f = open("RuleSet\\" + str(NowYear) + "_lemna.pkl", "wb")
    pickle.dump(RuleSet, f)
    f.close()


def DenasRule(x_test, puppetModel):
    RuleSet = []
    for i in range(len(x_test)):
        x = x_test[i : i + 1]
        print(i)
        rule = []
        activationState = calAcStateFromRule(rule, model)
        contributionVec = calContributionVec(puppetModel, activationState)
        while len(rule) <= n_fea_select:
            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:
                rule.append(newpt)
            contributionVec[pos] = 0
        RuleSet.append(rule)
    RuleSet = transferRuleSet(RuleSet)
    f = open("RuleSet\\" + str(NowYear) + "denas.pkl", "wb")
    pickle.dump(RuleSet, f)
    f.close()






def getFidelity():
    R1 = ReadRuleSet( fileName = "RuleSet\\" + str(NowYear) + "denas.pkl")
    R2 = ReadRuleSet( fileName = "RuleSet\\" + str(NowYear) + "_lemna.pkl")

    model = loadModel(year = NowYear)
    x, y = ReadData(year = NowYear)
    pred_y = model.predict(x, batch_size = 1000) > 0.5
    pos_x = x[np.where(pred_y > 0)[0]]
    neg_x = x[np.where(pred_y == 0)[0]]


    A = np.zeros([10, 2])
    B = np.zeros([10, 2])

    for k in range(10):
        tic = datetime.datetime.now()
        maxlength = k * 5 + 1
        for i in range(TestNum):
            rule1 = R1[i]
            rule1 = rule1[0 : maxlength]
            rule2 = R2[i]
            rule2 = rule2[0 : maxlength]
            x1 = pos_x[i:i + 1].copy()
            x2 = pos_x[i:i + 1].copy()

            xz1 = neg_x[0 : 1].copy()   #1 - x1  # np.zeros(np.shape(pos_x[i:i+1]))
            xz2 = neg_x[0 : 1].copy()   #1 - x2  # np.zeros(np.shape(pos_x[i:i+1]))
            for r in range(maxlength):
                p1 = rule1[r][0]
                v1 = rule1[r][1]
                x1[:, p1] = 1 - x1[:, p1]
                xz1[:, p1] = v1

                p2 = rule2[r][0]
                v2 = rule2[r][1]
                x2[:, p2] =  1 - x2[:, p2]
                xz2[:, p2] = v2
            A[k, 0] += (model.predict(x1) > 0.5)
            A[k, 1] += (model.predict(x2) > 0.5)
            B[k, 0] += (model.predict(xz1) > 0.5)
            B[k, 1] += (model.predict(xz2) > 0.5)
        toc = datetime.datetime.now()
        print(k, toc - tic)
    A = A / TestNum
    B = B / TestNum
    plt.plot(A[:, 0], 'r')
    plt.plot(A[:, 1], 'b')
    plt.show()
    plt.plot(B[:, 0], 'r')
    plt.plot(B[:, 1], 'b')
    plt.show()

    np.savetxt("data\\" + str(NowYear) + "_fidelity_A.csv", A, delimiter=',')
    np.savetxt("data\\" + str(NowYear) + "_fidelity_B.csv", B, delimiter=',')






if __name__ == "__main__":
    model = loadModel(year = NowYear)
    puppetModel = getPuppetModel("..\\model\\" + str(NowYear) + "\\MLP_model.h5")
    x, y = ReadData(year = NowYear, IsTrain = False)
    print('Data_num:', len(x))
    pred_y = (model.predict(x, batch_size = 1000) > 0.5)
    idx = np.nonzero(pred_y)[0]
    x_test = x[idx]
    x_test = x_test[0:TestNum]
    LemnaRule(x_test)
    # DenasRule(x_test, puppetModel)
    getFidelity()