import random
import numpy as np
import networkx as nx
from sklearn.linear_model import LinearRegression


def generate_Z(S):
    """
    S: number of clusters/nodes in the graph
    This function will randomly select half of the clusters/nodes as control, and the other half as variant
    """
    l = random.sample(range(S), int(S/2))
    Zc = [1 if i in l else 0 for i in range(S)]
    return Zc
    
    
def generate_Y(lambda0, lambda1, lambda2, G, shardMap, Zc, T=3):
    """
    lambda0: baseline effect;
    lambda1: treatment effect;
    lambda2: network effect;
    G: graph;
    shardMap: a numpy array with shape (num of nodes, ), each the cluster index for each node;
    Zc: a vector: treatment assignment at cluster level (len(Z) = num of clusters)
    t: number of steps
    """
    ## obtain assignment at individual level
    num_nodes = len(G.nodes)
    Z = np.array([Zc[shardMap[i]] for i in range(num_nodes)])
    Z = Z.reshape(-1,1)
    Y = np.zeros([num_nodes, 1])
    ## obtain adjacancy matrix: 
    Adj = nx.adjacency_matrix(G).todense()
    ## obtain degree of vertexes
    DegreeV = np.array(G.degree)[:,1].reshape([-1,1])
    for t in range(T):
        E = np.random.normal(size=num_nodes).reshape([-1,1])
        Ys = lambda0 + lambda1 * Z + lambda2 * np.matmul(Adj, Y) / DegreeV + E
        Y = (Ys > 0) + 0
    return Y
    
def estimators(Y, Z, Adj, DegreeV, q=0.7):
    """
    Y: numpy array with shape [num_nodes, 1], response for all nodes;
    Z: numpy array with shape [num_nodes, 1], treatment assignment at individual levels;
    G: graph;
    q: threshold for q-neighborhood exposed to treatment/control
    """
    ## obtain naive estimator:
    tau_naive = (np.sum(Y*Z)+0.0)/np.sum(Z) - (np.sum(Y*(1-Z))+0.0)/np.sum(1-Z)
    
    ## obtain neighbor estimator:
    # first compute percentage of neighbors in trt group for each node
    # sigma is the percent of i’s neighbors in treatment, i between 1 and num_nodes
    sigma = np.matmul(Adj, Z)/DegreeV
    # print('sigma:')
    # print(sigma)
    Nbor_exposure_diff = np.abs(Z - sigma)
    q_nbor_exposure = (Nbor_exposure_diff < (1-q)) + 0
    tau_neighbor = (np.sum(Y*Z*q_nbor_exposure) + 0.0)/(np.sum(Z*q_nbor_exposure)+0.0001) - (np.sum(Y*(1-Z)*q_nbor_exposure) + 0.0)/(np.sum((1-Z)*q_nbor_exposure)+0.0001)
    ## obtain lm estimator:
    ## linear formula: Y ~  beta0 + beta1*Z + beta2*sigma
    X_list = np.concatenate((Z, sigma), axis=1)
    ## fit linear regression to obtain beta0,beta1,beta2
    reg = LinearRegression().fit(X_list, Y)
    ## lm estimator = beta1_hat + beta2_hat
    tau_lm1 = np.sum(reg.coef_)
    
    return tau_naive, tau_neighbor, tau_lm1

def compute_p(G, shardMap, Y, Zc, S, M=1000):
    Adj = nx.adjacency_matrix(G).todense()
    DegreeV = np.array(G.degree)[:,1].reshape([-1,1])

    num_nodes = len(G.nodes)
    ## obtain assignment at individual level
    Z = np.array([Zc[shardMap[i]] for i in range(num_nodes)])
    Z = Z.reshape(-1,1)
    tau_naive, tau_neighbor, tau_lm1 = estimators(Y, Z, Adj, DegreeV)
    ## permutation test at individual level:
    Tau_stat_indiv = np.zeros([M, 3])
    for m in range(M):
        Z_m = generate_Z(num_nodes)
        Z_m = np.array(Z_m).reshape([-1,1])
        tau_naive_m, tau_neighbor_m, tau_lm1_m = estimators(Y, Z_m, Adj, DegreeV)
        Tau_stat_indiv[m,:] = [tau_naive_m, tau_neighbor_m, tau_lm1_m]
    ## permutation test at cluster level:
    Tau_stat_cluster = np.zeros([M, 3])
    for m in range(M):
        Zc_m = generate_Z(S)
        Z_m = [Zc_m[shardMap[i]] for i in range(num_nodes)]
        Z_m = np.array(Z_m).reshape([-1,1])
        tau_naive_m, tau_neighbor_m, tau_lm1_m = estimators(Y, Z_m, Adj, DegreeV)
        Tau_stat_cluster[m,:] = [tau_naive_m, tau_neighbor_m, tau_lm1_m]

    p_naive_i    = sum(tau_naive <  Tau_stat_indiv[:, 0])/M
    p_neighbor_i = sum(tau_neighbor < Tau_stat_indiv[:, 1])/M
    p_lm_i       = sum(tau_lm1      < Tau_stat_indiv[:, 2])/M

    p_naive_c =    sum(tau_naive <  Tau_stat_cluster[:, 0])/M
    p_neighbor_c = sum(tau_neighbor < Tau_stat_cluster[:, 1])/M
    p_lm_c       = sum(tau_lm1      < Tau_stat_cluster[:, 2])/M
    
    return p_naive_i, p_neighbor_i, p_lm_i, p_naive_c, p_neighbor_c, p_lm_c
