#########################################################
# TASK: To compare OR from master snps
# chisq, fisher test and logistic and adjusted logistic
#########################################################
getwd()
setwd('~/git/LSHTM_analysis/scripts')
getwd()

#install.packages("logistf")
library(logistf)
#########################################################
#%% variable assignment: input and output paths & filenames
drug = 'pyrazinamide'
gene = 'pncA'
gene_match = paste0(gene,'_p.')
cat(gene_match)

#===========
# input and output dirs
#===========
datadir = paste0('~/git/Data')
indir = paste0(datadir, '/', drug, '/', 'input')
outdir = paste0(datadir, '/', drug, '/', 'output')

#===========
# input and output files
#===========
in_filename  = 'original_tanushree_data_v2.csv'
#in_filename  = 'mtb_gwas_v3.csv'
infile = paste0(datadir, '/', in_filename)
cat(paste0('Reading infile1: raw snps', ' ', infile) )

# infile2: _gene associated meta snps file to extract valid snps and add calcs to.
# This is outfile3 from snps_extraction.py
in_filename_metadata = paste0(tolower(gene), '_metadata.csv')
infile_metadata = paste0(outdir, '/', in_filename_metadata)
cat(paste0('Reading infile2: gene associated metadata:', infile_metadata))

#===========
# output
#===========
out_filename = paste0(tolower(gene),'_', 'af_or.csv')
outfile = paste0(outdir, '/', out_filename)
cat(paste0('Output file with full path:', outfile))
#%% end of variable assignment for input and output files

#########################################################
# 1: Read master/raw snps stored in snps/
#####################################################

#===============
# Step 1: read raw snps (all remove entries with NA in pza column)
#===============
raw_data_all = read.csv(infile, stringsAsFactors = F)

# building cols to extract
dr_muts_col = paste0('dr_mutations_', drug)
other_muts_col = paste0('other_mutations_', drug)

cat('Extracting columns based on variables:\n'
    , drug
    , '\n'
    , dr_muts_col
    , '\n'
    , other_muts_col
    , '\n===============================================================')

raw_data = raw_data_all[,c("id"
                           , drug
                           , dr_muts_col
                           , other_muts_col)]
rm(raw_data_all)

rm(indir, in_filename, infile)

#===========
# 1a: exclude na
#===========
raw_data = raw_data[!is.na(raw_data[[drug]]),]

total_samples = length(unique(raw_data$id))
cat(paste0('Total samples without NA in', ' ', drug, 'is:', total_samples))

# sanity check: should  be true
is.numeric(total_samples) 

#===========
# 1b: combine the two mutation columns
#===========
#raw_data$all_mutations_pyrazinamide = paste(raw_data$dr_mutations_pyrazinamide, raw_data$other_mutations_pyrazinamide)

all_muts_colname = paste0('all_mutations_', drug)
raw_data[[all_muts_colname]] = paste(raw_data[[dr_muts_col]], raw_data[[other_muts_col]])
head(raw_data[[all_muts_colname]])

#===========
# 1c: create yet another column that contains all the mutations but in lower case
#===========
head(raw_data[[all_muts_colname]])
raw_data$all_muts_gene = tolower(raw_data[[all_muts_colname]]) 
head(raw_data$all_muts_gene)

# sanity checks
#table(grepl("gene_p",raw_data$all_muts_gene))
cat(paste0('converting gene match:', gene_match, ' ', 'to lowercase'))
gene_match = tolower(gene_match)

table(grepl(gene_match,raw_data$all_muts_gene))

# sanity check
if(sum(table(grepl(gene_match, raw_data$all_muts_gene))) == total_samples){
  cat('PASS: Total no. of samples match')
} else{
  cat('FAIL: No. of samples mismatch')
}

#########################################################
# 2: Read valid snps for which OR 
# can be calculated
#########################################################
cat(paste0('Reading metadata infile:', infile_metadata))

gene_metadata = read.csv(infile_metadata
                         #, file.choose()
                         , stringsAsFactors = F
                         , header = T)


# clear variables
rm(in_filename_metadata, infile_metadata)

# count na in pyrazinamide column
tot_pza_na = sum(is.na(gene_metadata$pyrazinamide))
expected_rows = nrow(gene_metadata) - tot_pza_na

# drop na from the pyrazinamide colum
gene_snps_or = gene_metadata[!is.na(gene_metadata[[drug]]),]

# sanity check
if(nrow(gene_snps_or) == expected_rows){
  cat('PASS: no. of rows match with expected_rows')
} else{
  cat('FAIL: nrows mismatch.')
}

# extract unique snps to iterate over for AF and OR calcs
gene_snps_unique = unique(gene_snps_or$mutation) 

cat(paste0('Total no. of distinct comp snps to perform OR calcs: ', length(gene_snps_unique)))
#=====================================
#OR calcs using the following 4
#1) chisq.test
#2) fisher
#3) modified chisq.test
#4) logistic
#5) adjusted logistic?
#6) kinship (separate script)

#======================================

################# modified chisq OR
# Define OR function
#x = as.numeric(mut)
#y = dst
custom_chisq_or = function(x,y){
  tab = as.matrix(table(x,y))
  a = tab[2,2]
  if (a==0){ a<-0.5}
  b = tab[2,1]
  if (b==0){ b<-0.5}
  c = tab[1,2]
  if (c==0){ c<-0.5}
  d = tab[1,1]
  if (d==0){ d<-0.5}
  (a/b)/(c/d)
  
}

#========================
# TEST WITH ONE
#========================
i = "pnca_p.trp68gly"
i = "pnca_p.gln10pro"
i = "pnca_p.leu159arg"

# IV
table(grepl(i,raw_data$all_muts_gene))
mut = grepl(i,raw_data$all_muts_gene)

# DV
#dst = raw_data$pyrazinamide
dst = raw_data[[drug]] # or raw_data[,drug]

# 2X2 table
table(mut, dst)

# CV
#c = raw_data$id[mut]
c = raw_data$id[grepl(i,raw_data$all_muts_gene)] 
#sid = grepl(raw_data$id[mut], raw_data$id) # warning
#argument 'pattern' has length > 1 and only the first element will be used
#grepl(raw_data$id=="ERR2512440", raw_data$id)

sid = grepl(paste(c,collapse="|"), raw_data$id)
table(sid)   

# 3X2 table
table(mut, dst, sid)

#============================
# compare OR
chisq.test(table(mut,dst))
chisq.test(table(mut,dst)) $ statistic

f = chisq.test(table(mut,dst)) $ statistic
chisq.test(dst, mut) $ statistic

fisher.test(table(mut, dst))
fisher.test(table(mut, dst))$p.value
fisher.test(table(mut, dst))$estimate
logistic_chisq_or(mut,dst)

# logistic or
summary(model<-glm(dst ~ mut, family = binomial))
or_logistic = exp(summary(model)$coefficients[2,1]); print(or_logistic)
pval_logistic_maxit = summary(model)$coefficients[2,4]; print(pval_logistic_maxit)

# extract SE of the logistic model for a given snp
logistic_se = summary(model)$coefficients[2,2]
print(paste0('SE:', logistic_se))
    
# extract Z of the logistic model for a given snp
logistic_zval = summary(model)$coefficients[2,3]
print(paste0('Z-value:', logistic_zval))
  
  
# extract confint interval of snp (2 steps, since the output is a named number)
ci_mod = exp(confint(model))[2,]
print(paste0('CI:', ci_mod))
#logistic_ci = paste(ci_mod[["2.5 %"]], ",", ci_mod[["97.5 %"]])
  
logistic_ci_lower = ci_mod[["2.5 %"]]
logistic_ci_upper = ci_mod[["97.5 %"]]
  
print(paste0('CI_lower:', logistic_ci_lower))
print(paste0('CI_upper:', logistic_ci_upper))


# adjusted logistic or: doesn't seem to make a difference
summary(model2<-glm(dst ~ mut + sid, family = binomial))
or_logistic2 = exp(summary(model2)$coefficients[2,1]); print(or_logistic2)
pval_logistic2 = summary(model2)$coefficients[2,4]; print(pval_logistic2)

#============ looping with sapply
#####################
# iterate: subset
#####################

snps_test = c("pnca_p.trp68gly", "pnca_p.leu4ser", "pnca_p.leu159arg","pnca_p.his57arg" )

snps = snps_test[1:4]

snps


ors = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  logistic_chisq_or(mut,dst)
})

ors

pvals = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  fisher.test(mut,dst)$p.value
})

pvals

afs = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  mean(mut)
})

afs

## logistic or
 ors_logistic = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut, family = binomial)
  or_logistic = exp(summary(model)$coefficients[2,1])

})
ors_logistic
head(ors_logistic); head(names(ors_logistic))

## logistic p-value
pvals_logistic = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut , family = binomial)
  pval_logistic = summary(model)$coefficients[2,4]
})

head(pvals_logistic); head(names(pvals_logistic))

## logistic se
se_logistic = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut , family = binomial)
  logistic_se = summary(model)$coefficients[2,2]
})

head(se_logistic); head(names(se_logistic))


## logistic z-value
zval_logistic = sapply(gene_snps_unique,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut , family = binomial)
  logistic_zval = summary(model)$coefficients[2,3]
})

head(zval_logistic); head(names(zval_logistic))

## logistic ci - lower bound
ci_lb_logistic = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut , family = binomial)
  ci_mod = exp(confint(model))[2,]
  logistic_ci_lower = ci_mod[["2.5 %"]]
})

head(ci_lb_logistic); head(names(ci_lb_logistic))

## logistic ci - upper bound
ci_ub_logistic = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut , family = binomial)
  ci_mod = exp(confint(model))[2,]
  logistic_ci_upper = ci_mod[["97.5 %"]]

})

head(ci_ub_logistic); head(names(ci_ub_logistic))

# logistic adj # Doesn't seem to make a difference
logistic_ors2 = sapply(snps,function(m){
  mut = grepl(m,raw_data$all_muts_gene)
  c = raw_data$id[mut]
  sid = grepl(paste(c,collapse="|"), raw_data$id)
  model2<-glm(dst ~ mut + sid, family = binomial)
  or_logistic2 = exp(summary(model2)$coefficients[2,1])
  #pval_logistic2 = summary(model2)$coefficients[2,4]
})

logistic_ors2

or_logistic2; pval_logistic2

head(logistic_ors)
#===========================================================
#%%
# sapply with multiple values

#https://gist.github.com/primaryobjects/33adabc337edd67b4a8d

snps_test = c("pnca_p.trp68gly", "pnca_p.leu4ser", "pnca_p.leu159arg","pnca_p.his57arg" )

snps = snps_test[1:4]

snps

# DV: pyrazinamide 0 or 1
dst = raw_data[[drug]]

# yayy works!
testdf = data.frame()

x = sapply(snps,function(m){

  df = data.frame()
  
  mut = grepl(m,raw_data$all_muts_gene)
  model<-glm(dst ~ mut, family = binomial)
  
  
  # allele frequency
  afs = mean(mut)
  
  # logistic model
  beta_logistic = summary(model)$coefficients[2,1]

  or_logistic = exp(summary(model)$coefficients[2,1])
  print(paste0('logistic OR:', or_logistic))
  
  pval_logistic = summary(model)$coefficients[2,4]
  print(paste0('logistic pval:', pval_logistic))
  
  se_logistic = summary(model)$coefficients[2,2]
  zval_logistic = summary(model)$coefficients[2,3]
  ci_mod = exp(confint(model))[2,]
  ci_lower_logistic = ci_mod[["2.5 %"]]
  ci_upper_logistic = ci_mod[["97.5 %"]]
  
  # custom_chisq and fisher: OR p-value and CI
  or_mychisq = custom_chisq_or(dst, mut)
  
  or_fisher = fisher.test(dst, mut)$estimate
  or_fisher = or_fisher[[1]]
  
  pval_fisher = fisher.test(dst, mut)$p.value
  
  ci_lower_fisher = fisher.test(dst, mut)$conf.int[1]
  ci_upper_fisher = fisher.test(dst, mut)$conf.int[2]
  
  # chi sq estimates 
  estimate_chisq = chisq.test(dst, mut)$statistic; estimate_chisq
  est_chisq = estimate_chisq[[1]]; print(est_chisq)
  
  pval_chisq = chisq.test(dst, mut)$p.value                                                        
                                                         
  #build a row to append to df
  row = data.frame(mutation = m
                   , af = afs
                   , beta_logistic = beta_logistic
                   , or_logistic = or_logistic
                   , pval_logistic = pval_logistic
                   , se_logistic = se_logistic
                   , zval_logistic = zval_logistic
                   , ci_low_logistic = ci_lower_logistic
                   , ci_hi_logistic = ci_upper_logistic
                   , or_mychisq = or_mychisq 
                   , or_fisher = or_fisher
                   , pval_fisher = pval_fisher
                   , ci_low_fisher= ci_lower_fisher
                   , ci_hi_fisher = ci_upper_fisher
                   , est_chisq = est_chisq
                   , pval_chisq = pval_chisq
                   )
  #print(row)
  
  testdf <<- rbind(testdf, row)
  
})

write.csv(testdf, 'test_ors.csv')
#=================================
####################
# iterate: subset
#####################
print(paste0('subset to iterate over;', snps))

# start loop
 perfectSeparation <- function(w) {
  if(grepl("fitted probabilities numerically 0 or 1 occurred",
           as.character(w))) {ww <<- ww+1}
}

for (i in snps){
  
  print(i)
  
  # IV
  #mut<-as.numeric(grepl(i,raw_data$all_muts_gene))
  mut = grepl(i,raw_data$all_muts_gene)
  table(mut)
  
  # DV
  #dst<-as.numeric(raw_data[[drug]])
  dst = raw_data[[drug]]
  
  # table
  print(table(dst, mut))
  
  #=====================
  # logistic regression, glm.control(maxit = n)
  #https://stats.stackexchange.com/questions/11109/how-to-deal-with-perfect-separation-in-logistic-regression
  #=====================
  #n = 1
  summary(model<-glm(dst ~ mut
                     , family = binomial
                     #, control = glm.control(maxit = 1)
                     #, options(warn = 1)
                     ))
                     #, warning = perfectSeparation))
  or_logistic = exp(summary(model)$coefficients[2,1]); print(or_logistic)
  pval_logistic_maxit = summary(model)$coefficients[2,4]; print(pval_logistic_maxit)
  logistic_se = summary(model)$coefficients[2,2]
  logistic_zval = summary(model)$coefficients[2,3]
  ci_mod = exp(confint(model))[2,]
  logistic_ci_lower = ci_mod[["2.5 %"]]
  logistic_ci_upper = ci_mod[["97.5 %"]]
  
  #=====================
  # fishers test
  #=====================
  #attributes(fisher.test(table(dst, mut)))
  or_fisher = fisher.test(table(dst, mut))$estimate
  or_fisher = or_fisher[[1]]; or_fisher
  
  pval_fisher = fisher.test(table(dst, mut))$p.value ; pval_fisher
  
  #=====================
  # chi square
  #=====================
  #chisq.test(y = dst, x = mut)
  #attributes(chisq.test(table(dst, mut)))
  est_chisq = chisq.test(table(dst, mut))$statistic 
  est_chisq = est_chisq[[1]]; est_chisq
  
  pval_chisq = chisq.test(table(dst, mut))$p.value; pval_chisq
  
  # all output 
  writeLines(c(paste0("mutation:", i)
               , paste0("=========================")
               , paste0("or_logistic:", or_logistic,"--->", "P-val_logistic_maxit:", pval_logistic_maxit )
               , paste0("OR_fisher:", or_fisher, "--->","P-val_fisher:", pval_fisher )
               , paste0("Chi_sq_estimate:", est_chisq, "--->","P-val_chisq:", pval_chisq)))
}
  #=====================
  # fishers test
  #=====================
  #attributes(fisher.test(table(dst, mut)))
  or_fisher = fisher.test(table(dst, mut))$estimate
  or_fisher = or_fisher[[1]]; or_fisher
  
  pval_fisher = fisher.test(table(dst, mut))$p.value ; pval_fisher
  
  
 # https://stats.stackexchange.com/questions/259635/what-is-the-difference-using-a-fishers-exact-test-vs-a-logistic-regression-for
  exact2x2(table(dst, mut),tsmethod="central")


#=====================================================================
# iterate over a df and then add these values  
#
my_data = as.data.frame(gene_snps_unique)
colnames(my_data) = "mutation"
print(colnames(my_data))

perfectSeparation <- function(w) {
  if(grepl("fitted probabilities numerically 0 or 1 occurred",
           as.character(w))) {ww <<- ww+1}
}

for(i in my_data$mutation) {
  print(paste0('snp to iterate over:', i))
}

for(i in my_data$mutation) {
  print(paste0('snp to iterate over:', i))

  #####
  # Run logistic regression
  #####
  
  #*************
  # start logistic regression model building
  # set the IV and DV for the logistic regression model and model
  #*************
  # IV: corresponds to each unique snp (extracted using grep)
  mut = as.numeric(grepl(i,raw_data$dr_muts_pza))
  
  # DV: pyrazinamide 0 or 1
  dst = as.numeric(raw_data$pyrazinamide)
  
  tab = table(mut, dst)
  print(tab)
  
  # glm model: with and without maxit
  model = tryCatch( glm(dst ~ mut
                        , family = binomial
                        #, control = glm.control(maxit = 1) # only used when required for one step estimator
  ), warning = perfectSeparation)
  
  model = glm(dst ~ mut, family = binomial)
  
  print(summary(model))
  
  #**********
  # extract relevant model output
  #**********
  # extract log OR i.e the Beta estimate of the logistic model for a given snp
  my_logor = summary(model)$coefficients[2,1]
  print(paste0('Beta:', my_logor))
  
  # Dervive OR i.e exp(my_or) from the logistic model for a given snp
  #my_or = round(exp(summary(model)$coefficients[2,1]), roundto)
  my_or = exp(summary(model)$coefficients[2,1])
  print(paste0('OR:', my_or))
  
  # extract SE of the logistic model for a given snp
  my_se = summary(model)$coefficients[2,2]
  print(paste0('SE:', my_se))
  
  # extract Z of the logistic model for a given snp
  my_zval = summary(model)$coefficients[2,3]
  print(paste0('Z-value:', my_zval))

  # extract P-value of the logistic model for a given snp
  my_pval = summary(model)$coefficients[2,4]
  print(paste0('P-value:', my_pval))
  
  # extract confint interval of snp (2 steps, since the output is a named number)
  ci_mod = exp(confint(model))[2,]
  #my_ci = paste(ci_mod[["2.5 %"]], ",", ci_mod[["97.5 %"]])
  
  my_ci_lower = ci_mod[["2.5 %"]]
  my_ci_upper = ci_mod[["97.5 %"]]
  
  print(paste0('CI_lower:', my_ci_lower))
  print(paste0('CI_upper:', my_ci_upper))
  
  #*************
  # Assign the regression output in the to df (meta_pza_pnca_snps_only)
  # you can use ('=' or '<-/->')
  #*************
  #my_data$logistic_logOR[my_data$mutation == i] = my_logor
 
  my_or -> my_data$OR[my_data$mutation == i]
  
  my_pval -> my_data$pvalue[my_data$mutation == i]
  
  my_zval -> my_data$zvalue[my_data$mutation == i]
  
  my_se -> my_data$logistic_se[my_data$mutation == i]
  
  my_ci_lower -> my_data$ci_lower[my_data$mutation == i]
  
  my_ci_upper -> my_data$ci_upper[my_data$mutation == i]
  
  #=#=#=#=#=#=#=#
  # COMMENT: This assigns the relevant extracted output 
  # to the df and fills NA where the mutation (row) doesn't exist 
  # in my mutation list I am iterating over
  #=#=#=#=#=#=#=#
  
}