import numpy as np
import pandas as pd

def generate_data(n, ncov, rho, K, Beta0, Beta, weights):
    mean = np.zeros(ncov)  # Mean vector of zeros
    cov = np.zeros((ncov, ncov))  # Variance-covariance matrix

    for i in range(ncov):
        for j in range(ncov):
            cov[i, j] = rho ** abs(i - j)  # Calculate covariance elements

    X = np.random.multivariate_normal(mean, cov, size=n)  # Generate data

    # Adjust mean column to be exactly zero
    X -= np.mean(X, axis=0)

    # Generate response variable Y using finite mixture of regression
    Y = np.zeros(n)  # Response variable

    # Select random components based on weights
    components = np.random.choice(K, size=n, p=weights)

    # Generate noise terms
    noise = np.random.normal(0, 1, size=n)

    for i in range(n):
        # Generate response variable Y for the selected component
        Y[i] = Beta0[components[i]] + np.dot(X[i], Beta[components[i]]) + noise[i]

    return X, Y

def save_to_csv(data, filename):
    df = pd.DataFrame(data)
    df.to_csv(filename, index=False)

# Set seed for reproducibility
np.random.seed(1980)

# Parameters
n = 500  # Number of data points
ncov = 10  # Dimension of the design matrix
rho = 0.25  # Covariance parameter
K = 2  # Number of regression components

# Define regression coefficients for each component
Beta0 = [2.0, -1.0]  # Intercepts
Beta = [
    np.array([2,-1,-2,1,2,0,0,0,0,0]),
    np.array([-1,1,2,0,0,0,0,-1,2,-2])
]  # Regression coefficients

# Define weights for each component (summing up to 1)
weights = [0.4, 0.6]

# Generate data
X, Y = generate_data(n, ncov, rho, K, Beta0, Beta, weights)

# Save data to CSV files
save_to_csv(X, 'Xdata_2.csv')
save_to_csv(Y, 'Ydata_2.csv')
