#' Remove columns
#'
#' @param kable_input Output of [knitr::kable()] with format specified
#' @param columns A numeric value or vector indicating in which column(s) rows
#' need to be removed
#'
#' @export
#'
#' @examples
#' remove_column(kable(mtcars), 1)
remove_column <- function (kable_input, columns) {
    if (is.null(columns)) return(kable_input)
    kable_format <- attr(kable_input, "format")
    if (!kable_format %in% c("html", "latex")) {
        warning("Please specify format in kable. kableExtra can customize",
                " either HTML or LaTeX outputs. See ",
                "https://haozhu233.github.io/kableExtra/ for details.")
        return(kable_input)
    }

    columns <- sort(unique(columns))
    if (kable_format == "html") {
        return(remove_column_html(kable_input, columns))
    } else if (kable_format == "latex") {
        stop("Removing columns was not implemented for latex kables yet")
    }
}

remove_column_html <- function (kable_input, columns) {
    kable_attrs <- attributes(kable_input)
    kable_xml <- kable_as_xml(kable_input)
    kable_tbody <- xml_tpart(kable_xml, "tbody")
    kable_thead <- xml_tpart(kable_xml, "thead")

    group_header_rows <- attr(kable_input, "group_header_rows")
    all_contents_rows <- seq(1, length(xml_children(kable_tbody)))

    if (!is.null(group_header_rows)) {
        warning("It's recommended to use remove_column after add_header_above.",
                "Right now some column span numbers might not be correct. ")
        all_contents_rows <- all_contents_rows[!all_contents_rows %in%
                                                   group_header_rows]
    }

    collapse_matrix <- attr(kable_input, "collapse_matrix")
    collapse_columns <- NULL
    if (!is.null(collapse_matrix)) {
        collapse_columns <- sort(as.numeric(sub("x", "",
                                                names(collapse_matrix))))
        collapse_columns_origin <- collapse_columns
    }

    while (length(columns) > 0) {
        xml2::xml_remove(xml2::xml_child(
            xml2::xml_child(kable_thead, xml2::xml_length(kable_thead)),
            columns[1]))
        if (length(collapse_columns) != 0 && collapse_columns[1] <= columns[1]){
            if (columns[1] %in% collapse_columns) {
                column_span <- collapse_matrix[[paste0('x', columns[1])]]
                non_skip_rows <- column_span != 0
                collapse_columns <- collapse_columns[
                    collapse_columns != columns[1]
                    ] - 1
            } else {
                non_skip_rows <- rep(TRUE, length(all_contents_rows))
            }
            prior_col <- which(collapse_columns_origin < columns[1])
            for (i in all_contents_rows[non_skip_rows]) {
                if (length(prior_col) == 0) {
                    pos_adj <- 0
                } else {
                    pos_adj <- sum(collapse_matrix[i, prior_col] == 0)
                }
                target_cell <- xml2::xml_child(
                    xml2::xml_child(kable_tbody, i),
                    columns[1] - pos_adj)
                xml2::xml_remove(target_cell)
            }
        } else {
            for (i in all_contents_rows) {
                target_cell <- xml2::xml_child(
                    xml2::xml_child(kable_tbody, i),
                    columns[1])
                xml2::xml_remove(target_cell)
            }
        }
        # not very efficient but for finite task it's probably okay
        columns <- (columns - 1)[-1]
    }

    # head_row <- xml2::xml_child(kable_thead, xml2::xml_length(kable_thead))
    # ncols <- xml2::xml_length(head_row)
    # body_nrows <- xml2::xml_length(kable_tbody)
    #
    # rowspan = matrix(1, nrow = body_nrows, ncol=ncols)
    # for(i in 1:body_nrows){
    #     target_row <- xml2::xml_child(kable_tbody, i)
    #     target_ncols <- xml2::xml_length(target_row)
    #     for(j in 1:target_ncols){
    #         target_cell <- xml2::xml_child(target_row, j)
    #         span <- as.numeric(xml2::xml_attr(target_cell, "rowspan"))
    #         span[is.na(span)] <- 0
    #         if(span>0){
    #             rowspan[i,j]=1
    #             rowspan[i+seq(from=1, to=span-1),j]=0
    #         }
    #     }
    # }
    #
    # for(i in 1:body_nrows){
    #     target_row <- xml2::xml_child(kable_tbody, i)
    #     for(j in rev(columns)){
    #         target_cell <- xml2::xml_child(target_row, j)
    #         if(rowspan[i,j]==1)
    #         xml2::xml_remove(target_cell)
    #     }
    # }
    #
    # for(j in columns){
    #     target_cell_head <- xml2::xml_child(head_row, j)
    #     xml2::xml_remove(target_cell_head)
    # }
    out <- as_kable_xml(kable_xml)
    attributes(out) <- kable_attrs
    if (!"kableExtra" %in% class(out))
        class(out) <- c("kableExtra", class(out))

    return(out)
}