import numpy as np
import pandas as pd

def generate_data(n, ncov, rho, K, Beta0, Beta, weights, umax):
    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

    rng = np.random.default_rng(123)  # Create a random number generator

    X = rng.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
    censored_time = np.log(rng.uniform(0, umax, size=n))  # Censored time

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

    # Generate noise terms
    noise = rng.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]

        # Apply censoring if applicable
        if Y[i] > censored_time[i]:
            Y[i] = censored_time[i]

    # Generate censoring indicator based on comparison with censored time
    censoring = (Y < censored_time).astype(int)

    return X, np.exp(Y), censoring, np.exp(censored_time)

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
umax = 40  # Maximum censoring time

# 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, censoring, censored_time = generate_data(n, ncov, rho, K, Beta0, Beta, weights, umax)

# Save data to CSV files
save_to_csv(X, 'Xdata_1.csv')
save_to_csv(Y, 'Ydata_1.csv')
save_to_csv(censoring, 'Ddata_1.csv')
save_to_csv(censored_time, 'Cdata_1.csv')
