#!/usr/bin/env sos-runner
# fileformat=SOS1.0
# Author: Gao Wang and Bo Peng

#
# Linear-model based prediction
#
# This script fits linear models
# using Lasso and Ridge regression
# and summarizes their prediction performance
# This script is written in a "mixed" style that
# takes advantage of both "forward" and "backward" style
# workflows

parameter: beta = [3, 1.5, 0, 0, 2, 0, 0, 0]
data_files = [("data_{}.train.csv".format(x + 1), "data_{}.test.csv".format(x + 1)) for x in range(5)]

[simulation: provides = ["data_{id}.train.csv", "data_{id}.test.csv"]]
# Simulate sparse data-sets
depends: R_library("MASS")
parameter: N = (40, 200) # training and testing samples
parameter: rstd = 3
output: "data_${id}.train.csv", "data_${id}.test.csv"
task:
R:
  set.seed(${id})
  N = sum(c(${N!,}))
  p = length(c(${beta!,}))
  X = MASS::mvrnorm(n = N, rep(0, p), 0.5^abs(outer(1:p, 1:p, FUN = "-")))
  Y = X %*% c(${beta!,}) + rnorm(N, mean = 0, sd = ${rstd})
  Xtrain = X[1:${N[0]},]; Xtest = X[(${N[0]}+1):(${N[0]}+${N[1]}),]
  Ytrain = Y[1:${N[0]}]; Ytest = Y[(${N[0]}+1):(${N[0]}+${N[1]})]
  write.table(cbind(Ytrain, Xtrain), ${_output[0]!r}, row.names = F, col.names = F, sep = ',')
  write.table(cbind(Ytest, Xtest), ${_output[1]!r}, row.names = F, col.names = F, sep = ',')

[ridge_1 (model fitting)]
# Ridge regression model implemented in R
# Build predictor via cross-validation and make prediction
parameter: nfolds = 5
input: data_files, group_by = 2
output: "${_input[0]!nn}.ridge.predicted.csv", "${_input[0]!nn}.ridge.coef.csv"
task:
R:
  train = read.csv(${_input[0]!r}, header = F)
  test = read.csv(${_input[1]!r}, header = F)
  model = glmnet::cv.glmnet(as.matrix(train[,-1]), train[,1], family = "gaussian", alpha = 0, nfolds = ${nfolds}, intercept = F)
  betahat = as.vector(coef(model, s = "lambda.min")[-1])
  Ypred = predict(model, as.matrix(test[,-1]), s = "lambda.min")
  write.table(Ypred, ${_output[0]!r}, row.names = F, col.names = F, sep = ',')
  write.table(betahat, ${_output[1]!r}, row.names = F, col.names = F, sep = ',')

[lasso_1 (model fitting)]
# LASSO model implemented in Python
# Build predictor via cross-validation and make prediction
parameter: nfolds = 5
input: data_files, group_by = 2
output: "${_input[0]!nn}.lasso.predicted.csv", "${_input[0]!nn}.lasso.coef.csv"
task:
python:
  import numpy as np
  from sklearn.linear_model import LassoCV
  train = np.genfromtxt(${_input[0]!r}, delimiter = ",")
  test = np.genfromtxt(${_input[1]!r}, delimiter = ",")
  model = LassoCV(cv = ${nfolds}, fit_intercept = False).fit(train[:,1:], train[:,1])
  Ypred = model.predict(test[:,1:])
  np.savetxt(${_output[0]!r}, Ypred)
  np.savetxt(${_output[1]!r}, model.coef_)

[ridge_2, lasso_2 (evaluate)]
# Evaluate predictors by calculating mean squared error
# of prediction vs truth (first line of output)
# and of betahat vs truth (2nd line of output)
depends: data_files
input: group_by = 2, group_with = "data_files"
output: "${_input[0]!nn}.mse.csv"
R:
  b = c(${beta!,})
  Ytruth = as.matrix(read.csv(${_data_files[1]!r}, header = F)[,-1]) %*% b
  Ypred = scan(${_input[0]!r})
  prediction_mse = mean((Ytruth - Ypred)^2)
  betahat = scan(${_input[1]!r})
  estimation_mse = mean((betahat - b) ^ 2)
  cat(paste(prediction_mse, estimation_mse), file = ${_output!r})
 
[default]
sos_run('ridge')
sos_run('lasso')