import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# Function to create a custom alpha matrix based on a causal graph
# def create_custom_alpha(l, non_zero_indices, alpha_min=0.4, alpha_max=0.8):
def create_custom_alpha(l, non_zero_indices, alpha_min=0.4, alpha_max=0.8):
    """
    l: Number of processes.
    non_zero_indices: List of tuples [(i, j), ...] where alpha[i, j] is non-zero, that is, j->i exists.
    alpha_min, alpha_max: Range for random non-zero values.
    """
    alpha = np.zeros((l, l))  # Initialize all elements to zero
    for i, j in non_zero_indices:
        alpha[i, j] = np.round(np.random.uniform(alpha_min, alpha_max),3)  # Assign random non-zero values
        # alpha[i,j] = ToBij()
        # alpha[0,7] = 0.3
        print("The coefficient alpha["+str(i)+","+str(j)+"] is",str(alpha[i, j]))
    return alpha

def compute_theta_ij_k(alpha, beta, delta, l, h):
    """"
    Function to generate excitation coefficients with exponential decay excitation function
    aplha: Decay rate of the exponential excitation function. lxl ndarray
    beta: Decay rate of the exponential excitation function. constant
    delta: Time step size. constant
    l: Number of processes
    h: Number of non-zero effective lag terms. manually set
    """
    theta = np.zeros((l, l, h))  # Shape (l, l, m), where m is the number of considered lag terms
    for i in range(l):
        for j in range(l):
            for k in range(1, h+1):
                theta[i, j, k-1] = (
                    alpha[i, j] * np.exp(-k * beta * delta) * (np.exp(beta * delta) - 1) / beta
                )
    return theta


# Function to calculate Phi and check spectral radius
def compute_Phi_and_check(alpha, beta):
    Phi = alpha / beta  # Compute Phi matrix
    eigvals = np.linalg.eigvals(Phi)  # Eigenvalues of Phi
    spectral_radius = max(np.abs(eigvals))  # Spectral radius
    return Phi, spectral_radius


# Function to simulate the multivariate Hawkes process
def simulate_hawkes_process(mu, alpha, beta, delta, l, h, n, noise_std=None):
    """
    h: Number of non-zero effective lag terms
    n: Number of total time steps
    """
    # Compute the excitation coefficients
    theta = compute_theta_ij_k(alpha, beta, delta, l, h)
    
    # Initialize the process data
    N = np.zeros((l, n))  # Shape (l, n), where l is the number of processes

    # If noise_std is not provided, default to a standard deviation of 0.5 for all processes
    if noise_std is None:
        noise_std = [0.5] * l

    # Simulate the process iteratively
    for t in range(1, n):
        for i in range(l):
            # Compute the weighted sum of past events
            summation = 0
            for j in range(l):
                if t <= h:
                    for k in range(1, t+1):
                        summation += theta[i, j, k-1] * N[j, t-k]
                else:
                    for k in range(1, h+1):
                        summation += theta[i, j, k-1] * N[j, t-k]
            
            # Add white noise and the background intensity
            noise = np.random.normal(0, noise_std[i])
            # if i == 7:
            #     noise = pow(np.random.exponential(scale=noise_std[i], size=1),2)
            # noise = np.random.normal(0, 0.5)
            N[i, t] = summation + mu[i] * delta + noise
    
    return N


def generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices):   
    # Create the custom alpha matrix
    alpha = create_custom_alpha(l, non_zero_indices)
    # Calculate Phi and check the spectral radius
    Phi, spectral_radius = compute_Phi_and_check(alpha, beta)
    # Ensure stability condition is met
    while spectral_radius >= 1:
        # Regenerate alpha while keeping the same causal graph
        alpha = create_custom_alpha(l, non_zero_indices)
        Phi, spectral_radius = compute_Phi_and_check(alpha, beta)
    # noise_std = np.random.uniform(0.5, 5, l)
    noise_std = []
    for i in range(l):
        # noise_std.append(np.random.uniform(6, 9)+np.random.uniform(0, 0.5) )
        noise_std.append(np.round((l-1-i)*10+np.random.uniform(1, 9),3))
    print("The noise_std generated are",str(noise_std),".")
    
    # Simulate the Hawkes process
    N = simulate_hawkes_process(mu, alpha, beta, delta, l, h, n, noise_std)
    # print(N.shape)
    observed_data = N[observed_indices,:]
    data = pd.DataFrame(observed_data.T,columns=[f'O_{i+1}' for i in range(len(observed_indices))])
    data = (data-data.mean())/data.std()
    # data = data-data.mean()
        
    return data  


def Case1(Num=5000):
    # Set parameters
    l = 3  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2: non_zero_indices = [(0, 1), (2, 1)]
    non_zero_indices = [(0, 1), (2, 1), (1,1)]
    # list of indices of observed component processes
    observed_indices = [0,1,2] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data  


def Case1_0(Num=5000):
    # Set parameters
    l = 3  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,0), (0, 1), (1,1), (1,2), (2, 1), (2,2)]
    # list of indices of observed component processes
    observed_indices = [0,1,2] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data  


def Case2(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (3,2), (3,3), (4, 2), (4,4)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data  



def Case3(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3), (2, 4), (4,4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 


def Case4(Num=5000):
    # Set parameters
    l = 7  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.5  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (3,2), (5,3), (5,5), (4, 2), (6,4), (6,6)]
    # list of indices of observed component processes
    observed_indices = [0,1,5,6] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 



def Case6(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3), (4,3), (4, 4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 



def Case7(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (4,2), (3,3), (4,3), (4, 4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 



def Case8(Num=5000):
    # Set parameters
    l = 8  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,0), (0,2), (2,4), (4,4), (1,3), (3,4), (6,5), (6,6), (7,7), (7,5), (5,4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,6,7] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 


def Case9(Num=5000):
    # Set parameters
    l = 8  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    # non_zero_indices = [(0,0), (0,2), (1,1), (1,2), (2,3), (3,3), (3,4), (4,4), (5,3), (6,5), (6,6), (7,5), (7,7)]
    non_zero_indices = [(0,0), (0,2), (1,1), (1,2), (2,3), (3,3), (5,3), (6,5), (6,6), (7,5), (7,7)]
    # list of indices of observed component processes
    observed_indices = [0,1,6,7] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data



def Case10(Num=5000):
    # Set parameters
    l = 9  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1 #0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    # non_zero_indices = [(0,0), (0,2), (1,1), (1,2), (2,3), (3,3), (3,4), (4,4), (5,3), (6,5), (6,6), (7,5), (7,7)]
    non_zero_indices = [(0,0), (0,7), (4,7), (1,4), (2,4), (1,1), (2,2), (8,7), (8,8)]
    # list of indices of observed component processes
    observed_indices = [0,1,2,8] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data



def CaseFig1b(Num=5000):
    # Set parameters
    l = 3  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,0), (0, 1), (1,1), (1,2), (2, 1), (2,2)]
    # list of indices of observed component processes
    observed_indices = [0,1,2] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data  




def CaseFig2a(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (1, 2), (2,2), (3,2), (3,3), (4, 2), (4,4)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data  



def CaseFig4a(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3), (4,3), (4, 4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 



def CaseFig4b(Num=5000):
    # Set parameters
    l = 5  # Number of processes
    h = 5 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (4,2), (3,3), (4,3), (4, 4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,3,4] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 



def CaseFig4c(Num=5000):
    # Set parameters
    l = 8  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    non_zero_indices = [(0,0), (0,2), (2,4), (4,4), (1,1), (1,3), (3,4), (4,4), (6,5), (6,6), (7,7), (7,5), (5,4)]
    # non_zero_indices = [(0,2), (0,0), (1, 2), (1,1), (2,2), (2,3), (3,3)]
    # list of indices of observed component processes
    observed_indices = [0,1,6,7] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data 


def CaseFig4d(Num=5000):
    # Set parameters
    l = 8  # Number of processes
    h = 50 # Number of non-zero effective lag terms
    n = Num  # Number of total time steps
    delta = 0.1  # Time step size
    beta = 1.0  # Decay rate of the exponential excitation function
    mu = np.random.uniform(15, 25, l)  # Background intensities
    # Define a custom causal graph (indices of non-zero elements in the alpha matrix)
    # Example: Process 1 influences 0, and Process 1 influences 2
    # non_zero_indices = [(0,0), (0,2), (1,1), (1,2), (2,3), (3,3), (3,4), (4,4), (5,3), (6,5), (6,6), (7,5), (7,7)]
    non_zero_indices = [(0,0), (0,2), (1,1), (1,2), (2,3), (3,3), (3,4), (4,4), (5,3), (6,5), (6,6), (7,5), (7,7)]
    # list of indices of observed component processes
    observed_indices = [0,1,6,7] # Example: [0,2] the first and third component processes are observed.
    data = generate_simulation(l, h, n, delta, beta, mu, non_zero_indices, observed_indices)
    return data






def main():
    data = Case1()

if __name__ == '__main__':
    main()