### =========================================================================
### XString Display Colors
### -------------------------------------------------------------------------

### Return a named character vector where all the names are single letters.
### Colors for A, C, G, and T were inspired by
###   https://en.wikipedia.org/wiki/Nucleotide#Structure
### Called in .onLoad() to initialize DNA_AND_RNA_COLORED_LETTERS.
make_DNA_AND_RNA_COLORED_LETTERS <- function()
{
    ## Not sure why but the built-in white() style in the crayon package
    ## produces some kind of light grey text color. So we define a style
    ## that produces a text color that is 100% white.
    whiter <- make_style(rgb(1, 1, 1))
    dark_grey_bg <- make_style(rgb(0.5,0.5,0.5), bg=TRUE)

    ## All the IUPAC ambiguity letters minus N.
    dark_grey_bg_letters <- c("M", "R", "W", "S", "Y", "K", "V", "H", "D", "B")

    c(
        A=make_style(rgb(1, 0.5, 0.5), bg=TRUE)(black("A")),
        C=make_style(rgb(0.5, 1, 0.5), bg=TRUE)(black("C")),
        G=make_style(rgb(0.5, 1, 1), bg=TRUE)(black("G")),
        T=make_style(rgb(1, 0.8, 0.5), bg=TRUE)(black("T")),
        U=make_style(rgb(1, 1, 0.5), bg=TRUE)(black("U")),
        setNames(sprintf(dark_grey_bg(whiter("%s")), dark_grey_bg_letters),
                 dark_grey_bg_letters),
        N=make_style("grey", bg=TRUE)(whiter("N"))
    )
}

### Return a named character vector where all the names are single letters.
### Colors amino acids by similarity
### Colors groupins by
###   https://www.jalview.org/help/html/colourSchemes/zappo.html
### Called in .onLoad() to initialize AA_COLORED_LETTERS.
make_AA_COLORED_LETTERS <- function(){
    whiter <- make_style(rgb(1, 1, 1))
    dark_grey_bg <- make_style(rgb(0.5,0.5,0.5), bg=TRUE)

    ## All the IUPAC ambiguity letters minus X.
    dark_grey_bg_letters <- c("U","O","B","J","Z")

    cp <- c("#fbf8cc", "#ffcfd2", "#cfbaf0",
            "#a3c4f3", "#8eecf5", "#b9fbc0", "#f1c0e8")

    c(
        # Cysteine
        C=make_style(cp[1], bg=TRUE)(black("C")),

        # Aliphatic/hydrophobic
        A=make_style(cp[2], bg=TRUE)(black("A")),
        V=make_style(cp[2], bg=TRUE)(black("V")),
        M=make_style(cp[2], bg=TRUE)(black("M")),
        L=make_style(cp[2], bg=TRUE)(black("L")),
        I=make_style(cp[2], bg=TRUE)(black("I")),

        # Conformationally Special
        P=make_style(cp[3], bg=TRUE)(black("P")),
        G=make_style(cp[3], bg=TRUE)(black("G")),

        # Positive
        K=make_style(cp[4], bg=TRUE)(black("K")),
        R=make_style(cp[4], bg=TRUE)(black("R")),
        H=make_style(cp[4], bg=TRUE)(black("H")),

        # Hydrophilic
        N=make_style(cp[5], bg=TRUE)(black("N")),
        T=make_style(cp[5], bg=TRUE)(black("T")),
        Q=make_style(cp[5], bg=TRUE)(black("Q")),
        S=make_style(cp[5], bg=TRUE)(black("S")),

        # Aromatic
        F=make_style(cp[6], bg=TRUE)(black("F")),
        Y=make_style(cp[6], bg=TRUE)(black("Y")),
        W=make_style(cp[6], bg=TRUE)(black("W")),

        # Negative
        E=make_style(cp[7], bg=TRUE)(black("E")),
        D=make_style(cp[7], bg=TRUE)(black("D")),

        # Ambiguity
        setNames(sprintf(dark_grey_bg(whiter("%s")), dark_grey_bg_letters),
                 dark_grey_bg_letters),

        # Any code
        X=make_style("grey", bg=TRUE)(whiter("X"))
    )
}

### 'x' must be a character vector.
## env_var_name is the name of the corresponding palette in .pkgenv
.add_xstring_colors <- function(x, env_var_name)
{
    if (!isTRUE(getOption("Biostrings.coloring", default=FALSE)))
        return(x)
    color_palette <- get(env_var_name, envir=.pkgenv)
    ans <- vapply(x,
        function(xi) {
            xi <- safeExplode(xi)
            m <- match(xi, names(color_palette))
            match_idx <- which(!is.na(m))
            xi[match_idx] <- color_palette[m[match_idx]]
            paste0(xi, collapse="")
        },
        character(1),
        USE.NAMES=FALSE
    )
    x_names <- names(x)
    if (!is.null(x_names))
        names(ans) <- x_names
    ans
}

.update_X_palette <- function(colors=NULL, env_var_name,
                                alphabet, default_palette_function){
    ## passing default_palette_function as a function pointer so we don't
    ## have to evaluate it unless necessary
    palette <- get(env_var_name, envir=.pkgenv)
    if(is.null(colors))
        palette <- default_palette_function()
    if(!is.null(colors)){
        if(!is.list(colors)){
            stop("'colors' should be NULL or a named list of entries with 'bg' ",
                    "and optionally 'fg' values.")
        }

        n <- names(colors)
        if(!is.null(alphabet) && length(setdiff(n, alphabet)) != 0){
            ## non-BStrings: checking if the characters are valid
            stop("Invalid codes specified.")
        } else if(is.null(alphabet)){
            ## BStrings: checking for single characters (0:255 in raw)
            name_nchars <- vapply(n, \(x) length(charToRaw(x)), integer(1L))
            if(!all(name_nchars == 1L))
                stop("Invalid codes specified.")
        }

        for(i in seq_along(colors)){
            fg <- colors[[i]]$fg
            bg <- colors[[i]]$bg
            if(is.null(fg) && is.null(bg)){
                palette[n[i]] <- n[i]
            } else if(is.null(bg)) {
                palette[n[i]] <- make_style(fg)(n[i])
            } else {
                if(is.null(fg)) fg <- rgb(1,1,1)
                palette[n[i]] <- make_style(bg, bg=TRUE)(make_style(fg)(n[i]))
            }
        }
    }

    assign(env_var_name, palette, envir=.pkgenv)
}

update_DNA_palette <- function(colors=NULL){
    .update_X_palette(colors, "DNA_AND_RNA_COLORED_LETTERS",
                        union(DNA_ALPHABET, RNA_ALPHABET),
                        make_DNA_AND_RNA_COLORED_LETTERS)
}

update_RNA_palette <- update_DNA_palette

update_AA_palette <- function(colors=NULL){
    .update_X_palette(colors, "AA_COLORED_LETTERS",
                        AA_ALPHABET,
                        make_AA_COLORED_LETTERS)
}

update_B_palette <- function(colors=NULL){
    ## BStrings don't have a default palette
    ## thus their default palette function is just \() return(character(0L))
    .update_X_palette(colors, "B_COLORED_LETTERS",
                        NULL,
                        \(){ character(0L) })
}

add_colors <- function(x) UseMethod("add_colors")
add_colors.default <- identity
add_colors.DNA <- add_colors.RNA <- function(x){ .add_xstring_colors(x, "DNA_AND_RNA_COLORED_LETTERS") }
add_colors.AA <- function(x){ .add_xstring_colors(x, "AA_COLORED_LETTERS") }
add_colors.B <- function(x) { .add_xstring_colors(x, "B_COLORED_LETTERS") }
