# glmnet function part a
cv.glmnet.balanced = function(x, y, nfolds = 10, ...) {
  
  foldid = rep(NA_integer_, length(y))
  sink("/dev/null")
  ll = pamr:::balanced.folds(y, nfolds) 
  sink()
  for (i in seq_along(ll)) { foldid[ ll[[i]] ] = i }
  stopifnot(!any(is.na(foldid)), nrow(x) == length(y))
  
  weights = rep(NA_real_, length(y))
  tab = table(y)
  stopifnot(setequal(names(tab), levels(y)))
  for (nm in names(tab)) weights[ y == nm ] = 1/tab[nm]

  cv.glmnet(x = x, y = y, foldid = foldid, weights = weights, ...) 
}

entities = c("rLN", "MZL", "FL", "MCL", "DLBCL") 

# glmnet function part b
my_glmnet = function(df) {
  require("glmnet")
  require("pamr")
  x = dplyr::select(df, all_of(cell_types)) |> as.matrix()
  y = factor(df$Entity, levels = entities)
  
  ## estimate prediction performance by LOO CV
  confusion_table = lapply(seq_along(y), function(i) {
    pr = cv.glmnet.balanced(x[-i, ], y[-i], family = "multinomial") |>
      predict(newx = x[i,, drop = FALSE], type = "response")
    tibble(truth   = y[i],
           predicted = factor(names(which.max(pr[1,,1])), levels = entities))
  }) |> bind_rows() |> table()
  
  ## final fit
  fit  = cv.glmnet.balanced(x, y, family = "multinomial")
  cfit = coef(fit)
  coefs = lapply(names(cfit), function(ent) { 
    m = cfit[[ent]]
    tibble(Entity    = ent,
           cell_type = rownames(m), 
           beta      =  m[, 1]) 
  }) |> bind_rows() |> mutate(Entity = factor(Entity, levels = entities))
  
  list(fit = fit, 
       confusion_table = confusion_table,
       coefs = coefs)
}

# This function reads and handles TCR data
readTCR <- function(files=NULL){
  
  lapply(files, data.table::fread) %>% 
    bind_rows() %>% 
    as_tibble() %>% 
    mutate(PatientID=strsplit(barcode, split = "_") %>% sapply("[[", 2)) %>% 
    select(PatientID, 1:(ncol(.)-1)) %>% 
    rename(Barcode_full=barcode)

} 

# This function adds entity to df
add_entity <- 
  function(data){
      data %>% 
        left_join(., df_meta %>% select(PatientID, Entity) %>% distinct, by="PatientID") %>% 
        dplyr::mutate(Entity=gsub(Entity, pattern = ", GCB|, non-GCB", replacement = ""))
  }


# This function calculates the proportion of variables
add_prop <- 
  function(data=NULL, vars=NULL, group.vars=NULL, ungroup=TRUE, keep.n=FALSE, prop.name=NULL) {
    
    group.vars <-  vars[group.vars]
    
    dftmp <- 
      data %>% 
        dplyr::count(!!!syms(unique(c(vars, group.vars)))) %>% 
        group_by(!!!syms(c(group.vars))) %>% 
        dplyr::mutate(Prop=n/sum(n))
    
    if(keep.n==FALSE){
      dftmp <- dftmp %>% select(-n)
    }
    
    if(ungroup==TRUE){
      dftmp <- dftmp %>% ungroup()
    }
    
    if(!is.null(prop.name)){
      colnames(dftmp)[which(colnames(dftmp)=="Prop")] <- prop.name
      
    }
    
    return(dftmp)
  }

# This functions expands data frames and fills missing values with 0
fill_zeros <- 
  function(data=NULL ,names_from=NULL, values_from=NULL) {
    data %>% 
      pivot_wider(names_from = all_of(names_from), values_from = all_of(values_from), values_fill = 0) %>% 
      pivot_longer(cols = (ncol(data)-1):ncol(.), names_to = names_from, values_to = values_from)
  }
  
# Run standard Seurat processing pipeline
SeuratProc_T <- 
  function(sobj, verbose=FALSE, dims.clustering=NULL, resolution.clustering=NULL, dims.umap=NULL) {
    
    # Remove 
    sobj <- DietSeurat(sobj)
    DefaultAssay(sobj) <- "RNA"
    
    # Filter data set based on RNA
    sobj <- FindVariableFeatures(sobj, selection.method = "vst", nfeatures = 2000, verbose=verbose)
    
    # Scale data (RNA and ADT)
    sobj <- ScaleData(sobj, features = rownames(sobj), verbose=verbose)
    
    # Assess cell cycle
    sobj <- CellCycleScoring(sobj, s.features = cc.genes$s.genes, g2m.features = cc.genes$g2m.genes, set.ident = TRUE)
    sobj <- ScaleData(sobj, vars.to.regress = c("S.Score", "G2M.Score", "percent.mt"), verbose=verbose)
    
    # Run PCA
    sobj <- RunPCA(sobj, features = VariableFeatures(sobj), nfeatures.print=5, ndims.print=1:2,
                   reduction.name = "pcaRNA", reduction.key = "pcaRNA_")
    
    # Run clustering based on transcriptome
    sobj <- FindNeighbors(sobj, dims = dims.clustering, verbose=verbose, reduction = "pcaRNA")
    sobj <- FindClusters(sobj, resolution = resolution.clustering, verbose=verbose)
    
    # Run UMAP based on transcriptome
    sobj <- RunUMAP(sobj, dims = dims.umap, verbose=verbose, reduction.key = "umapRNA_",
                    reduction.name = "umapRNA", reduction = "pcaRNA")
    
    return(sobj)
    
  }

SeuratProcADT_T <- 
  function(sobj, verbose=FALSE, dims.clustering=NULL, resolution.clustering=NULL, dims.umap=NULL) {
    
    DefaultAssay(sobj) <- "ADT"
    VariableFeatures(sobj, assay="ADT") <- rownames(sobj@assays$ADT)
    
    sobj <- ScaleData(sobj, assay = "ADT", verbose=verbose)
    
    #### Run PCA and print ElbowPlot
    sobj <- RunPCA(sobj, features = rownames(sobj@assays$ADT), nfeatures.print=5, ndims.print=1:2, 
                   reduction.name = "pcaADT", reduction.key = "pcaADT_")
    
    #### Run clustering based on ADT
    sobj <- FindNeighbors(sobj, dims = dims.clustering, verbose=verbose,  reduction = "pcaADT")
    sobj <- FindClusters(sobj, resolution = resolution.clustering, verbose=verbose)
    
    #### Run UMAP based on ADT
    sobj <- RunUMAP(sobj, dims = dims.umap, verbose=verbose, reduction = "pcaADT",
                    reduction.name = "umapADT", reduction.key = "umapADT_")
    
    DefaultAssay(sobj) <- "RNA"
    
    return(sobj)
    
  }