"""
This file performs the following tasks:
    1) Generate random DAGs using the function: generate_graph()
    2) Generate AR model coefficients for the random DAGs using the function: create_A_B_coeffs_6_nodes()
    3) Generate data using the AR models using teh function: continuous_data_generation_arbitrary_lags()
    4) Use the data and the three algorithms FD based W-PC, TD based W-PC, and 
       Winer-phase to tune thresholds at each specified sample size using the functions: tune_freq_threshold_deg(),
       tune_time_threshold_deg(), and tune_wiener_phase_threshold_deg()       
    5) Run the three algorithms at specified sample sized for each of the networks 
       with chosen thresholds to calculate error rates and computation time using the functions: compute_PC_error_FFT_avg(),
       compute_time_PC_error(), and compute_Wiener_phase_error_avg()
    6) Plot the average error rate and average computation time at each sample size.
"""

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import time
from itertools import combinations
import random



def generate_graph(num_nodes,max_in_degree):
    # generates a random dag with maximum in degree equal to max_in_degree
    G = nx.DiGraph()
    nodes = list(range(num_nodes))
    G.add_nodes_from(node for node in nodes)
    for node in nodes:
        # in_degree = max_in_degree 
        in_degree = random.randint(1, max_in_degree)
        if node - in_degree < 0:
            in_degree = node  # Avoid in-degrees exceeding node indices
        in_neighbors = random.sample(nodes[:node], in_degree)
        G.add_edges_from((neighbor, node) for neighbor in in_neighbors)
    Adj = nx.adjacency_matrix(G).toarray().T
    print("Number of nodes:", G.number_of_nodes())
    print("Number of edges:", G.number_of_edges())
    # print("Adjacency matrix:", Adj.shape," \n", Adj)
    return G,Adj

    

# Find moral graph and skeleton for a given graph given its adjacency matrix
# Input:  (1) adjacency_matrix: Adjacency matrix of directed graph
# Output: (1) MG: Adjacency matrix of the estimated moral graph
#         (2) Skeleton: Adjacency matrix of the estimated skeleton
def find_moral_graph_and_skeleton_from_adjacency_matrix(Adj):
    nNodes=Adj.shape[0]
    Skeleton=(Adj+Adj.T)
    MG=np.zeros([nNodes,nNodes],dtype=bool)
    MG=np.copy(Skeleton)
    for ind1 in range(nNodes):
        for ind2 in range(nNodes):
            for ind3 in range(nNodes):
                if (Adj[ind1,ind2]==1) and ((Adj[ind1,ind3]==1)) and ind2!=ind3:
                    MG[ind2,ind3]=1
                    MG[ind3,ind2]=1
    return MG,Skeleton




def create_A_B_coeffs_6_nodes(Adj,filter=None):
    # Create A and B matrix for IIR transfer functions
    # A is coeffs of interaction with the past of self, 
    # B interaction with the other nodes 
    # returns A,B
    nNodes=Adj.shape[0]    
    A = np.array([[1.739381261594696859e-02,	1.043628756956818150e-02,	5.218143784784090751e-03],
                  [3.415208172002098808e-01,	1.366083268800839523e-01,	6.830416344004197615e-02],
                  [8.289102927104288199e-01,	5.802372048973001295e-01,	3.315641170841715502e-01],
                  [8.846579905691620560e-02,	4.423289952845810280e-02,	2.653973971707486099e-02],
                  [8.562817629518102436e-01,	5.137690577710861684e-01,	2.568845288855430842e-01],
                  [1.630296867565628194e-01,	1.141207807295939597e-01,	4.890890602696884581e-02]])

    B=np.random.uniform(0.2,0.4,size=(nNodes,nNodes,3))
    Adj=np.repeat(Adj[:, :, np.newaxis], 3, axis=2)
    B=B*Adj
    B[:,:,1:3]=0
    return A,B



# Generate data according to the VAR model
# Inputs: (1) nSamples: Total number of samples
#         (2) B: p x p x DELAY-1 across other variables
#         (3) A: Matrix with coefficients representing self dynamics
#         (4) noise_pow: Noise power (noise is i.i.d. Gausian)
#         (5) nlags: required lag on other variables
# Output: (1) Timeseries data for a generative graph entailed by A adn B arrays
def continuous_data_generation_arbitrary_lags(nSamples,B,A=None,noise_pow=-1,nlags=0):
    nNodes = B.shape[0]
    if noise_pow.all() == -1:
        noise_pow = np.ones(B.shape[0])
    print("Generating continous data.. nSamples: ",nSamples," for nNodes :",nNodes)
    x = np.zeros((nSamples,nNodes))
    for ind in range(nlags):
        x[ind,:] = np.random.randn(nNodes)*np.sqrt(noise_pow)
    for ind_Samp in range(nlags,nSamples):
        x[ind_Samp,:] = np.random.randn(nNodes)*np.sqrt(noise_pow) - A[:,0]*x[ind_Samp-1,:] - A[:,1]*x[ind_Samp-2,:] - A[:,2]*x[ind_Samp-3,:]
        x[ind_Samp,:] = x[ind_Samp,:] + np.dot(B[:,:,0],x[ind_Samp-nLags,:])
    return x



# Compute least squates solution: Project X on to Z, Z=a0x_0+a1x1+...am xm, where m is the number of columns of X
def Least_Squares_Solution(X,Z):
    X_star = X.conj().T
    XX=np.matmul(X_star, X)
    if np.isscalar(XX):
        return np.matmul(np.divide( X_star,XX), Z)
    else:
        return np.matmul(np.matmul(np.linalg.inv(np.matmul(X_star, X)), X_star), Z)



# Estimate the Wiener filters in time domain.
# Inputs: (1) x: The data set to be used for calculating projection
#         (2) i: The node on which to be projected
#         (3) L: The number of past and future samples to be used for calculating the projections
# Output: (1) w: the computed Wiener projection coefficients.
def Wiener_time(x,i,L):
    # ||xi-YB||^2
    m = (2*L+1)
    T = x.shape[0]-1
    xi = x[L:T-L+1,i]
    C = [ind for ind in range(x.shape[1])]
    C = np.delete(C,i)
    y = np.zeros([T-2*L+1,m*C.shape[0]])
    for ind in range(L,T-L+1,L):
        y[ind-L,:] = x[ind-L:ind+L+1,C].reshape(m*len(C))
    w1 = Least_Squares_Solution(y,xi)
    # w1 = np.linalg.lstsq(y,xi)[0]
    w = w1.reshape(m,len(C))
    return w



# Function for estimating the Wiener filters in time domain.
# Inputs:  (1) data: time-series data.
#          (2) i: First node.
#          (3) j: Second node.
#          (4) z: Conditioning set of nodes.
# Outputs: (1) W: The estimated Wiener filters
#          (2) W_sum: Sum of all the time domain Wiener coefficients
def compute_partial_Wiener_coeffs_time(data,i,j,z=[],nlags=0):
    c=np.append(j,np.int32(z))
    indexes = np.unique(c, return_index=True)[1]
    c = [c[index] for index in sorted(indexes)]
    X = data[:,np.append(i,c)] # Wiener_time computes projection of x_i on x_ibar. So need to pass x_{i,c}
    W = Wiener_time(X,0,nlags)
    W_sum = np.sum(abs(W),axis=0)
    W_time_avg = W_sum/W.shape[0]
    return W,W_time_avg



# Algorithm for estimatin the d-separation sets using Wiener projections estimated in time domain.
# Inputs:  (1) data: time-series data
#          (2) nlags: The value (L) of the number of samples from past and future to be used in Wiener filter estimation
#          (3) thresh: Thresold for comparing the Wiener filter extimates to check for d-separation
#          (4) q: The maximum number of elements in the conditioning set
# Outputs: (1) C: A dictionary containing node pairs as key and the corresponding d-separation sets as the elements
def Run_PC_Wiener_test_time(data,nlags,thresh,Q):
    # print("Run_PC_Wiener_test_time")
    nNodes=data.shape[1]
    C={}
    for ind1 in range(nNodes):
        for ind2 in range(nNodes):
            if ind1>ind2:
                # continue
                flag=False
                # the complement set through which we condition
                D=[i for i in range(nNodes) if i != ind1 and i !=ind2]
                # iterate over the combinations of various size in the increasing order
                for ind_cond in range(Q+1):
                    if flag==True:
                        break
                    combination_set=[i for i in combinations(D,ind_cond)]
                    for c in combination_set:
                        # print("\n combination set: ",c)
                        W_i_jZ,W_sum_i_jZ=compute_partial_Wiener_coeffs_time(data,ind1,ind2,c,nlags)
                        # print("Projection coefficient of X"+str(ind1+1)+" on "+str(np.append(ind2,c)+1)+"=",W_sum_i_jZ)
                        if abs(W_sum_i_jZ[0])<thresh: #.1
                            # print("Separation set of ( X"+str(ind1+1)+", X"+str(ind2+1),") is : ",np.array(c)+1)
                            C[ind1,ind2]=c
                            flag=True
                            break
                            c=np.array([ind1])
                        if np.shape(c)[0]==min(Q,nNodes-2): #checking if |c|==q, which means |{c,q}|>q
                            # print("No separation set for ( X"+str(ind1+1)+", X"+str(ind2+1),")") #%%
                            C[ind1,ind2]=None
                            break

    return C




# Compute the FFT of the data
# Inputs:  (1) data: time-series data
#          (2) nfft: numer of FFT points
# Output:  (1) y: the FFT of the provided data samples
def compute_fft(data,nfft):
    nSamples,nNodes=data.shape
    nTrajectories=np.int32(nSamples/nfft) 
    y=np.zeros((nTrajectories,nNodes,nfft),dtype=complex)
    for ind in range(nTrajectories-1): # discard final few residual samples
        x=data[ind*(nfft):(ind+1)*nfft,:]
        X=np.fft.fft(x,axis=0) #/np.sqrt(nfft)
        y[ind,:,:]=np.transpose(X) #X is (nfft,nNodes)
    return y



def Wiener_Proj(X,Z):
    # Project X on to Z, Z=a0x_0+a1x1+...am xm, 
    # where m is the number of columns of X
    X_star = X.conj().T
    XX=np.matmul(X_star, X)
    if np.isscalar(XX):
        return np.matmul(np.divide( X_star,XX), Z)
    else:
        return np.matmul(np.matmul(np.linalg.inv(np.matmul(X_star, X)), X_star), Z)



def compute_Wiener_coeffs(nNodes,data):
    # Returns W and W_mag
    # W: nNodes X nSamples X nNodes
    # W_mag: returns h-infinity norm over frequencies
    nSamples=data.shape[2]
    W=np.zeros([nNodes,nSamples,nNodes],dtype='complex')
    W_mag=np.zeros([nNodes,nNodes])
    for proj_ind in range(nNodes):
        W1=np.zeros([nSamples,nNodes],dtype='complex')
        for freq_index in range(nSamples):
            fft_data=data[:,:,freq_index]
            Z=fft_data[:,proj_ind]
            X_bar=np.delete(fft_data,proj_ind,1)
            W1[freq_index,:]=np.insert(Wiener_Proj(X_bar,Z),proj_ind,0)
        W[proj_ind,:,:]=W1
        W_mag[proj_ind,:]=np.amax(abs(W1),axis=0)
    return W,W_mag



def compute_partial_Wiener_coeffs_avg(data,i,j,z=[],freq_choice=[]):
    # Returns W and W_mag
    # W: nNodes X n_FFT X nNodes
    # W_mag: returns h-infinity norm over frequencies
    c=np.append(j,np.int32(z))
    # c=np.unique(c)
    indexes = np.unique(c, return_index=True)[1]
    c=[c[index] for index in sorted(indexes)]
    n1=len(c)
    n_FFT=data.shape[2]
    W=np.zeros([n_FFT,n1],dtype='complex')
    # W_mag=np.zeros([n1])
    W_mag_avg = np.zeros([n1],dtype=float)
    for freq_index in freq_choice:
        fft_data=data[:,:,freq_index]
        # print(fft_data.shape)
        # print(freq_index)
        X=fft_data[:,i]
        X_bar=fft_data[:,c]
        # print(Z.shape,X_bar.shape)
        # print(X_bar[:3,:])
        W[freq_index,:]=Wiener_Proj(X_bar,X)
        W_mag_avg = W_mag_avg + np.abs(W[freq_index,:])
    W_mag_avg = np.divide(W_mag_avg,len(freq_choice))
    # W_mag=np.amax(abs(W),axis=0)
    return W,W_mag_avg




def Run_PC_Wiener_test_FFT_avg(data,nFFT,threshd):
    # print("Run_PC_Wiener_test_FFT_avg")
    nNodes=data.shape[1]
    C={}
    data_FFT = compute_fft(data,nfft=nFFT)
    freq_choice = range(nFFT)
    for ind1 in range(nNodes):
        for ind2 in range(nNodes):
            if ind1>ind2:
                # print('Chocsen node pair: [',ind1+1,',',ind2+1,']')
                flag=False
                # the complement set through which we condition
                D=[i for i in range(nNodes) if i != ind1 and i !=ind2]
                # print("D ",np.array(D)+1)
                # iterate over the combinations of various size in the increasing order
                for ind_cond in range(nNodes-1):
                    if flag==True:
                        break
                    combination_set=[i for i in combinations(D,ind_cond)]
                    # print("combination"+str(ind_cond)+" =",combination_set)
                    for c in combination_set:
                        # print("\n combination set: ",np.array(c)+1)
                        # # The partial Wiener coefficients does not seem symmetric wrt the projection. So, taking the average of two.
                        W_i_jZ1,W_mag_i_jZ1=compute_partial_Wiener_coeffs_avg(data_FFT,ind1,ind2,c,freq_choice)
                        
                        W_mag_i_jZ=W_mag_i_jZ1
                        
                        # print("Projection coefficient of X"+str(ind1+1)+" on "+str(np.append(ind2,c)+1)+"=",W_mag_i_jZ)
                        if W_mag_i_jZ[0]<threshd:
                            # print("Separation set of ( X"+str(ind1+1)+", X"+str(ind2+1),") is : ",np.array(c)+1)
                            # c=[l+1 for l in c]
                            C[ind1,ind2]=c
                            flag=True
                            break
                            c=np.array([ind1])
                        # if np.shape(c)[0]==3:#nNodes-2:
                        #     print("No separation set for ( X"+str(ind1+1)+", X"+str(ind2+1),")")
                        C[ind1,ind2]=None
            
    return C



def estimate_essential_graph(Top_est,C):
    nNodes = Top_est.shape[0]
    #% Collider Identification
    collider_graph = np.copy(Top_est)
    collider_set = []#np.ones((nNodes,1))*(-1)
    for i in range(nNodes):
        for j in range(nNodes):
            for k in range(nNodes):
                if i>j and j!=k and i!=k and Top_est[i,k]==1 and Top_est[j,k]==1 and Top_est[i,j]!=1 and (C[(i,j)]!=None) and (k not in C[(i,j)]):
                    collider_graph[i,k] = 0
                    collider_graph[j,k] = 0
                    collider_set.append(k)


    # % Orienting other edges if possible
    final_graph = np.copy(collider_graph)
    for i in range(nNodes):
        for j in range(nNodes):
            if i>j and collider_graph[i,j]==1 and collider_graph[j,i]==1:
                if i in collider_set:
                    final_graph[i,j] = 0
                if j in collider_set:
                    final_graph[j,i] = 0

    
    return final_graph,collider_graph



# Inputs:  (1) data: Time-series data
#          (2) nSamplles_list: The samples list to be used at which the errors are to be calculated
#          (3) nFFT: numer of FFT points
#          (4) thresh: Thresold for checking d-separation condition
#          (5) q: Maximum size of conditioning set
#          (6) true_ess_graph: Adjacency matrix of the true essential graph
# Outputs: (1) error_freq_hinf_PC: Array of network estimation errors at different sample sizes
#          (2) computation_freq: The Computation times at different sample sizes
def compute_PC_error_FFT_avg(data,nSamples_list,nFFT,true_ess_graph,thresh_freq,q):
    num_nodes=data.shape[1]
    error_freq_hinf_PC = np.zeros(len(nSamples_list))
    computation_freq = np.zeros(len(nSamples_list))
    for ind_samples in range(len(nSamples_list)):
        start_PC_freq = time.time()
        if nSamples_list[ind_samples] > 10*nFFT:
            C_FFT = Run_PC_Wiener_test_FFT_avg(data[:nSamples_list[ind_samples]],nFFT,threshd=thresh_freq[ind_samples])
        else:
            continue
        Top_est_half = np.zeros([num_nodes,num_nodes])
        for c in C_FFT:
            if C_FFT[c] == None: Top_est_half[c]=1
        Top_est_FFT = (Top_est_half + Top_est_half.T)
        est_essential_graph_fft,est_colliders = estimate_essential_graph(Top_est_FFT,C_FFT)
        error_freq_hinf_PC[ind_samples] = sum(sum(est_essential_graph_fft!=true_ess_graph))/sum(sum(true_ess_graph))
        computation_freq[ind_samples] = time.time()-start_PC_freq
    return error_freq_hinf_PC, computation_freq




# Function for tunning threshold for estimation in time domain.
def tune_time_threshold_deg(data,thresh_range,Lag,avg_deg_lim,max_deg_lim,min_deg_lim,nSamplesTune=32000,q=3):
    nNodes=data.shape[1]
    thresh_time = 0.01 # Default value of threshold
    thresh_vec = np.arange(thresh_range[0],thresh_range[1],thresh_range[2])
    for ind in range((len(thresh_vec))):
        C_time = Run_PC_Wiener_test_time(data[:nSamplesTune], Lag, thresh=thresh_vec[ind], Q = q)
        Top_est_half = np.zeros([nNodes,nNodes])
        for c in C_time:
            if C_time[c] == None: Top_est_half[c]=1
        Top_est_time = (Top_est_half + Top_est_half.T)
        min_deg, max_deg, avg_deg = compute_undir_graph_degrees(Top_est_time)
        if avg_deg >= avg_deg_lim[0] and  avg_deg <= avg_deg_lim[1] and max_deg <= max_deg_lim and min_deg >= min_deg_lim:
            thresh_time = thresh_vec[ind]
            break
    print(f"Chosen time threshold={thresh_time}")
    
    return thresh_time



# Function for tunning threshold for estimation in FFT domain.
def tune_freq_threshold_deg(data,nFFT,thresh_range,avg_deg_lim,max_deg_lim,min_deg_lim,nSamplesTune=32000,q=3):
    nNodes=data.shape[1]
    thresh_freq = 0.01 # Default value of threshold
    thresh_vec = np.arange(thresh_range[0],thresh_range[1],thresh_range[2])
    for ind in range((len(thresh_vec))):
        C_FFT = Run_PC_Wiener_test_FFT_avg(data[:nSamplesTune],nFFT,threshd=thresh_vec[ind])
        Top_est_half = np.zeros([nNodes,nNodes])
        for c in C_FFT:
            if C_FFT[c] == None: Top_est_half[c]=1
        Top_est_FFT = (Top_est_half + Top_est_half.T)
        min_deg, max_deg, avg_deg = compute_undir_graph_degrees(Top_est_FFT)
        if avg_deg >= avg_deg_lim[0] and  avg_deg <= avg_deg_lim[1] and max_deg <= max_deg_lim and min_deg >= min_deg_lim:
            thresh_freq = thresh_vec[ind]
            break
    print(f"Chosen fft={nFFT} threshold={thresh_freq}") 
    
    return thresh_freq


# Function for tunning threshold for Wiener phase.
def tune_wiener_phase_threshold_deg(data,nFFT,thresh_range,avg_deg_lim,max_deg_lim,min_deg_lim,nSamplesTune=32000,q=3):
    thresh_vec = np.arange(thresh_range[0],thresh_range[1],thresh_range[2])
    thresh_freq = 0.01 # Default value of theshold
    for ind in range((len(thresh_vec))):
        est_essential_graph_FFT = Run_Wiener_Phase_avg(data[:nSamplesTune],nFFT,thresh=thresh_vec[ind])
        MG_wph,Top_est_win_ph = find_moral_graph_and_skeleton_from_adjacency_matrix(est_essential_graph_FFT)
        min_deg, max_deg, avg_deg = compute_undir_graph_degrees(Top_est_win_ph)
        if avg_deg >= avg_deg_lim[0] and  avg_deg <= avg_deg_lim[1] and max_deg <= max_deg_lim and min_deg >= min_deg_lim:
            thresh_freq = thresh_vec[ind]
            break
    print(f"Chosen win-ph threshold={thresh_freq}")
    
    return thresh_freq



def Run_Wiener_Phase_avg(data,nFFT,thresh=0.1):
    nNodes=data.shape[1]
    data_FFT=compute_fft(data,nfft=nFFT)
    W,Wmag=compute_Wiener_coeffs(nNodes, data_FFT)
    
    
    
    # Moral Graph estimation using absolute value of average magnitude of Wiener projection over frequencies
    W_mag_avg = np.zeros((nNodes,nNodes),dtype=float)

    freq_choice = range(nFFT)
    for i in freq_choice:
        W_mag_avg = W_mag_avg + np.abs(W[:,i,:])
        
    W_mag_avg = np.divide(W_mag_avg,len(freq_choice))
    W_mag_avg = np.divide(W_mag_avg+np.transpose(W_mag_avg),2)
    moral_graph_adj=W_mag_avg>thresh
    
    # Skeleton estimation using h-inf of imaginary part of Wiener projection coefficients
    W_imag=W.imag
    W_imag_avg = np.zeros((nNodes,nNodes),dtype=float)
    freq_choice = range(nFFT)
    for i in freq_choice:
        W_imag_avg = W_imag_avg + np.abs(W_imag[:,i,:])
        
    W_imag_avg = np.divide(W_imag_avg,len(freq_choice))
    W_imag_avg = np.divide(W_imag_avg+np.transpose(W_imag_avg),2)
    skeleton_adj=W_imag_avg>thresh
    
    # Essential graph estimation
    spurious_adj=abs(moral_graph_adj^skeleton_adj)
    C={}
    collider_graph=np.zeros_like(skeleton_adj)
    collider_graph[:]=skeleton_adj[:]
    collider_set = []
    for i in range(nNodes):
        for j in range(nNodes):
            if i>j and spurious_adj[i,j]==1:
                potential_colider=[]
                for k in range(nNodes):
                    if skeleton_adj[i,k]==1 and skeleton_adj[j,k]==1:
                        potential_colider.append(k)
                C[i,j]=potential_colider
                if len(C[i,j])==1:
                    collider_graph[i,C[i,j]]=False
                    collider_graph[j,C[i,j]]=False
                    collider_set.append(C[i,j][0])
                else:
                    for l in C[i,j]:
                        W_ij_c,W_ij_c_mag=compute_partial_Wiener_coeffs_avg(data_FFT, i, j,l,freq_choice = range(nFFT))
                        if W_ij_c_mag[0]>thresh:
                            collider_graph[i,l]=False
                            collider_graph[j,l]=False
                            collider_set.append(l)
                            
                            
    # % Orienting other edges if possible
    final_graph = np.copy(collider_graph)
    for i in range(nNodes):
        for j in range(nNodes):
            if i>j and collider_graph[i,j]==1 and collider_graph[j,i]==1:
                if i in collider_set:
                    final_graph[i,j] = 0
                if j in collider_set:
                    final_graph[j,i] = 0
                        
    
    return final_graph




def compute_Wiener_phase_error_avg(data,nSamples_list,nFFT,trueEssGraph,threshd):
    error_freq_wiener_phase=np.zeros(len(nSamples_list))
    computation_freq=np.zeros(len(nSamples_list))
    for ind_samples in range(len(nSamples_list)):
        start_wiener_phase_freq=time.time()
        # print(nSamples_list[ind_samples])
        est_essential_graph_fft=Run_Wiener_Phase_avg(data[:nSamples_list[ind_samples]],nFFT,thresh=threshd[ind_samples])
        
        error_freq_wiener_phase[ind_samples]=sum(sum(est_essential_graph_fft!=trueEssGraph))/sum(sum(trueEssGraph))
        computation_freq[ind_samples]=time.time()-start_wiener_phase_freq
    return error_freq_wiener_phase, computation_freq



def compute_time_PC_error(data,nSamples_list,lag_time_WF,trueEssGraph,threshd):
    nNodes=data.shape[1]
    error_time_PC = np.zeros(len(nSamples_list))
    computation_PC_time = np.zeros(len(nSamples_list))
    for ind_samples in range(len(nSamples_list)):
        start_PC_time = time.time()
        C_time = Run_PC_Wiener_test_time(data[:nSamples_list[ind_samples]], lag_time_WF, thresh=threshd[ind_samples], Q = nNodes-2)
        Top_est_half = np.zeros([nNodes,nNodes])
        for c in C_time:
            if C_time[c] == None: Top_est_half[c]=1
        Top_est_time = (Top_est_half + Top_est_half.T)
        est_essential_graph_time,est_colliders_time = estimate_essential_graph(Top_est_time,C_time)
        error_time_PC[ind_samples] = sum(sum(est_essential_graph_time!=trueEssGraph))/sum(sum(trueEssGraph))
        computation_PC_time[ind_samples] = time.time()-start_PC_time
        
    return error_time_PC, computation_PC_time




def find_true_Essential_graph(Skeleton, Adj):
    collider_graph = np.copy(Skeleton)
    collider_nodes = []
    for ind1 in range(num_nodes): # ind1 is the collider node index
        for ind2 in range(num_nodes):    
        #ind2 and ind3 are the indices for the parent of colliders
            for ind3 in range(num_nodes):
                if (Adj[ind1,ind2] == 1) and ((Adj[ind1,ind3] == 1)) and Skeleton[ind2,ind3] !=1 and ind2 != ind3 and ind1 != ind3 and ind1 != ind2:
                    print( ind1," is a collider")
                    collider_graph[ind2,ind1]=0
                    collider_graph[ind3,ind1]=0
                    collider_nodes.append(ind1)                       
    # % Orienting other edges if possible
    true_ess_graph = np.copy(collider_graph)
    for i in range(num_nodes):
        for j in range(num_nodes):
            if i>j and collider_graph[i,j]==1 and collider_graph[j,i]==1:
                if i in collider_nodes:
                    true_ess_graph[i,j] = 0
                if j in collider_nodes:
                    true_ess_graph[j,i] = 0
                    
    return true_ess_graph



def compute_undir_graph_degrees(Adj):
    nNodes = Adj.shape[0]
    node_degs = []
    for node in range(nNodes):
        node_degs.append(sum(Adj[:,node]))
    max_deg = max(node_degs)
    min_deg = min(node_degs)
    avg_deg = sum(node_degs)/nNodes
    
    return min_deg, max_deg, avg_deg


   


# Main function ##################################################################################################################
##################################################################################################################################
if __name__ == '__main__':
    
    ## Initialize the variables
    start_sim_time=time.time()
    print("Start time: ",start_sim_time)
    nNetworks = 25
    nFFT = 32
    nSamples_list = [640, 3200, 6400, 32000, 128*10**3, 32*10**4]
    num_nodes = 6
    max_in_degree = 2
    max_out_degree = 1
    noise_pow1 = 1*np.ones(num_nodes)
    nLags = 1
    avg_deg_lim = [1.6,2.7]
    max_deg_lim = 5
    min_deg_lim = 1
    random.seed(1)
    np.random.seed(1)

    error_freq_hinf_Networks  = np.zeros([nNetworks, len(nSamples_list)])
    error_freq_wiener_phase_networks = np.zeros([nNetworks, len(nSamples_list)])
    error_time_Networks = np.zeros([nNetworks, len(nSamples_list)])
    
    computation_freq_Networks = np.zeros([nNetworks, len(nSamples_list)])
    computation_wiener_phase_networks = np.zeros([nNetworks, len(nSamples_list)])
    computation_time_Networks = np.zeros([nNetworks, len(nSamples_list)])

    # Iterating through the randomly generated networks
    for ind_network in range(nNetworks):
        print("ind_network: ",ind_network,"-------------------------------------------------------")
        ########### ########### ########### ########### ########### ########### ########### 
        # Network and Data Generation
        ########### ########### ########### ########### ########### ########### ###########
        #  Random network generation
        G,Adj = generate_graph(num_nodes,max_in_degree)
        MG,Skeleton = find_moral_graph_and_skeleton_from_adjacency_matrix(Adj)
        A1,B1 = create_A_B_coeffs_6_nodes(Adj)
        true_ess_graph = find_true_Essential_graph(Skeleton, Adj)
        
        ########### ########### Generate data ############################################  
        start_data=time.time()
        print("Data generation started at ", start_data)
        data=continuous_data_generation_arbitrary_lags(nSamples_list[-1],B=B1,A=A1, noise_pow=noise_pow1,nlags=nLags)
        print("Time taken to generate data = ", time.time()-start_data)


        ############ ########### Tuning the thresholds ########### ###########
        thresh_freq = np.zeros([nNetworks, len(nSamples_list)])
        thresh_win_phase = np.zeros([nNetworks, len(nSamples_list)])
        thresh_time = np.zeros([nNetworks, len(nSamples_list)])
        lag_time_WF = 15
        thresh_range_FFT = [0.001,0.2,0.005]
        thresh_range_time = [0.001,0.05,0.001]
        thresh_range_win_ph = [0.001,0.2,0.01]
        for ind_samp in range(len(nSamples_list)):
            print("Tuning at sample size = ",nSamples_list[ind_samp],"**************************")
            thresh_freq[ind_network,ind_samp] = tune_freq_threshold_deg(data, nFFT, thresh_range_FFT,avg_deg_lim,max_deg_lim,min_deg_lim, nSamplesTune = nSamples_list[ind_samp], q = num_nodes-2)
            thresh_time[ind_network,ind_samp] = tune_time_threshold_deg(data, thresh_range_time, lag_time_WF,avg_deg_lim,max_deg_lim,min_deg_lim, nSamplesTune = nSamples_list[ind_samp], q= num_nodes-2)
            thresh_win_phase[ind_network,ind_samp] = tune_wiener_phase_threshold_deg(data, nFFT, thresh_range_win_ph,avg_deg_lim,max_deg_lim,min_deg_lim, nSamplesTune = nSamples_list[ind_samp])


        ########### ########### ########### ########### ########### ########### 
        print("PC algorithms")
        #################################### #################################### 
        ## PC algorithm in time domain
        #################################### #################################### 
        error_time_PC,computation_PC_time = compute_time_PC_error(data,nSamples_list,lag_time_WF,true_ess_graph,threshd=thresh_time[ind_network,:])
        error_time_Networks[ind_network,:]=error_time_PC
        computation_time_Networks[ind_network,:]=computation_PC_time
        

        #################################### #################################### 
        ## PC algorithm in Freq domain
        #################################### #################################### 
        error_freq_hinf_PC ,computation_freq =compute_PC_error_FFT_avg(data,nSamples_list,nFFT, true_ess_graph, thresh_freq=thresh_freq[ind_network,:],q=num_nodes-2)
        error_freq_hinf_Networks[ind_network,:] = error_freq_hinf_PC
        computation_freq_Networks[ind_network,:] = computation_freq
        
        # #################################### #################################### 
        # ## Wiener-Phase algorithm
        # #################################### ####################################
        error_freq_hinf_wiener_phase,computation_wiener_phase = compute_Wiener_phase_error_avg(data,nSamples_list,nFFT,true_ess_graph,threshd=thresh_win_phase[ind_network,:])
        error_freq_wiener_phase_networks[ind_network,:] = error_freq_hinf_wiener_phase
        computation_wiener_phase_networks[ind_network,:] = computation_wiener_phase
        
        

# Plotting the results
plt.style.use('seaborn-v0_8-colorblind')
FONTSIZE=20
WIDTH_SIZE=8
HEIGHT_SIZE=6
LEGENDSIZE=14
LINEWIDTH=3

plt.figure(figsize=(WIDTH_SIZE,HEIGHT_SIZE))
plt.semilogx(nSamples_list,sum(error_freq_hinf_Networks)/nNetworks,'b>-',label="nFFT="+str(nFFT ),linewidth=LINEWIDTH,markersize=11)
plt.semilogx(nSamples_list,sum(error_freq_wiener_phase_networks)/nNetworks,'cx--',label="Win-Ph",linewidth=LINEWIDTH,markersize=11)
plt.semilogx(nSamples_list,sum(error_time_Networks)/nNetworks,'gD-.',label="time",linewidth=LINEWIDTH,markersize=11)

plt.legend(loc="upper right", fontsize=LEGENDSIZE)
plt.title("Time vs Freq Average Error Rate Comparison, nNodes="+str(num_nodes), fontsize=15)
plt.xlabel(r"Number of samples",fontsize=FONTSIZE)
plt.ylabel(r"Error rate", fontsize=FONTSIZE)
# plt.savefig("error_comp_PC_Wiener_Phase_IID_noise,"+"nNodes="+str(num_nodes)+".pdf")
plt.show()

plt.figure(figsize=(WIDTH_SIZE,HEIGHT_SIZE))
plt.semilogx(nSamples_list,sum(computation_freq_Networks)/nNetworks,'b>-',label="nFFT="+str(nFFT),linewidth=LINEWIDTH,markersize=11)
plt.semilogx(nSamples_list,sum(computation_wiener_phase_networks)/nNetworks,'cx--',label="Win-Ph",linewidth=LINEWIDTH,markersize=11)
plt.semilogx(nSamples_list,sum(computation_time_Networks)/nNetworks,'gD-.',label="time")
plt.legend(loc="upper left", fontsize=LEGENDSIZE)
plt.title("Time vs Freq  Average Computation Time, nNodes="+str(num_nodes), fontsize=15)
plt.xlabel(r"Number of samples",fontsize=FONTSIZE)
plt.ylabel(r"Computation time", fontsize=FONTSIZE)
# plt.savefig("computation_time_PC_Wiener_Phase_random_IID_noise,"+"nNodes="+str(num_nodes)+".pdf")
plt.show()
# 
print("stop")
    
