
## Lookup table with IACS grain colours
grainDict <- data.frame(grain_type = c('PP','DF','RG','FC','DH','SH','MF','IF','FC','PPgp','FCxr','MFcr'),
                        colour = c('#00FF00','#228B22','#FFB6C1','#ADD8E6','#0000FF','#FF00FF','#FF0000','#00FFFF','#ADD8E6','#00FF00','#ADD8E6','#FF0000'), stringsAsFactors = F)


#' Grain colour lookup
#'
#' @param Grains grain type
#'
#' @return R colour code
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#' Grains <- c('PP', 'DF', 'RG', 'FC', 'FCxr', 'DH', 'SH', 'MF', 'MFcr', 'IF')
#' Colours <- getColoursGrainType(Grains)
#' plot(1:length(Grains), col = Colours, pch = 20, cex = 3)
#' text(1:length(Grains), 1:length(Grains), Grains, pos = 1)
#'
#' @export
#'
getColoursGrainType <- function(Grains) {
  Col <- sapply(Grains, function(x) grainDict$colour[grainDict$grain_type == x][1])
  Col[is.na(Col)] <- 'gray'
  return(Col)
}

#' Convert numeric aspects to factors
#'
#' @param Aspects vector of numeric aspects
#' @param Angles vector of slope angles
#'
#' @return vector of aspect factors
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#'
#' Aspects <- c(0, 90, 270, 0)
#' Angles <- c(0, 40, 20, 15)
#'
#' Aspects <- cardinalAspects(Aspects, Angles)
#' print(Aspects)
#'
#' @export
#'
cardinalAspects <- function(Aspects, Angles){
  Dir <- rep(NA, length(Aspects))
  Dir[Angles == 0] <- 'Flat'
  Dir[Aspects == 0 & Angles > 0] <- 'N'
  Dir[Aspects == 90] <- 'E'
  Dir[Aspects == 180] <- 'S'
  Dir[Aspects == 270] <- 'W'
  Dir <- factor(Dir, levels = c('Flat', 'N', 'E', 'S', 'W'))
  return(Dir)
}



#' Summary of a single snowprofile
#'
#' @param Profile snowprofile object
#'
#' @return dataframe
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#' summary(SP2)
#'
#' @export
#'
summary.snowprofile <- function(Profile)
{

  # Initialize dataframe with a 1 row long dummy variable
  Metadata <- data.frame(init = NA)

  # Loop through each element in profile
  for (col in names(Profile))
  {
    Element <- Profile[col]

    # Copy element if it's length is one
    if (lengths(Element) == 1) {
      Metadata[col] <- Element

      # Special treatment for latlon
    } else if (col == "latlon") {
      Metadata$lat <- Profile$latlon[1]
      Metadata$lon <- Profile$latlon[2]

      # Place any summary stats of the layers here
    } else if (col == "layers") {
      Metadata$nLayers <- nrow(Profile$layers)
    }

  }

  # Add cardinal aspect
  Metadata$dir <- cardinalAspects(Profile$aspect, Profile$angle)

  # Delete dummy variable
  Metadata$init <- NULL

  # Return metadata dataframe
  return(Metadata)

}


#' Merge multiple profiles with columns for metadata
#'
#' @importFrom data.table rbindlist
#'
#' @param Profiles list of profiles
#'
#' @return dataframe
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#' Rprofiles <- proRbind(SP2)
#' head(Rprofiles)
#'
#' @export
#'
proRbind <- function(Profiles)
{

  # If it's a single profile place it in a list so below code works
  if (is(Profiles, "snowprofile")) Profiles <- list(Profiles)

  # Confirm profiles is a list of snowprofile objects
  if (any(sapply(Profiles, class) != "snowprofile")) stop("proRbind requires a list of snowprofile objects")

  # Merge metadata with layer data for each profile
  AllProfiles <- lapply(Profiles, function(p) merge(summary(p), p$layers, all.y = T))

  # Rbind into large dataframe
  AllProfiles <- as.data.frame(rbindlist(AllProfiles, fill = T))

  # Return big dataframe
  return(AllProfiles)
}


#' Summarize multiple snowprofiles
#'
#' Wrapper for summary.snowprofile
#'
#' @importFrom data.table rbindlist
#'
#' @param profiles list of snowprofile objects
#'
#' @return dataframe
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#' Metadata <- proSummary(SP2)
#' head(Metadata)
#'
#' @export
#'
proSummary <- function(Profiles)
{
  Summaries <- lapply(Profiles, summary)
  Summaries <- as.data.frame(rbindlist(Summaries, fill = T))
  return(Summaries)
}


#' Plot hardness profile
#'
#' @param Profile snowprofile object
#' @param TempProfile draw unscaled temperature profile (default = TRUE)
#' @param Col vector of colors corresponding to the grain types in the profile (defaults to a lookup table)
#' @param bg background color of the plot (e.g. "white" (default), "transparent", etc..)
#' @param ... other parameters to barplot
#'
#' @examples
#'
#' require(SarpSnowGeneral)
#' plot(SP2)
#' plot(SP2, col = 'black')
#'
#' @export
#'
plot.snowprofile <- function(Profile,
                             TempProfile = TRUE,
                             Col = sapply(Profile$layers$grain_type, getColoursGrainType),
                             bg = 'white',
                             ...)
{
  ## Extract snowprofile layers
  Layers <- Profile$layers

  par(las = T, bg = bg)

  ## Draw horizontal barplot with hardness profile
  barplot(Layers$hardness,
          width = c(Layers$height[1], diff(Layers$height)),
          col = Col,
          horiz = T,
          border = NA,
          space = 0,
          xlim = c(0,5),
          xaxt = 'n',
          xlab = '',
          beside = T,
          ...)

  ## Draw scaled temperature profile
  if (TempProfile) {
    if ('temperature' %in% names(Layers))
    {
      lines(Layers$temperature/(max(Layers$temperature) - min(Layers$temperature) + 0.1)*2 + 2,
            Layers$height, col = 'red')
    }
  }

  ## Add harndess and height axis
  axis(1, at = 1:5, labels = c('F', '4F', '1F', 'P', 'K'))
  axis(2)
}


#' Function for adding vertical grid lines and date axis to chart
#'
#' Function for adding vertical grid lines and date axis to chart.
#' @param DateArray Date array.
#' @param WithXAxis Switch for plotting the x-axis (default=TRUE).
#' @param HighlightDateStart Array with starts of highlighted time period (default=NA).
#' @param HighlightDateEnd Array with ends of highlighted time period (default=HighlightDateStart).
#' @param HighlightCol Array with colors for highlighted time period (default="light grey").
#' @param cex.axis cex.axis (default=1.0).
#' @param cex.lab cex.lab (default=1.0).
#' @export

addTimeGridAndHighlight <- function(DateArray,
                                    WithXAxis=T,
                                    HighlightDateStart=NA,
                                    HighlightDateEnd=HighlightDateStart,
                                    HighlightCol="light grey",
                                    cex.axis=1.0,
                                    cex.lab=1.0) {

  ## Create axis
  DateArray <- seq(from=as.Date(min(DateArray)), to=as.Date(max(DateArray)), by="days")
  Saturdays <- DateArray[weekdays(DateArray)=="Saturday"]
  if(WithXAxis) {
    axis(1, at=Saturdays, labels=format(Saturdays, "%m/%d"), cex.axis=cex.axis)
  }

  ## Expand HighlightCol is necessary
  if(length(HighlightDateStart) > 1 & length(HighlightCol)==1){
    HighlightCol <- rep(HighlightCol, times=length(HighlightDateStart))
  }

  ## Add Highlighted areas
  if(all(!is.na(HighlightDateStart))) {
    HighlightDateStart <- as.Date(HighlightDateStart)
    HighlightDateEnd   <- as.Date(HighlightDateEnd)
    for (i in 1:length(HighlightDateStart)){
      rect(HighlightDateStart[i]-0.5,-1000, HighlightDateEnd[i]+0.5, 1000, col=HighlightCol[i], border=NA)
    }
  }

  ## Vertical gridlines
  abline(v=Saturdays, lty=3, col="dark grey")
  box()

}


#' Plot time series of snow profiles
#'
#' Plot time series of snow profiles
#' @param Profiles List of snow profiles created with getProfiles() function
#' @param Type TYpe of profile to be plotted. So far the following types are available: graintype (default), hardness, temperature, density, grainsize.
#' @param WithXAxis Switch for plotting x-axis (default = TRUE).
#' @param DateStart Start date for plot. If not provided, the function takes the date range from Profiles (default = NA).
#' @param DateEnd End date for plot. If not provided, the function takes the date range from Profiles (default = NA).
#' @param TopDown Option to plot by depth instead of height, zero on top of plot (default = FALSE)
#' @param cex.axis cex.axis (default = 1.0).
#' @param cex.lab cex.lab (default = 1.0).
#' @param ylab Label for y-axis. Default is 'HS (cm)'.
#' @param SiteName SiteName (default = NULL).
#'
#' @examples
#'
#' require(SarpSnowAnalysis)
#'
#' ## Get data
#' Site <- "VIR155942"
#' DateStart <- "2015-12-01"
#' DateEnd <- "2016-04-31"
#'
#' Pro <- getProfiles("GNP", Sites = Site, DateStart = DateStart, DateEnd = DateEnd)
#'
#' ## Plot snowprofile
#' plotTSSnowProfile(Pro, TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked")
#' plotTSSnowProfile(Pro, Type = "density", TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked: DENSITY")
#' plotTSSnowProfile(Pro, Type = "temperature", TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked: SNOW TEMPERATURE")
#' plotTSSnowProfile(Pro, Type = "grainsize", TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked: GRAIN SIZE")
#' plotTSSnowProfile(Pro, Type = "hardness", TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked: HARDNESS")
#' plotTSSnowProfile(Pro, Type = "ssi", TrackingDates = "2016-01-03", SiteName = Site, main = "Seasonal Snowprofile with 2016-01-03 layer tracked: SSI")
#' plotTSSnowProfile(Pro, TopDown = TRUE)
#'
#' @export
#'
plotTSSnowProfile <- function(Profiles,
                              Type = "graintype",
                              WithXAxis = T,
                              DateStart = NA,
                              DateEnd = NA,
                              TopDown = FALSE,
                              OutlineLyrs = FALSE,
                              cex.axis = 1.0,
                              cex.lab = 1.0,
                              ylab = "HS (cm)",
                              SiteName = NULL,
                              main = NA,
                              ylim = NULL,
                              ...) {

  print("... starting plotTSSnowprofile")

  ## Check type
  Type <- tolower(Type)
  if (Type != "graintype" & Type != "hardness" & Type != "density" & !startsWith(Type,  "temp") & Type != "grainsize" & Type != "ssi") {
    stop(paste0("Profile type: '", Type, "' not supported!"))
  }

  ## Unpack profile metadata
  Meta <- proSummary(Profiles)

  ## Assign date range for chart
  if (is.na(DateStart)) DateStart <- substr(min(Meta$datetime), 1, 10)
  if (is.na(DateEnd)) DateEnd <- substr(max(Meta$datetime), 1, 10)
  DateStart <- as.Date(DateStart)
  DateEnd <- as.Date(DateEnd)

  ## Subset data to date range
  Profiles <- Profiles[which(Meta$date >= DateStart & Meta$date <= DateEnd)]
  Meta <- Meta[(Meta$date >= DateStart & Meta$date <= DateEnd),]

  ## Rbind all the profiles to get table with all layers
  Lyrs <- proRbind(Profiles)

  ## Overwrite height with depth values
  if (TopDown) {
    Lyrs$height <- -Lyrs$depth
    ylab = 'Depth (cm)'
  }

  ## Define the bottom of layers
  Lyrs$index <- Lyrs[c('station_id','datetime')]
  Lyrs$bottom <- c(0, Lyrs$height[1:(nrow(Lyrs) - 1)])
  Lyrs$bottom[which(!duplicated(Lyrs$index))] <- 0
  if (TopDown) Lyrs$bottom[which(!duplicated(Lyrs$index))] <- -Lyrs$hs[which(!duplicated(Lyrs$index))]

  ## Select colour scheme
  if (Type == "graintype") {
    Lyrs$Col <- getColoursGrainType(Lyrs$grain_type)
  } else if (Type == "hardness") {
    Lyrs$Col <- getColoursHardness(Lyrs$hardness)
  } else if (Type == "density") {
    Lyrs$Col <- getColoursDensity(Lyrs$density)
  } else if (startsWith(Type, 'temp')) {
    Lyrs$Col <- getColoursSnowTemp(Lyrs$temperature)
  } else if (Type == "grainsize") {
    Lyrs$Col <- getColoursGrainSize(Lyrs$grain_size)
  } else if (Type == "ssi") {
    Lyrs$Col <- getColoursSSI(Lyrs$ssi)
  } else {
    Lyrs$Col <- "#E8E8E8"
  }

  ## Calculate plot dimensions
  if (is.null(ylim)) {
    ymin <- ifelse(TopDown, -max(Meta$hs) - 10, 0)
    ymax <- ifelse(TopDown, 0, max(Meta$hs) + 10)
    ylim <- c(ymin, ymax)
  }


  ## Set up plot
  plot(NA, NA, type = 'n',
       xlim = c(DateStart, DateEnd), ylim = ylim,
       xlab = '', xaxt = 'n', yaxt = 'n', ylab = ylab,
       las = 1, ...)
  addTimeGridAndHighlight(c(DateStart, DateEnd), WithXAxis = WithXAxis, cex.axis = cex.axis, cex.lab = cex.lab)


  ## Draw rectangles
  rect(xleft = Lyrs$date - 0.5,
       ybottom = Lyrs$bottom - ifelse(Type == 'ssi', 0, 0.5),
       xright = Lyrs$date + 0.5,
       ytop = Lyrs$height,
       col = Lyrs$Col,
       border = ifelse(OutlineLyrs == FALSE, NA, "#606060"))


  ## Adding grid lines on top

  ## Date grid
  DateArray <- seq(from = DateStart, to = DateEnd, by = "days")
  Saturdays <- DateArray[weekdays(DateArray) == "Saturday"]
  abline(v = Saturdays, lty = 3, col = "dark grey")

  ## Height grid and y-axis
  HeightGrid <- pretty(ylim, n = 4)
  if (TopDown) {
    axis(2, at = HeightGrid, labels = -HeightGrid)
  } else {
    axis(2, at = HeightGrid, labels = HeightGrid)
  }
  abline(h = HeightGrid, lty = 3, col = "dark grey")
  box()


  ## Adding title
  if (!is.na(main)) {
    title(main)
  }

  ## Adding station name
  if (!is.null(SiteName)) {
    mtext(text = SiteName, side = 3, line = 0.2, adj = 1, cex = cex.axis*0.5)
  }

}


#' Plot a single parameter in multiple profiles side-by-side
#'
#' @param Profiles List of snow profiles created with getProfiles() function
#' @param Type Type of profile to be plotted. So far the following types are available: graintype (default), hardness, temperature, density, grainsize.
#' @param SortMethod Sort profiles along x-axis in existing order (SortMethod = ''), sorted by HS ('hs'), or elevation ('elev')
#' @param TopDown Option to plot by depth instead of height, zero on top of plot (default = FALSE)
#' @param OutlineLyrs
#' @param ylim
#' @param main
#' @param ...
#'
#' @examples
#'
#' require(SarpSnowAnalysis)
#'
#' ## Get data
#' DateStart <- "2016-01-01"
#' Profiles <- getProfiles("GNP", DateStart = DateStart)
#'
#' ## Plot snowprofiles
#' plotSideBySideProfiles(Profiles)
#' plotSideBySideProfiles(Profiles, TopDown = TRUE)
#'
#' @export
#'
plotSideBySideProfiles <- function(Profiles, Type = "graintype", SortMethod = "hs", TopDown = FALSE, OutlineLyrs = FALSE, ylim = NULL, main = NULL, ...) {

  ## Check
  if (Type != "graintype" & Type != "hardness" & Type != "density" & substr(Type, 1, 4) != "temp" & Type != "grainsize" & Type != "ssi") {
    stop(paste0("Profile type: '", Type, "' not supported!"))
  }

  ## Sort profiles
  Metadata <- proSummary(Profiles)
  if (SortMethod == "hs") {
    Profiles <- Profiles[order(Metadata$hs)]
  } else if (SortMethod == "elev") {
    Profiles <- Profiles[order(Metadata$elev)]
  } else if (SortMethod == "") {
    Profiles <- Profiles
  }

  ## Rbind all the profiles to get table with all layers
  Lyrs <- proRbind(Profiles)

  ## Overwrite height with depth values
  if (TopDown) Lyrs$height <- -Lyrs$depth

  ## Add index for each unique profile
  Lyrs$index <- cumsum(!duplicated(Lyrs[c('station_id','datetime')]))

  ## Define the bottom of layers
  Lyrs$bottom <- c(0, Lyrs$height[1:(nrow(Lyrs) - 1)])
  Lyrs$bottom[which(!duplicated(Lyrs$index))] <- 0
  if (TopDown) Lyrs$bottom[which(!duplicated(Lyrs$index))] <- -Lyrs$hs[which(!duplicated(Lyrs$index))]

  ## Select colour scheme
  if (Type == "graintype") {
    Lyrs$Col <- getColoursGrainType(Lyrs$grain_type)
  } else if (Type == "hardness") {
    Lyrs$Col <- getColoursHardness(Lyrs$hardness)
  } else if (Type == "density") {
    Lyrs$Col <- getColoursDensity(Lyrs$density)
  } else if (startsWith(Type, 'temp')) {
    Lyrs$Col <- getColoursSnowTemp(Lyrs$temperature)
  } else if (Type == "grainsize") {
    Lyrs$Col <- getColoursGrainSize(Lyrs$grain_size)
  } else if (Type == "ssi") {
    Lyrs$Col <- getColoursSSI(Lyrs$ssi)
  } else {
    Lyrs$Col <- "#E8E8E8"
  }

  ## Calculate plot dimensions
  if (is.null(ylim)) {
    ymin <- ifelse(TopDown, -max(Metadata$hs) - 10, 0)
    ymax <- ifelse(TopDown, 0, max(Metadata$hs) + 10)
    ylim <- c(ymin, ymax)
  }


  ## Set up plot
  plot(NA, NA, type = 'n',
       xlim = c(0, length(Profiles) + 1), ylim = ylim,
       xlab = '', xaxt = 'n', yaxt = 'n', las = 1, ...)


  ## Draw rectangles
  rect(xleft = Lyrs$index - 0.5,
       ybottom = Lyrs$bottom,
       xright = Lyrs$index + 0.5,
       ytop = Lyrs$height,
       col = Lyrs$Col,
       border = ifelse(OutlineLyrs == FALSE, NA, "#606060"))


  ## Height grid and y-axis
  HeightGrid <- pretty(ylim, n = 4)
  if (TopDown) {
    axis(2, at = HeightGrid, labels = -HeightGrid)
  } else {
    axis(2, at = HeightGrid, labels = HeightGrid)
  }
  abline(h = HeightGrid, lty = 3, col = "dark grey")
  box()

}
