rm(list = ls()); gc();

library(data.table)
library(lightgbm)
set.seed(52)

train <- fread("../input/train.csv")
test <- fread("../input/test.csv")
structure <- fread("../input/structures.csv")

names(structure) <- c("molecule_name","atom_index_0", "atom_0", "x_0", "y_0", "z_0")
train[, names(structure)[3:6] := structure[train, on=.(molecule_name, atom_index_0), 
                                    mget(paste0("x.", names(structure)[3:6]))]]
test[, names(structure)[3:6] := structure[test, on=.(molecule_name, atom_index_0), 
                                    mget(paste0("x.", names(structure)[3:6]))]]                                    

names(structure) <- c("molecule_name","atom_index_1", "atom_1", "x_1", "y_1", "z_1")                                    
train[, names(structure)[3:6] := structure[train, on=.(molecule_name, atom_index_1), 
                                    mget(paste0("x.", names(structure)[3:6]))]]
test[, names(structure)[3:6] := structure[test, on=.(molecule_name, atom_index_1), 
                                    mget(paste0("x.", names(structure)[3:6]))]]
test[, scalar_coupling_constant:=0]

rm(structure); gc();

train_test <- rbind(train,test)

rm(train, test); gc();

train_test[, ":="(diff_structure_x = x_0 - x_1, diff_structure_y = y_0 - y_1, diff_structure_z = z_0 - z_1,
                  dist_xyz = sqrt((x_0 - x_1) ^ 2 + (y_0 - y_1) ^ 2 + (z_0 - z_1) ^ 2))]

train_test[, (c("count_molecule", "min_molecule_dist_xyz", "mean_molecule_dist_xyz", 
                "median_molecule_dist_xyz", "max_molecule_dist_xyz", "sd_molecule_dist_xyz")) 
                := list(.N, min(dist_xyz), mean(dist_xyz), median(dist_xyz), max(dist_xyz), sd(dist_xyz)), 
                by = molecule_name]

train_test[, (c("min_molecule_atom_0_dist_xyz", "mean_molecule_atom_0_dist_xyz", "median_molecule_atom_0_dist_xyz", 
                "max_molecule_atom_0_dist_xyz", "sd_molecule_atom_0_dist_xyz")) 
                := list(min(dist_xyz), mean(dist_xyz), median(dist_xyz), max(dist_xyz), sd(dist_xyz)), 
                by = .(molecule_name, atom_index_0)]
                
train_test[, (c("min_molecule_atom_1_dist_xyz", "mean_molecule_atom_1_dist_xyz", "median_molecule_atom_1_dist_xyz", 
                "max_molecule_atom_1_dist_xyz", "sd_molecule_atom_1_dist_xyz")) 
                := list(min(dist_xyz), mean(dist_xyz), median(dist_xyz), max(dist_xyz), sd(dist_xyz)), 
                by = .(molecule_name, atom_index_1)]                

train_test[, ":="(diff_mean_molecule_dist_xyz = dist_xyz - mean_molecule_dist_xyz,
                  diff_median_molecule_dist_xyz = dist_xyz - median_molecule_dist_xyz,
                  diff_mean_molecule_atom_0_dist_xyz = dist_xyz - mean_molecule_atom_0_dist_xyz,
                  diff_median_molecule_atom_0_dist_xyz = dist_xyz - median_molecule_atom_0_dist_xyz,
                  diff_mean_molecule_atom_1_dist_xyz = dist_xyz - mean_molecule_atom_1_dist_xyz,
                  diff_median_molecule_atom_1_dist_xyz = dist_xyz - median_molecule_atom_1_dist_xyz)]                                              

train <- train_test[scalar_coupling_constant!=0]
test <- train_test[scalar_coupling_constant==0]
train_scc<- train$scalar_coupling_constant

rm(train_test); gc();

train[, (c("id","molecule_name","atom_0","atom_1","scalar_coupling_constant")):=list(NULL,NULL,NULL,NULL,NULL)]
test[, (c("id","molecule_name","atom_0","atom_1","scalar_coupling_constant")):=list(NULL,NULL,NULL,NULL,NULL)]

params <- list(learning_rate = 0.2,
               boosting_type = 'gbdt',
               objective = "regression",
               metric = 'mae',
               nthread = 8,
               sub_feature = 1,
               bagging_fraction = 0.85,
               bagging_freq = 1,
               num_leaves = 128, 
               min_split_gain = 0.01,
               min_data = 50,
               reg_alpha = 0.13,
               reg_lambda = 0.36,
               min_hessian = 0.001)

submit <- fread("../input/sample_submission.csv")

for (mol_type in unique(train$type)) {
    train_type_idx <- train$type == mol_type
    test_type_idx <- test$type == mol_type
    nrounds <- round(sum(1*(train_type_idx==TRUE))/50)

    lgb_train <- lgb.Dataset(data = data.matrix(train[train_type_idx, setdiff(names(train),"type"), with = FALSE]), label = train_scc[train_type_idx])  
    lgbm_mod <- lgb.train(params = params, lgb_train, nrounds = nrounds)  
    submit$scalar_coupling_constant[test_type_idx] <- predict(lgbm_mod, data.matrix(test[test_type_idx, setdiff(names(test),"type"), with = FALSE]))
}

fwrite(submit, "submission.csv")