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

### Load dataset
# Log-transform unit sales so we're working with something closer to what
# will be scored. I'm not sure what to do with negative values, so I'm
# just setting their log-transform to zero.
train_data <- fread('../input/train.csv',header=TRUE,sep=',') %>% 
  mutate(log_sales=ifelse(unit_sales > 0,log(unit_sales+1),0)) %>%
  select(-unit_sales)

test_data <- fread('../input/test.csv',header=TRUE,sep=',')

print('DEBUG 1')
### Feature engineering 1: Date-based features
all_dates <- data.frame(date=seq(ymd('2013-01-01'),
                                 ymd('2017-08-31'),
                                 by='1 day')) %>%
  mutate(year=lubridate::year(date),
         month=lubridate::month(date),
         wday=lubridate::wday(date),
         payday=(lubridate::month(date+1) > lubridate::month(date)) |
           lubridate::day(date) == 15,
         payday_close=data.table::shift(payday,n=0:7,fill=0) %>% 
           data.frame %>% rowSums,
         date=as.character(date)) %>%
  select(-payday) # Interested in the week after payday

holidays <- read.csv('../input/holidays_events.csv') %>%
  filter(locale=='National') %>%
  filter(transferred=='False') %>% 
  mutate(date=as.character(date)) %>%
  select(date,type)
all_dates <- left_join(all_dates,holidays,by='date') %>%
  mutate(type=as.character(type),
         type=ifelse(is.na(type),'None',type)) %>%
  group_by(date) %>% # need this to deal with multiple events on same date
  summarize(year=first(year),month=first(month),
            wday=first(wday),type=first(type)) %>%
  mutate(type=as.factor(type))


oil <- read.csv('../input/oil.csv') %>%
  mutate(date=as.character(date))
all_dates <- left_join(all_dates,oil,by='date')
# Need to impute missing oil prices; this is a simple way to do it
all_dates$imputed <- sapply(1:nrow(all_dates), function(i) {
  start <- max(1,i-3)
  end <- min(i+3,nrow(all_dates))
  all_dates[start:end,'dcoilwtico'] %>% colMeans(na.rm=TRUE)
})
all_dates <- all_dates %>% 
  mutate(dcoilwtico=ifelse(is.na(dcoilwtico),imputed,dcoilwtico)) %>%
  select(-imputed)
print('DEBUG 2')
### Feature engineering 2: Store-based features
stores <- read.csv('../input/stores.csv') %>%
  mutate(cluster=as.factor(cluster)) %>%
  select(store_nbr,type,cluster)

### Feature engineering 3: Item-based features
items <- read.csv('../input/items.csv') %>%
  select(-class) # too many for a reasonable factor variable

### Join everything to test and training data
join_data <- function(d) {
  d %>%
    left_join(all_dates,by='date') %>%
    left_join(stores,by='store_nbr') %>%
    left_join(items,by='item_nbr') %>%
    select(-id,-store_nbr,-item_nbr)
}

### Convert to sparse matrices for xgBoost 
df2sparse <- function(d) {
  has_na <- names(d)[is.na(d) %>% colSums > 0]
  for (n in has_na) {
    d[,n] <- ifelse(is.na(d[,n]),'NA',as.character(d[,n])) %>% 
      factor(levels=levels(d[,n]))
  }
  model.matrix(~ .+0, data=d, 
               contrasts.arg = lapply(d[,sapply(d, is.factor)], 
                                      contrasts, contrasts=FALSE)) %>%
    Matrix(sparse=TRUE)
}
print('DEBUG 3')
train <- train_data %>%
  tail(nrow(test_data)*7/3)  # keep enough for 70:30 ratio with test

# monitor training using the 8/16-8/31 window from another year
watch_test <- train_data %>%
  filter(ymd(date) >= ymd('2016-08-16')) %>%
  filter(ymd(date) <= ymd('2016-08-31')) 
  
# need to free up some memory
rm(train_data)
  
print('DEBUG 4')

train_sp <- train %>%
  join_data %>%
  select(-date,-log_sales) %>%
  df2sparse

watch_test_sp <- watch_test %>%
  join_data %>%
  select(-date,-log_sales) %>%
  df2sparse

### Set up and train xgBoost model
d_train <- xgb.DMatrix(data = train_sp,
                       label = train$log_sales)
d_watch_test <- xgb.DMatrix(data = watch_test_sp,
                            label = watch_test$log_sales)
                            
print('DEBUG 5')

watchlist <- list(train=d_train, test=d_watch_test)
xgb1 <- xgb.train(params=list(max_depth=7),
                  data = d_train, 
                  watchlist=watchlist,
                  objective = "reg:linear",
                  nthread=8, # for AWS r4.2xlarge
                  nrounds=100)

### Prepare output for submission

test_sp <- test_data %>%
  join_data %>%
  select(-date) %>%
  df2sparse

pred <- predict(xgb1, test_sp)
output <- data.frame(id=test_data$id,unit_sales=exp(pred)-1)
write.csv(output,'xgb_kernel.csv',row.names=FALSE,quote=FALSE)