#' Annotate RNA duplexes with features
#'
#' Overlays RNA duplexes with `GRanges` annotation object.
#'
#' @param gi `GInteraction` object to annotate
#' @param anno_gr `GRanges` object with the `keys` columns in the metadata
#' @param keys vectors with the names of the features to use for annotation.
#' @param save_ambig When RNA duplex overlaps multiple features, ambiguous annotation for a single key
#' can be stored in `ambig_list.A` and `ambig_list.B`. `ambig.A` and `ambig.B' fields will be added as the 0/1 label
#' @param ambig_key a which feature to use for recording the annotation ambiguity. Determines the values in `ambig_list.A` and `ambig_list.B`.
#' @param order_key Experimental. In case RNA duplex overlaps multiple features, this key will be used to sort the overlapping features.
#' @param order_vec Experimental. An ordered vector of values in `order_key` annotation feature, which sets the priority in case of feature overlap.
#' @export
#' @return `GInteractions` object with new fields
#' @details
#' For each annotation feature in `keys`, i.e if keys=c(keyname1),
#' then `<keyname1>.A`, `<keyname1>.B` annotation fields will be created, containing the
#' names of overlapping features
#' If no overlap is found for the feature, then filed will have NA
#' @examples
#' data("RNADuplexesSampleData")
#' annotateGI(gi = RNADuplexSampleDGs, anno_gr = SampleGeneAnnoGR)
#'
#' # Prioritisation of the snRNA and lncRNA before mRNA if genes overlap
#' annotateGI(
#'     gi = RNADuplexSampleDGs, anno_gr = SampleGeneAnnoGR,
#'     keys = c("gene_id", "gene_name", "gene_type"),
#'     order_key = "gene_type",
#'     order_vec = c("snRNA", "lncRNA", "protein_coding"),
#'     save_ambig = TRUE
#' )
annotateGI <- function(gi, anno_gr,
    keys = c("gene_name", "gene_type", "gene_id"),
    ambig_key = keys[1],
    save_ambig = TRUE,
    order_key = NULL,
    order_vec = NULL) {
    orininal_columns <- colnames(mcols(gi))

    if (!all(keys %in% colnames(mcols(anno_gr)))) {
        ms <- str_c(setdiff(keys, colnames(mcols(anno_gr))), collapse = ", ")
        message()
        stop("Not all provided keys are found in the anotation GRanges \nmissing keys: ", ms)
    }
    anno_gr <- anno_gr[, keys]
    anno_gr$feature_id <- seq_len(length(anno_gr))
    gi@regions$region_id <- seq_len(length(gi@regions))
    po <- findOverlapPairs(gi@regions, anno_gr)

    # df will be grouped by region from below
    df <- tibble(
        "region_id" = po@first$region_id,
        "overlap_w" = width(pintersect(po))
    ) %>%
        bind_cols(data.frame(mcols(po@second))) %>%
        group_by(region_id) %>%
        mutate(ambig = vctrs::vec_cast(n_distinct(.data[[ambig_key]], na.rm = TRUE) > 1, integer()))


    # save ambig mapping regions (use pre-defined key to save ids)
    if (save_ambig) {
        message("Will save all redundant values in: ", ambig_key)
        ambig_save <- df %>% dplyr::filter(ambig == 1)
        ambig_save <- ambig_save %>%
            dplyr::filter(!is.na(.data[[ambig_key]])) %>%
            group_by(region_id) %>%
            dplyr::summarise(ambig_list = paste(unique(.data[[ambig_key]]), collapse = ",")) %>%
            ungroup()
    }

    # if we use some column to sort the annotation for resolution of ambiguities
    if (!is.null(order_key) & !is.null(order_vec)) {
        # Order by the overlap amount first, then by the order key (gene_type)
        message("Sorting the annotation layers by: ", order_key)
        if ((order_key %in% colnames(df) && (order_key %in% keys))) {
            df <- df %>%
                mutate(sort_col = factor(.data[[order_key]], levels = order_vec, ordered = TRUE)) %>%
                arrange(sort_col) %>%
                select(-sort_col) # sort by provided order
        } else {
            warning(
                "Provided annotaion order key", order_key,
                "does not exist in annotation or not part of the keys, skipping"
            )
        }
    }

    df <- df %>%
        tidyr::fill(all_of(keys), .direction = "up") %>%
        slice(1) %>%
        ungroup()

    dt <- left_join(as_tibble(mcols(gi@regions)), df, by = "region_id")

    if (save_ambig) {
        dt <- left_join(dt, ambig_save, by = "region_id")
    }

    dt <- dt %>%
        select(-c(feature_id, overlap_w, region_id))

    cnames <- c(str_c(colnames(dt), ".A"), str_c(colnames(dt), ".B"))

    dt <- cbind(dt[gi@anchor1, ], dt[gi@anchor2, ])
    colnames(dt) <- cnames
    dt <- dt %>% dplyr::select(order(colnames(dt)))

    if (save_ambig) {
        dt <- dt %>%
            relocate(ambig.A, ambig.B, ambig_list.A, ambig_list.B, .after = last_col())
    } else {
        dt <- dt %>% dplyr::select(-c(ambig.A, ambig.B, ambig_list.A, ambig_list.B))
    }

    dt <- dt %>% data.frame()
    mcols(gi) <- cbind(mcols(gi)[orininal_columns], dt)
    gi@regions$region_id <- NULL
    return(gi)
}



#' Annotate RNA-RNA interactions as cis- and trans-
#'
#' @description
#' Annotated each entry gi object as cis, if the .A and .B arms correspond to
#' the same feature (i.e transcript_id or gene_id)
#' If the values are are equal, then annotated with value: `cis` = 1, If not equal
#' or `NA`: `cis` = 0
#'
#' @param gi `GInteractions` object containing two metadata columns as feature annotation
#' @param id_col_base base name of the feature id columns to use. Function will
#' look for <id_col_base>.A and <id_col_base>.B columns and compare them
#' @keywords internal
#' @return gi `GInteractions` object containing `cis` field with 0/1 values
.annotateCisTrans <- function(gi, id_col_base = "gene_id") {
    gi$cis <- 0

    name_col1 <- paste0(id_col_base, ".A")
    name_col2 <- paste0(id_col_base, ".B")

    gi$c1 <- mcols(gi)[, name_col1]
    gi$c2 <- mcols(gi)[, name_col2]


    if (length(gi[!is.na(gi$c1) & !is.na(gi$c2) &
        (gi$c1 == gi$c2)]) != 0) {
        gi[!is.na(gi$c1) & !is.na(gi$c2) &
            (gi$c1 == gi$c2)]$cis <- 1
        gi$c1 <- NULL
        gi$c2 <- NULL
        return(gi)
    } else {
        message("Could not do cis/trans annotation because columns content is missing ")
        return(gi)
    }
}

#' Helper function to add count data to metadata of `GInteractions`
#'
#' Merges the count dataframe and interactions metadata by `id_col`
#' If key is not found, in metadata throws error

#' @param gi `GInteractions`
#' @param df_counts dataframe with read counts
#' @param id_col key to use in merge
#' @return `GInteractions` with added counts
.addGeneCounts <- function(gi, df_counts, id_col = "gene_id") {
    df_counts <- df_counts %>% as.data.frame()
    df_all_cts <- as_tibble(data.frame("RNA" = unname(df_counts[1]), "n" = unname(df_counts[2])))

    id_columns <- str_c(id_col, c(".A", ".B"))

    if (!all(id_columns %in% colnames(mcols(gi)))) {
        stop(id_columns, " are not found in the input gi")
    }
    colnames(df_all_cts) <- c(id_columns[1], "gene_count.A")
    gi$gene_count.A <- left_join(
        as_tibble(mcols(gi)[id_columns[1]]),
        df_all_cts
    ) %>% pull(gene_count.A)
    colnames(df_all_cts) <- c(id_columns[2], "gene_count.B")
    gi$gene_count.B <- left_join(
        as_tibble(mcols(gi)[id_columns[2]]),
        df_all_cts
    ) %>% pull(gene_count.B)
    return(gi)
}
#' Calculate ligation-based significance statistics for RNA duplexes
#'
#' Calculates significance statistics for each gene/transcript pair represented
#' in a `GInteractions` object. Two optional models are supported:
#'
#' 1. A binomial random-ligation model, where pair probability is derived from
#'    per-gene/transcript counts and tested using the right-tail binomial test.
#' 2. An odds-ratio / Fisher exact test based on a 2x2 contingency table using
#'    chimeric counts and singleton counts. This test is used in RBP-dependent and
#'    RNA probing methods like RIL-seq
#'
#' The function also outputs abundance fractions for each duplex pair and counts
#' of other RNA duplexes formed by either or neither gene/transcript in the pair.
#'
#' @param gi `GInteractions` object annotated with gene/transcript names.
#' @param df_counts `data.frame`. A two-column data frame with gene/transcript
#'   counts. The first column should match the `<id_col>.A` / `<id_col>.B`
#'   identifiers in `gi`. The second column is the corresponding count.
#' @param id_col `character(1)`. Prefix for gene/transcript metadata fields in
#'   `gi`. Two fields, `<id_col>.A` and `<id_col>.B`, are expected.
#' @param compute_binom `logical(1)`. Whether to compute the binomial
#'   random-ligation test. Default is `TRUE`.
#' @param compute_odds `logical(1)`. Whether to compute the odds-ratio /
#'   Fisher exact test. Default is `TRUE`.
#'
#' @return `GInteractions` object with added metadata columns.
#'
#' @details
#' **Binomial model**
#'
#' Null hypothesis: an observed duplex is generated by random ligation.
#'
#' Alternative hypothesis: the duplex is enriched above the random-ligation
#' expectation.
#'
#' The random-ligation probability \eqn{P(a, b)} is defined as:
#'
#' \deqn{
#' P(a, b) \propto
#' \begin{cases}
#'     2 \cdot P(a) \cdot P(b) & \text{if } a:b \text{ is observed and } a \neq b \\
#'     P(a)^2                  & \text{if } a:b \text{ is observed and } a = b \\
#'     0                       & \text{otherwise}
#' \end{cases}
#' }
#'
#' where:
#'
#' \deqn{
#' P(a) = \frac{\text{N reads}(a)}{\text{total N reads}}
#' }
#'
#' The observed-pair probabilities are renormalized to sum to one and the
#' p-value is computed from the right tail of the binomial distribution.
#'
#' **Odds-ratio / Fisher model**
#'
#' For a pair `(A, B)`, the contingency table is:
#'
#' \tabular{lll}{
#'                 \tab B present \tab B absent\cr
#'   A present     \tab C1        \tab C2\cr
#'   A absent      \tab C3        \tab C4\cr
#' }
#'
#' with:
#' \itemize{
#'   \item `C1`: chimeric fragments containing both `A` and `B`
#'   \item `C2`: fragments containing `A` but not `B`
#'   \item `C3`: fragments containing `B` but not `A`
#'   \item `C4`: fragments containing neither `A` nor `B`
#' }
#'
#' A one-sided Fisher exact test (`alternative = "greater"`) is used.
#'
#' @examples
#' data("RNADuplexesSampleData")
#' gi <- calculateLigationPvalues(
#'     RNADuplexSampleDGs,
#'     df_counts = RNADuplexesGeneCounts
#' )
#' hist(gi$p.adj, breaks = 20)
#'
#' @export
calculateLigationPvalues <- function(gi,
    df_counts,
    id_col = "gene_id",
    compute_binom = TRUE,
    compute_odds = FALSE) {
    df_counts <- df_counts %>% as.data.frame()
    if (!(compute_binom || compute_odds)) {
        stop("Error in calculateLigationPvalues at least one compute_binom or compute_odds must be TRUE")
    }


    df_all_cts <- as_tibble(data.frame(
        RNA = unname(df_counts[[1]]),
        n = unname(df_counts[[2]])
    ))

    id_columns <- str_c(id_col, c(".A", ".B"))

    if (!all(id_columns %in% colnames(mcols(gi)))) {
        stop(id_columns, " are not found in the input gi")
    }

    chim_df <- as_tibble(mcols(gi)) %>%
        dplyr::select(c(n_reads, !!!id_columns))
    colnames(chim_df) <- c("n", "Araw", "Braw")
    chim_df$chim_id <- seq_len(nrow(chim_df))
    # filter not annotated chimeras and those for which we don't have record in counts df
    unique_ids <- unique(df_all_cts$RNA)


    # keep original row-to-pair mapping so we can join back later
    chim_df$AB <- apply(chim_df[c("Araw", "Braw")], 1, function(row) {
        paste(sort(row), collapse = "<>")
    })
    chimdf_save <- chim_df[c("chim_id", "AB")]
    chim_df <- chim_df %>%
        group_by(AB) %>%
        summarise(n_reads_chim_total = sum(n), .groups = "drop")

    ssplit <- str_split_fixed(chim_df$AB, "<>", n = 2)
    chim_df$A <- ssplit[, 1]
    chim_df$B <- ssplit[, 2]

    chim_df <- chim_df %>%
        left_join(
            tibble(idA = df_all_cts$RNA, nA = df_all_cts$n),
            by = c("A" = "idA")
        ) %>%
        left_join(
            tibble(idB = df_all_cts$RNA, nB = df_all_cts$n),
            by = c("B" = "idB")
        ) %>%
        dplyr::filter((A %in% unique_ids) & (B %in% unique_ids))

    N_chim_total <- sum(chim_df$n_reads_chim_total)
    N_total <- sum(df_all_cts$n)
    N_frag_total <- N_total + N_chim_total

    # pair-level abundance
    chim_df <- chim_df %>%
        mutate(
            Pa = nA / N_total,
            Pb = nB / N_total,
            Pab = ifelse(A != B, 2 * Pa * Pb, Pa * Pa)
        )
    df <- data.frame(chim_df[c("A", "B", "n_reads_chim_total")]) %>%
        dplyr::rename(n = "n_reads_chim_total")

    # For odds- test count chimeras with A but not B, with B but not A, and with neither
    chim_df$chim_withA_noB <- apply(df, 1, function(row) {
        df %>%
            dplyr::filter(
                (A == row["A"] | B == row["A"]) &
                    (A != row["B"] & B != row["B"])
            ) %>%
            pull(n) %>%
            sum()
    })

    chim_df$chim_withB_noA <- apply(df, 1, function(row) {
        df %>%
            dplyr::filter(
                (A == row["B"] | B == row["B"]) &
                    (A != row["A"] & B != row["A"])
            ) %>%
            pull(n) %>%
            sum()
    })

    chim_df$chim_noB_noA <- apply(df, 1, function(row) {
        df %>%
            dplyr::filter(
                (A != row["B"] & B != row["B"]) &
                    (A != row["A"] & B != row["A"])
            ) %>%
            pull(n) %>%
            sum()
    })

    # singleton fpor odds
    chim_df$nonchimA <- vapply(
        vctrs::vec_cast(chim_df$A, character()),
        function(x) {
            df_all_cts %>%
                dplyr::filter(RNA == x) %>%
                pull(n)
        },
        numeric(1)
    )

    chim_df$nonchimB <- vapply(
        vctrs::vec_cast(chim_df$B, character()),
        function(x) {
            df_all_cts %>%
                dplyr::filter(RNA == x) %>%
                pull(n)
        },
        numeric(1)
    )

    # unify output
    chim_df$pval <- NA_real_
    chim_df$p.adj <- NA_real_
    chim_df$odds_ratio <- NA_real_
    chim_df$odds_pval <- NA_real_
    chim_df$odds_p.adj <- NA_real_

    if (compute_binom) {
        chim_df$Pab_norm <- chim_df$Pab / sum(chim_df$Pab)

        chim_df$pval <- pbinom(
            chim_df$n_reads_chim_total - 1,
            size = N_chim_total,
            prob = chim_df$Pab_norm,
            lower.tail = FALSE
        )
        chim_df$p.adj <- p.adjust(chim_df$pval, method = "BH")
    } else {
        chim_df$Pab_norm <- NA_real_
    }

    if (compute_odds) {
        # 2x2 table - odds test similar to the one in bacterial methods :
        #            B present                  B absent
        # A present  C1 = AB chimera           C2 = A only
        # A absent   C3 = B only               C4 = neither A nor B
        chim_df$C1 <- chim_df$n_reads_chim_total
        chim_df$C2 <- chim_df$chim_withA_noB + chim_df$nonchimA
        chim_df$C3 <- chim_df$chim_withB_noA + chim_df$nonchimB
        chim_df$C4 <- N_frag_total - chim_df$C1 - chim_df$C2 - chim_df$C3

        odds_res <- lapply(seq_len(nrow(chim_df)), function(i) {
            mat <- matrix(
                c(
                    chim_df$C1[i], chim_df$C2[i],
                    chim_df$C3[i], chim_df$C4[i]
                ),
                nrow = 2,
                byrow = TRUE
            )

            ft <- stats::fisher.test(mat, alternative = "greater")

            data.frame(
                odds_ratio = unname(ft$estimate),
                odds_pval = ft$p.value
            )
        })

        odds_res <- dplyr::bind_rows(odds_res)
        chim_df$odds_ratio <- odds_res$odds_ratio
        chim_df$odds_pval <- odds_res$odds_pval
        chim_df$odds_p.adj <- p.adjust(chim_df$odds_pval, method = "BH")
    } else {
        chim_df$C1 <- NA_real_
        chim_df$C2 <- NA_real_
        chim_df$C3 <- NA_real_
        chim_df$C4 <- NA_real_
    }

    counts_stats <- chim_df %>%
        dplyr::select(
            AB,
            pval, p.adj, Pa, Pb,
            chim_withA_noB, chim_withB_noA, chim_noB_noA,
            nonchimA, nonchimB,
            C1, C2, C3, C4,
            odds_ratio, odds_pval, odds_p.adj
        ) %>%
        dplyr::rename(
            p_val = pval,
            count.chim.A.notB = chim_withA_noB,
            count.chim.notA.B = chim_withB_noA,
            count.chim.notA.notB = chim_noB_noA
        )

    chimdf_save <- chimdf_save %>%
        left_join(counts_stats, by = "AB")

    statcounts_final <- tibble(chim_id = seq_len(length(gi))) %>%
        left_join(chimdf_save, by = "chim_id") %>%
        dplyr::select(-c(chim_id, AB)) %>%
        data.frame()

    mcols_old <- mcols(gi) %>%
        data.frame() %>%
        as_tibble() %>%
        dplyr::select(-any_of(colnames(statcounts_final))) %>%
        data.frame()

    mcols(gi) <- cbind(mcols_old, statcounts_final)

    gi <- .addGeneCounts(gi, df_counts)

    return(gi)
}

.checkRNAduplexinstalled <- function() {
    returncode <- system2("RNAduplex", args = c("-h"), stdout = NULL, stderr = NULL)
    if (returncode != 0) {
        return(1)
    } else {
        return(0)
    }
}

.runRNAduplex <- function(RNA1, RNA2) {
    input <- paste(RNA1, RNA2, sep = "\n")
    result <- system2("RNAduplex", input = input, stdout = TRUE)
    return(result)
}

.getDuplexString <- function(gi, fafile) {
    s <- Biostrings::readBStringSet(fafile)
    s <- Biostrings::RNAStringSet(s)
    newnames <- str_squish(str_sub(names(s), 1, 5))
    names(s) <- newnames
    gi$seq1 <- s[get_arm_a(gi)]
    gi$seq2 <- s[get_arm_b(gi)]
    return(gi)
}


.getGCContent <- function(seqrna) {
    # seqrna is a RNAstringset based on the arm
    gc_freq <- Biostrings::letterFrequency(seqrna, c("G", "C"))
    gc_content <- rowSums(gc_freq) / width(seqrna) * 100
    return(gc_content)
}
