library(data.table)
library(dplyr)

train <- fread('../input/train.csv', 
               select = c('Cliente_ID', 'Producto_ID', 'Agencia_ID', 'Ruta_SAK', 'Demanda_uni_equil'))
test <- fread('../input/test.csv', 
              select = c('id', 'Cliente_ID', 'Producto_ID', 'Agencia_ID', 'Ruta_SAK'))
train$log_demand = log1p(train$Demanda_uni_equil) 

setkey(train, Producto_ID, Cliente_ID, Ruta_SAK)
setkey(test, Producto_ID, Cliente_ID, Ruta_SAK)

mean_total <- mean(train$log_demand) #overall mean
mean_Prod <-  train[, .(mean_prod = mean(log_demand)), by = .(Producto_ID)]
mean_Prod_Ruta <- train[, .(mean_prod_ruta = mean(log_demand)),
                        by = .(Producto_ID, Ruta_SAK)] #mean by product and ruta
mean_Client_Prod_Agencia <- train[, .(mean_client_prod = mean(log_demand)),
                                  by = .(Producto_ID, Cliente_ID, Agencia_ID)] #mean by product, client, agencia

submit <- merge(test, mean_Client_Prod_Agencia, all.x = TRUE, by = c("Producto_ID", "Cliente_ID", "Agencia_ID"))
submit <- merge(submit, mean_Prod_Ruta, all.x = TRUE, by = c("Producto_ID", "Ruta_SAK"))
submit <- merge(submit, mean_Prod, all.x = TRUE, by = "Producto_ID")

preds = submit$mean_client_prod
preds[is.na(preds)] = submit$mean_prod_ruta[is.na(preds)]
preds[is.na(preds)] = submit$mean_prod[is.na(preds)]
preds[is.na(preds)] = mean_total

preds = expm1(preds)

submission = data.frame(id = submit$id, Demanda_uni_equil = preds)
write.csv(submission, "submission.csv", row.names = FALSE)
print("Successfully completed")