'''
    6 - Output the final image
'''
# Import libraries
from rpy2 import robjects
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri

# Activate pandas2ri to allow Dataframes to be translated between
# languages on the fly
pandas2ri.activate()

# Create a method that allows users to execute r code on the fly
r = robjects.r


# Configuration
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Create your own equation for use in model_train
equation = robjects.Formula('amount~.')

# Select the method to use in caret.train
method = "rf"

# Replace [image] with the path to the image
# Import raster into R globalenv
r('img <- raster::stack("[image]")')

# Set output format
r_format = "raster"

# Set output datatype
r_datatype = "INT1U"


# Output the final image
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
print('* 6 - Outputting final image')
# Reference r libraries in python
caret = importr('caret')

# Get the train model in Python and insert into R globalenv
robjects.globalenv['train_model'] = caret.train(equation, method=method, data=train_full)

# Within the R environment, create the raster and save to ./Data/outputImage
r('raster::predict(object=img, model=train_model, filename="./Data/outputImage", fun=predict, format="%s", datatype="%s", overwrite=TRUE, na.rm=TRUE)' % (r_format, r_datatype))

print('  (Script 6 finished. Raster saved as "./Data/outputImage")')