library(dplyr)
library(data.table)
library(xgboost)

## low memory version of data load - only keeps adjusted demand
train <- fread('../input/train.csv', 
               select = c('Semana', 'Cliente_ID', 'Producto_ID', 
                          'Agencia_ID', 'Ruta_SAK', 'Canal_ID',
                          'Demanda_uni_equil'),
               header = T, sep = ',', data.table = F)
train <- rename(train, week = Semana, client_ID = Cliente_ID, 
                product_ID = Producto_ID, depot_ID = Agencia_ID, 
                route_ID = Ruta_SAK, channel_ID = Canal_ID,
                adjDemx = Demanda_uni_equil)

## use this section instead if you have more memory and want ALL the variables
# train <- fread('../input/train.csv', header = T, sep = ',', data.table = F)
# names(train) <- c('week', 'depot_ID', 'channel_ID', 'route_ID', 'client_ID',
#                    'product_ID', 'salesx', 'salesPesosx', 'returnsx',
#                    'returnsPesosx', 'adjDemx')

# pull time series store/product/route/depot/channel data into single rows
wk <- 3
for (lagName in c('_L7', '_L6', '_L5', '_L4', '_L3', '_L2', '_L1')) { # time lagged column names
  addDF <- filter(train, week == wk) %>% select(-week)
  train <- filter(train, week != wk) # remove week and save some memory
  colnames(addDF) <- gsub('x$', lagName, colnames(addDF)) # change colnames to "lagged" names
  if (wk == 3) {
    featDF <- addDF
  } else {
    featDF <- full_join(featDF, addDF) 
  }
  wk <- wk + 1
}

# load test data set
test <-  fread('../input/test.csv', header = T, sep = ',', data.table = F)
names(test) <- c("id", "week", "depot_ID", "channel_ID", "route_ID", "client_ID",
                 "product_ID") 

### grab just the test rows from training set
testRows <- group_by(test, depot_ID, channel_ID, route_ID, client_ID, product_ID) %>%
  summarise(sumWeeks = sum(week)) %>% ungroup() 
testRows$adjDem_p1 <- NA
testRows$adjDem_p2 <- NA
testRows$adjDem_p1[(testRows$sumWeeks == 10) | (testRows$sumWeeks == 21)] <- 0 # 0 if week10 prediction needed
testRows$adjDem_p2[(testRows$sumWeeks == 11) | (testRows$sumWeeks == 21)] <- 0 # 0 if week11 prediction needed
testRows$sumWeeks = NULL # temporary column for above logic, remove
submitDF <- left_join(testRows, featDF)
rm(featDF) # free up memory

# filter rows that need week 10 prediction and label as week = 10
predDF_week10 <- filter(submitDF, !is.na(adjDem_p1)) %>% mutate(week = 10)

# most basic feature set
model_features <- c("adjDem_L7", "adjDem_L6", "adjDem_L5", 
                    "adjDem_L4", "adjDem_L3", "adjDem_L2")

# prepare feature and label matrices for xgboost
X <- as.matrix(predDF_week10[, model_features])
X[is.na(X)] <- 0
Y <- log1p(predDF_week10$adjDem_L1)
X <- X[!is.na(Y),]
Y <- Y[!is.na(Y)]

model <- xgboost(data = X, label = Y, nrounds = 25)

# make predictions, use features from week+1
predict_features <- c("adjDem_L6", "adjDem_L5", "adjDem_L4", 
                      "adjDem_L3", "adjDem_L2", "adjDem_L1")
X <- as.matrix(predDF_week10[, predict_features])
X[is.na(X)] <- 0
predDF_week10$Demanda_uni_equil <- expm1(predict(model, X))

# build submission file
submitDF <- left_join(test, predDF_week10) %>% 
  select(id, Demanda_uni_equil)
submitDF$Demanda_uni_equil[is.na(submitDF$Demanda_uni_equil)] = 0 # week 11 "prediction"

# write result
write.csv(submitDF, 'submit.csv', row.names = FALSE)