# loss_functions_plot.R
require(ggplot2)
require(cowplot)

plot_activation_function <- function(f, title, range, target){
  ggplot(data.frame(x=range), mapping=aes(x=x)) + 
    geom_hline(yintercept=0, colour='red', alpha=1/4) +
    geom_vline(xintercept=0, colour='red', alpha=1/4) +
    stat_function(fun=f, colour = "dodgerblue3",) +
    ggtitle(title) +
    theme(text = element_text(size = 16), plot.title = element_text(hjust = 0.5), axis.text.x = element_text(angle=45, hjust = 1)) +
    xlab("Prediction") + ylab("Loss") +
    scale_x_continuous(breaks=c(0, target, 0.25, 0.5, 0.75, 1))
}

plot_all_activation_functions <- function(f1, f2, f3, title, range, target){
  ggplot(data.frame(x=range), mapping=aes(x=x)) + 
    geom_hline(yintercept=0, color='red', alpha=1/4) +
    geom_vline(xintercept=0, color='red', alpha=1/4) +
    stat_function(fun=f1, aes(colour = "MSE")) +
    stat_function(fun=f2, aes(colour = "MAE")) +
    stat_function(fun=f3, aes(colour = "RMSE")) +
    ggtitle(title) +
    theme(text = element_text(size = 18, face = "bold"),
          plot.title = element_text(hjust = 0.5),
          axis.text.x = element_text(angle=45, hjust = 1),
          legend.position = "top") +
    xlab("Prediction") + ylab("Loss") +
    scale_x_continuous(breaks=c(0, target, 0.25, 0.5, 0.75, 1)) +
    scale_colour_manual("", 
                        breaks = c("MSE", "MAE", "RMSE"),
                        values = c("MSE"="dodgerblue3", "MAE"="red3", "RMSE"="orange1"))
    
}

set.seed(42)
yhats <- runif(32, min=0, max=1)
mse <- function(yhats, y=0) {((y - yhats) ** 2)/2}
p_mse <- plot_activation_function(mse, 'MSE', c(0,1), 0)
  
mae <- function(yhats, y=0) {abs(y - yhats)/2}
p_mae <- plot_activation_function(mae, 'MAE', c(0,1), 0)

rmse <- function(yhats, y=0) {sqrt(((y - yhats) ** 2)/2)}
p_rmse <- plot_activation_function(rmse, 'RMSE', c(0,1), 0)

plot_all_activation_functions(mse, mae, rmse, '', c(0,2), 0)
ggsave("Plots/All_Loss_Functions_Plot.pdf")

p_grid <- cowplot::plot_grid(p_mse, p_mae, p_rmse, ncol = 3)
  # theme(plot.margin = unit(c(1,31,3,3), "lines"))
ggsave("Plots/Loss_Function_Plots.pdf", plot = p_grid)