"""Machine learning processes"""
from pyppl import Proc, Box
from . import params, rimport
from . import delefactory, procfactory
from modkit import Modkit
Modkit().delegate(delefactory())

@procfactory
def _pRegressTrain():
	"""
	@name:
		pRegressTrain
	@description:
		Train a regression model
	@input:
		`infile:file`: The input file (regression done on columns)
		`fmfile:file`: The file defining the regression formula.
			```
			Case1   lm   Y ~ X1 + X2
			Case2   glm  Y ~ X1 + X2 + X1:X2
			```
			- `Case` column can be omitted, if so, `Case1`, `Case2` will be added automatically.
			- No head.
		`casefile:file`: Defining cases:
			```
			R1   Case1
			R2   Case1
			... ...
			Rm   Case1
			Rm   Case2
			Rm+1 Case2
			... ...
			Rn   Case2
			```
			- One row can bed used for multiple cases
			- No head.
	@output:
		`outfile:file`: The output file with the summary of the models
		`outdir:dir`  : The output directory containing the summaries, models and plots
	@args:
		`plot`   : Plot the model? Default: `False`
		`inopts` : Input options for input file. Default: `Box(cnames = True, rnames = True)`
		`devpars`: The device parameters for the plot. Default: `Box(res = 300, height = 4000, width = 4000)`
	"""
	pRegressTrain        = Proc(desc = 'Train a regression model')
	pRegressTrain.input  = 'infile:file, fmfile:file, casefile:file'
	pRegressTrain.output = [
		'outfile:file:{{i.infile | fn2}}.regress/{{i.infile | fn2}}.regress.txt',
		'outdir:dir:{{i.infile | fn2}}.regress'
	]
	pRegressTrain.args.plot    = False
	pRegressTrain.args.cov     = ''
	pRegressTrain.args.yval    = 'numeric'
	pRegressTrain.args.save    = True
	pRegressTrain.args.glmfam  = 'binomial'
	pRegressTrain.args.inopts  = Box(rnames = True, cnames = True)
	pRegressTrain.args.devpars = Box(res = 300, height = 4000, width = 4000)
	pRegressTrain.envs.rimport = rimport
	pRegressTrain.lang         = params.Rscript.value
	pRegressTrain.script       = "file:scripts/mlearn/pRegressTrain.r"
	return pRegressTrain

@procfactory
def _pRegressPred():
	"""
	@name:
		pRegressPred
	@description:
		Use a trained linear regression model to predict
	@input:
		`infile:file`: The input file
		`model:file` : The trained model by `pRegressTrain`.
			- It can also be a directory containing models generated by `pRegressTrain`.
			- We will scan the model with file name: `*.<model>.rds`
			- Models have to be on the same column.
	@output:
		`outfile:file`: The output file, a copy of infile with Case, predicted value and items in `args.out` added.
		`outdir:dir`  : The output directory
	@args:
		`inopts` : The input options.
		`out`    : What to include in the output file
			- `prob`: Output the probabilities. Default: `True`
			- `auc` : Output AUCs. Default: `True` (only when Y column exists in infile)
			- If `auc` is True, then ROC is anyway plotted.
		`plot`   : Parameters to plot the ROC. Use `False` to disable. Default: `False`
			- Enabled only when Y column exists in infile
			- Could be a dict (Box) with keys:
			- `labels`:  show labels or not
			- `showAUC`: show AUC or not
			- `combine`: combine all roc in one figure?
		`ggs`    : The ggs items to be added for the plot. Default:
			```Box({
				'style_roc': {},
				# show legend at bottom right corner
				'theme#auc': {'legend.position': [1, 0], 'legend.justification': [1, 0]}
			})```
		`devpars`: The device parameters for the plot. Default: `Box(res = 300, height = 2000, width = 2000)`
		`inopts` : Options for reading the input file. Default: `Box(cnames = True, rnames = True, delimit = "\t")`
	"""
	pRegressPred           = Proc(desc = 'Use a trained linear regression model to predict')
	pRegressPred.input     = 'infile:file, model:file'
	pRegressPred.output    = [
		'outfile:file:{{i.infile | fn2}}.pred/{{i.infile | fn2}}.pred.txt',
		'outdir:dir:{{i.infile | fn2}}.pred'
	]
	pRegressPred.args.out  = Box(prob = True, auc = True)
	pRegressPred.args.plot = Box(labels = False, showAUC = True, combine = True)
	pRegressPred.args.ggs  = Box({
		'style_roc': {},
		# show legend at bottom right corner
		'theme#auc': {'legend.position': [1, 0], 'legend.justification': [1, 0]}
	})
	pRegressPred.args.devpars = Box(res = 300, height = 2000, width = 2000)
	pRegressPred.args.inopts  = Box(
		cnames  = True,
		rnames  = True,
		delimit = "\t"
	)
	pRegressPred.envs.rimport = rimport
	pRegressPred.lang         = params.Rscript.value
	pRegressPred.script       = "file:scripts/mlearn/pRegressPred.r"
	return pRegressPred

@procfactory
def _pGlmTrain():
	"""
	@input:
		infile: The input file (Last column as Y)
	@output:
		outmodel: The output model (RData file)
		outdir  : The output directory containing model, plots and other files
	@args:
		plot   : Whether plot the glm probability. Default: `True`
		formula: The formula to perform the regression. Default: `None`.
			- If `None`, will use all first N-1 columns as features.
		inopts : The input options.
		yval   : The type of y values. Default: `categ`
			- `categ`  : categorical values
			- `prob`   : probabilities
			- `numeric`: numeric values
		devpars: The device parameters for the plot.
	"""
	return Box(
		desc   = 'Train a logistic regression model',
		lang   = params.Rscript.value,
		input  = 'infile:file',
		output = [
			'outmodel:file:{{i.infile | stem}}.glm/{{i.infile | stem}}.glm.RDS',
			'outdir:dir:{{i.infile | stem}}.glm'
		],
		args = Box(
			plot    = True,
			formula = None,
			devpars = Box(res = 300, height = 2000, width = 2000),
			params  = Box(family = 'binomial'),
			ggs     = Box(
				geom_smooth = Box({
					"method"     : "glm",
					"method.args": Box(family = "binomial"),
					"se"         : True
				})
			),
			inopts  = Box(
				cnames  = True,
				rnames  = True,
				delimit = "\t"
			),
			yval    = 'categ'))

@procfactory
def _pGlmTest():
	"""
	@input:
		`infile:file`: The input file
		`model:file` : The trained model by `pLogitRegTrain`
	@output:
		`outdir:dir`: The output directory
	@args:
		`inopts` : The input options.
		`outprob`: Also output probabilities? Default: True
	"""
	return Box(
		desc   = 'Test trained logistic regression model',
		lang   = params.Rscript.value,
		input  = 'infile:file, model:file',
		output = 'outdir:dir:{{i.infile | stem}}.glm.test',
		args   = Box(
			outprob = True,
			outauc  = True,
			params  = Box(),
			ggs     = Box(),
			devpars = Box(res = 300, height = 2000, width = 2000),
			inopts  = Box(
				cnames  = True,
				rnames  = True,
				delimit = "\t"
			)))

@procfactory
def _pRandomForestTrain():
	"""
	@input:
		`infile:file`: The input file (Last column as Y)
	@output:
		`outmodel:file`: The output model (RData file)
		`outdir:dir`   : The output directory containing model, plots and other files
	@args:
		`plot`   : Whether plot the feature importance. Default: `True`
		`formula`: The formula to perform the regression. Default: `None`.
			- If `None`, will use all first N-1 columns as features.
		`inopts` : The input options.
		`na`     : Replace NAs with? Default: `0`
		`devpars`: The device parameters for the plot.
	@requires:
		`r-randomForst`
	"""
	pRandomForestTrain = Proc(desc = "Train a random forest model.")
	pRandomForestTrain.input = 'infile:file'
	pRandomForestTrain.output = [
		'outmodel:file:{{i.infile | stem}}.rforest/{{i.infile | stem}}.rforest.rds',
		'outdir:dir:{{i.infile | stem}}.rforest'
	]
	pRandomForestTrain.args.plot    = True
	pRandomForestTrain.args.formula = None
	pRandomForestTrain.args.na      = 0
	pRandomForestTrain.args.devpars = Box(res = 300, height = 2000, width = 2000)
	pRandomForestTrain.args.inopts  = Box(
		cnames  = True,
		rnames  = True,
		delimit = "\t"
	)
	pRandomForestTrain.envs.rimport = rimport
	pRandomForestTrain.lang         = params.Rscript.value
	pRandomForestTrain.script       = "file:scripts/mlearn/pRandomForestTrain.r"
	return pRandomForestTrain

@procfactory
def _pDecisionTreeTrain():
	"""
	@name:
		pDecisionTreeTrain
	@description:
		Train a decision tree model
	@input:
		`infile:file`: The input file (Last column as Y)
	@output:
		`outmodel:file`: The output model (RData file)
		`outdir:dir`   : The output directory containing model, plots and other files
	@args:
		`plot`   : Whether plot the feature importance. Default: `True`
		`formula`: The formula to perform the regression. Default: `None`.
			- If `None`, will use all first N-1 columns as features.
		`inopts` : The input options.
		`na`     : Replace NAs with? Default: `0`
		`devpars`: The device parameters for the plot. Default: `Box(res = 300, height = 2000, width = 2000)`
	@requires:
		`r-rpart`
	"""
	pDecisionTreeTrain = Proc(desc = "Train a decision tree model")
	pDecisionTreeTrain.input = 'infile:file'
	pDecisionTreeTrain.output = [
		'outmodel:file:{{i.infile | stem}}.dtree/{{i.infile | stem}}.dtree.rds',
		'outdir:dir:{{i.infile | stem}}.dtree'
	]
	pDecisionTreeTrain.args.plot    = True
	pDecisionTreeTrain.args.formula = None
	pDecisionTreeTrain.args.na      = 0
	pDecisionTreeTrain.args.devpars = Box(res = 300, height = 2000, width = 2000)
	pDecisionTreeTrain.args.inopts  = Box(
		cnames  = True,
		rnames  = True,
		delimit = "\t"
	)
	pDecisionTreeTrain.envs.rimport = rimport
	pDecisionTreeTrain.lang         = params.Rscript.value
	pDecisionTreeTrain.script       = "file:scripts/mlearn/pDecisionTreeTrain.r"
	return pDecisionTreeTrain

@procfactory
def _pCrossValid():
	"""
	@input:
		`infile:file`: The input data file.
	@output:
		`outmodel:file`: The trained model in RDS format
		`outdir:dir`   : The output directory containing output model and plots.
	@args:
		`inopts` : The options to read the input file.
		`ctrl`   : Arguments for `trainControl`.
			- See https://topepo.github.io/caret/model-training-and-tuning.html#the-traincontrol-function
			- Available methods: "boot", "cv", "LOOCV", "LGOCV", "repeatedcv", "timeslice", "none" and "oob"
		`train`  : Arguments for `train` other than `data` and `trControl`.
			- See https://topepo.github.io/caret/model-training-and-tuning.html#basic-parameter-tuning
			- form: formula of the model,
			- method: rf, glm, etc
		`seed`   : The seed.
		`nthread`: # threads to use.
		`plots`  : Do types of plots.
			- `varimp` also available
			- You can also concatenate them using comma (`,`)
	@requires:
		`r-caret`
	"""
	return Box(
		desc   = 'Cross validation on the model',
		lang   = params.Rscript.value,
		input  = 'infile:file',
		output = [
			'outmodel:file:{{i.infile | fn2}}.{{args.train.method}}/{{i.infile | fn2}}.{{args.train.method}}.RDS',
			'outdir:dir:{{i.infile | fn2}}.{{args.train.method}}'
		],
		args = Box(inopts  = Box(cnames = True, rnames = True),
			ctrl    = Box(method = '', savePredictions = True, classProbs = True, verboseIter = True),
			train   = Box(form = None, method = '', metric = 'ROC'),
			seed    = None,
			nthread = 1,
			plots   = ['model', 'roc'], # varimp
			devpars = Box(res = 300, height = 2000, width = 2000)))

