import numpy as np
import pandas as pd



# read data
df = pd.read_csv('../input/tabular-playground-series-jul-2022/data.csv')
x  = np.array(df.loc[:,'f_22':'f_28']) # these are the only relevant float vars



# power transform integer vars to make them look more 'normal'
power = 0.47 # best value
x2 = ((1 + np.array(df.loc[:,'f_07':'f_13']))**power - 1) / power



# init EM means, covariances, support vars
nc = 7 * 6 # 7 groups of 6 clusters each
ci = np.zeros([nc, 7, 7], dtype=np.float32) # (nc, 7, 7) covariance of integer vars - init to diagonal matrix
for i in range(7): ci[:,i,i] = 1
c  = np.zeros([nc, 5, 5], dtype=np.float32) # (nc, 5, 5) covariance of float vars - init to diagonal matrix. Here 2 uncorrelated vars are removed
for i in range(5): c[:,i,i] = 1
# hardcode float means to points selected by GMM(42) for float vars only - to improve convergence with full data
m  = np.array([1,-1,-1,-1,-1,1,-1,1,1,-1,-1,-1,1,1,1,1,1,-1,1,-1,-1,1,1,-1,-1,-1,1,1,1,1,1,-1,1,-1,-1,-1,1,-1,1,1,1,1,-1,1,-1,-1,-1,1,-1,-1,
1,-1,-1,-1,-1,-1,1,1,-1,1,-1,1,1,-1,1,-1,1,1,1,-1,-1,1,1,-1,-1,-1,-1,1,-1,-1,-1,-1,1,-1,1,-1,1,-1,1,-1,
-1,-1,-1,1,1,1,1,-1,1,-1,1,-1,1,1,-1,1,1,1,1,1,-1,1,1,-1,-1,-1,1,-1,-1,-1,-1,-1,-1,-1,1,-1,-1,-1,-1,-1,
1,-1,-1,-1,1,-1,-1,1,-1,-1,-1,1,-1,1,-1,-1,1,-1,-1,1,-1,1,-1,-1,-1,1,1,1,1,1,-1,-1,1,1,-1,1,1,1,-1,-1,
-1,-1,-1,-1,-1,1,-1,-1,-1,1,1,1,1,-1,1,1,-1,1,-1,1,-1,1,-1,-1,1,1,-1,1,1,1,-1,-1,-1,1,1,1,1,1,-1,-1], dtype=np.float32).reshape(nc,5)
mi = 3.33 * np.ones([nc, 7], dtype=np.float32) # (nc, 7) mean of integer vars. Init to mean of them all
non_corr_vars = np.array([0,1, 3,4, 0,5, 1,5, 2,5, 4,5, 4,6]).reshape(7,2) # for 7 cluster groups, indices of 2 vars that are uncorrelated with the rest. They have 0 mean and 1 std.
prob = np.zeros([x.shape[0], nc], dtype=np.float32) # predicted probability of each point for each cluster



# custom EM loop
for ii in range(20):# 20 iters are enough to converge
    # 1. Expectation
    for i in range(nc): # 42 clusters
        # ints
        cov = ci[i,:,:]
        mu  = mi[i,:]
        p1  = ((x2 - mu.reshape(1,-1)).dot(np.linalg.inv(cov)) * (x2 - mu.reshape(1,-1))).sum(axis=1)
        p1  = np.exp(-p1/2) / np.linalg.det(cov)**(1/2)
        # floats part 1 - 5 correlated vars
        cov = c[i,:,:]
        mu  = m[i,:]
        ncv = list(non_corr_vars[i//6,:]) # indices of 2 vars that are uncorrelated with the rest
        cv  = list(set([0,1,2,3,4,5,6]) - set(ncv)) # indices of 5 vars that are correlated with each other
        x1  = x[:,cv]
        p2  = ((x1 - mu.reshape(1,-1)).dot(np.linalg.inv(cov)) * (x1 - mu.reshape(1,-1))).sum(axis=1)
        p2  = np.exp(-p2/2) / np.linalg.det(cov)**(1/2)
        # floats part 2 - 2 uncorrelated vars with 0 mean and 1 std
        x3  = x[:,ncv]
        p3  = (x3 * x3).sum(axis=1)
        p3  = np.exp(-p3/2)
        prob[:,i] = p1 * p2 * p3 # final probability is the product of 3 pieces
    gm2 = prob.argmax(axis=1)
    LP  = np.log(prob.max(axis=1)).mean() # log prob - to track convergence
    print(ii, np.round(LP, 4))

    # 2. Maximization
    for i in range(7):# ints - 7 cluster groups *******************************************
        x_cl = x2[(gm2//6)==i,:] # points for this cluster block only
        # set means to means of x
        mi[i*6:(i+1)*6,:] = x_cl.mean(axis=0) # int mean is the same for all subclusters
        # covariance
        x_cl = x_cl - mi[i*6,:].reshape(1,-1)
        ci[i*6:(i+1)*6,:,:] = x_cl.T.dot(x_cl) / x_cl.shape[0] # int cov is the same for all subclusters
        # set cov to 0 if it is close to zero (<.01)
        c2 = ci[i*6:(i+1)*6,:,:].ravel()
        c2[(np.abs(c2) < .01)] = .0
        ci[i*6:(i+1)*6,:,:] = c2.reshape(-1, 7, 7)
    for i in range(nc):# floats - 42 clusters ******************************************
        ncv = list(non_corr_vars[i//6,:]) # indices of 2 vars that are uncorrelated with the rest
        cv  = list(set([0,1,2,3,4,5,6]) - set(ncv)) # indices of 5 vars that are correlated with each other
        x_cl = x[gm2==i,:] # points for this cluster only
        # covariance (mean is hardcoded so does not need to be updated)
        x_cl = x_cl[:,cv] - m[i,:].reshape(1,-1)
        c[i,:,:] = x_cl.T.dot(x_cl) / x_cl.shape[0]



# submission
sub = pd.read_csv('../input/tabular-playground-series-jul-2022/sample_submission.csv')
sub['Predicted'] = gm2//6
sub.to_csv('submission.csv', index=False)