library(stringr)

# statcheck function adapted from the statcheck package by Michele Nuijten
# see github.com/michelenuijten/statcheck
statcheck <- function(
  x
){
  # Create empty objects
  raw_result = NULL
  genotype = NULL
  snp = NULL
  dv = NULL
  or_comparison = NULL
  or_result = NULL
  ci_confidence = NULL
  ci_lb = NULL
  ci_ub = NULL
  p_comparison = NULL
  p_result  =  NULL
  p_recalc  =  NULL
  test_stat_recalc <- NULL
  se_recalc = NULL

  if (length(x)==0) return(NULL)

  if (is.null(names(x))) names(x) <-  1:length(x)

  # create counter for loops
  index <- 1

  message("Extracting statistics...")
  pb <- txtProgressBar(max=length(x),style=3)
  for (i in 1:length(x)){

    txt <- x[i]

    # identify sequence of results and extract their text
    resLoc <- gregexpr("((genotype).*?rs[0-9]{1,10}|rs[0-9]{1,10}.*?(genotype)).*?(associate|relate|correlate)[a-z]{1,}?\\s(with|to).*?(odds\\sratio|\\(?OR\\)?).*?[0-9]{2}\\%\\s(confidence\\sinterval|\\(?ci\\)?).*?(p.*?\\s?[0-9]?.[0-9]{1,5})",
                       txt,
                       ignore.case = TRUE)[[1]]
    resContext <- substring(txt,
                            resLoc,
                            resLoc + attr(resLoc, "match.length") + 3)

    # locate data in context strings
    # genotype
    locator_genotype <- gregexpr("genotype.?[A-Za-z]{2}", resContext, ignore.case = TRUE)

    # SNP
    locator_snp <- gregexpr("\\srs[0-9]{1,10}\\s", resContext, ignore.case = TRUE)

    # DV
    locator_ass <- gregexpr("(associat[a-z]{1,}|
                           correlat[a-z]{1,}|
                           relat[a-z]{1,})\\s(with|to)\\s", resContext, ignore.case = TRUE)

    locator_dv <- gregexpr("(associat[a-z]{1,}|
                           correlat[a-z]{1,}|
                           relat[a-z]{1,})\\s(with|to)\\s[a-z]{1,}", resContext, ignore.case = TRUE)

    # Odds ratio
    locator_or <- gregexpr("(odds ratio.*|or.*?).*?[<>=]", resContext, ignore.case = TRUE)

    # Confidence interval
    locator_ci <- gregexpr("[0-9]{2}\\%.*?(confidence interval.*|ci.*?)[=:;].*?([0-9]{1,}?[.]?[0-9]{1,}.*?[0-9]{1,}?[.]?[0-9]{1,})", resContext, ignore.case = TRUE)

    # P-value (if present)
    locator_p <- gregexpr("p.?[<>=]", resContext, ignore.case = TRUE)

    for (j in 1:length(resContext)){
      # extract data from each context with use of locators
      # genotype
      genotype_ind <- substr(resContext[[j]], locator_genotype[[j]] + 10 - 1, locator_genotype[[j]] + 10)
      # SNP
      snp_temp <- stringr::str_match_all(resContext[[j]],
                                "\\srs[0-9]{1,10}\\s")[[1]][1]
      snp_ind <- gsub(pattern = "\\s", x = snp_temp, replacement = "")
      # DV
      dv_ind <- gsub(pattern = "\\s",
                     x = substr(resContext[[j]], locator_ass[[j]] + attr(locator_ass[[j]], "match.length"), locator_dv[[j]] + attr(locator_dv[[j]], "match.length")),
                     "")
      # OR
      or_temp <- substr(resContext[[j]], locator_or[[j]], locator_or[[j]] + attr(locator_or[[j]], "match.length") + 5)
      or_comp <- str_match_all(or_temp,
                               "[<>=]")[[1]][1]
      or_ind <- as.numeric(str_match_all(or_temp,
                              "[0-9]{1,}?[.][0-9]{1,3}")[[1]][1])
      # CI
      ci_temp <- substr(resContext[[j]], locator_ci[[j]], locator_ci[[j]] + attr(locator_ci[[j]], "match.length") + 5)
      ci_conf <- as.numeric(substr(str_match_all(ci_temp, "[0-9]{1,2}\\%\\s?[cC]")[[1]][,1], 0, 2))
      ci_conc <- str_match_all(ci_temp, "[0-9]{1,}?[.][0-9]{1,3}")[[1]][,1]
      ci_lb_ind <- as.numeric(min(ci_conc))
      ci_ub_ind <- as.numeric(max(ci_conc))
      # P-value
      p_temp <- substr(resContext[[j]], locator_p[[j]], locator_p[[j]] + attr(locator_p[[j]], "match.length") + 5)
      p_comp <- str_match_all(p_temp,
                              "[<>=]")[[1]][1]
      p_ind <- as.numeric(str_match_all(p_temp,
                             "[0-9]{1,}?[.][0-9]{1,3}")[[1]][1])

      # recalculate p-value
      # from Altman, D. G., & Bland, J. M. (2011). How to obtain the P value from a confidence interval. BMJ , 343, d2304.
      se_ind <- (ci_ub_ind - ci_lb_ind) / (2 * qnorm((1 - (ci_conf / 100)) / 2, lower.tail = FALSE))
      z_ind <- or_ind / se_ind
      p_recalc_ind <- pnorm(z_ind, lower.tail = FALSE) * 2

      # write back results into main objects
      raw_result[index] = resContext[[j]]
      genotype[index] = genotype_ind
      snp[index] = snp_ind
      dv[index] = dv_ind
      or_comparison[index] = or_comp
      or_result[index] = or_ind
      ci_confidence[index] = ci_conf
      ci_lb[index] = ci_lb_ind
      ci_ub[index] = ci_ub_ind
      p_comparison[index] = p_comp
      p_result [index] =  p_ind
      p_recalc [index] =  p_recalc_ind
      test_stat_recalc[index] = z_ind
      se_recalc[index] = se_ind

      index = index + 1
    }
  }

  # data frame
  Res <- data.frame(raw_result = raw_result,
                    genotype = genotype,
                    snp = snp,
                    dv = dv,
                    or_comparison = or_comparison,
                    or_result = or_result,
                    ci_confidence = ci_confidence,
                    ci_lb = ci_lb,
                    ci_ub = ci_ub,
                    p_comparison = p_comparison,
                    p_result  = p_result,
                    p_recalc  = p_recalc,
                    test_stat_recalc = test_stat_recalc,
                    se_recalc = se_recalc)

class(Res) <- c("statcheck","data.frame")

###---------------------------------------------------------------------
# Return message when there are no results
if(nrow(Res)>0){

  return(Res)
} else {
  Res <- cat("statcheck did not find any results\n")
}
}


## Setup R error handling to go to stderr
#options(show.error.messages=F,
        #error = function () { cat( geterrmessage(), file=stderr() ); q( "no", 1, F ) } )
warnings()

library(RColorBrewer)
library(lattice)
library(latticeExtra)
library(grid)
library(gridExtra)
library(optparse)

option_list <- list(
    make_option("--h_dataframe", type="character", help="path to h-signature dataframe"),
    make_option("--z_dataframe", type="character", help="path to z-signature dataframe"),
    make_option("--plot_method", type = "character", help="How  data should be plotted (global or lattice)"),
    make_option("--pdf", type = "character", help="path to the pdf file with plots"),
    make_option("--title", type = "character", help="Graph Title")
    )

parser <- OptionParser(usage = "%prog [options] file", option_list = option_list)
args = parse_args(parser)

# data frames implementation
h_dataframe = read.delim(args$h_dataframe, header=F)
colnames(h_dataframe) = c("chrom", "overlap", "sig", "z-score")
h_dataframe$sig = h_dataframe$sig * 100  # to get probs in %
z_dataframe = read.delim(args$z_dataframe, header=F)
colnames(z_dataframe) = c("chrom", "overlap", "sig", "z-score")

# functions
      globalgraph = function () {
        pdf(args$pdf)
        par(mfrow=c(2,2),oma = c(0, 0, 3, 0))

        plot(z_dataframe[z_dataframe$chrom == "all_chromosomes", c(2,3)],
             type = "h", main="Numbers of pairs", cex.main=1, xlab="overlap (nt)",
             ylab="Numbers of pairs", col="darkslateblue", lwd=4)

        plot(z_dataframe[z_dataframe$chrom == "all_chromosomes", c(2,4)],
             type = "l", main="Number of pairs Z-scores", cex.main=1, xlab="overlap (nt)",
             ylab="z-score", pch=19, cex=0.2, col="darkslateblue", lwd=2)

        plot(h_dataframe[h_dataframe$chrom == "all_chromosomes", c(2,3)],
             type = "l", main="Overlap probabilities", cex.main=1, xlab="overlap (nt)",
             ylab="Probability [%]", ylim=c(0,50), pch=19, col="darkslateblue", lwd=2)

        plot(h_dataframe[h_dataframe$chrom == "all_chromosomes", c(2,4)],
             type = "l", main="Overlap Probability Z-scores", cex.main=1,
             xlab="overlap (nt)", ylab="z-score", pch=19, cex=0.2,
             col="darkslateblue", lwd=2)

        mtext(args$title, outer = TRUE, cex=1)
        dev.off()
      }

      treillisgraph = function (df, ...) {
          pdf(args$pdf, paper="special", height=11.69, width=6 ) # 8.2677
          p = xyplot(sig ~ overlap|factor(method, levels=unique(method))+chrom, data = df,
                   type = "l",
                   col='darkblue',
                   cex=0.5,
                   scales=list(y=list(tick.number=4, relation="free", cex=0.6, rot=0), x=list(cex=0.6, alternating=FALSE)),
                   xlab = "Overlap",
                   ylab = "signature (Nbr of pairs / Overlap prob.)",
                   main = args$title,
                   par.strip.text=list(cex=.5),
                   pch=19, lwd =2,
                   as.table=TRUE,
                   layout=c(2,12),
                   newpage = T,
                   ...)
           plot(p)
           dev.off()
      }

# main

if (args$plot_method=="global") {
    globalgraph()
    }

if(args$plot_method=="lattice") {
    # rearrange dataframes
    h_sig = h_dataframe[,c(1,2,3)]
    h_sig = cbind(rep("Overlap Prob (%)", length(h_sig[,1])), h_sig)
    colnames(h_sig) = c("method", "chrom", "overlap", "sig")
    z_pairs = z_dataframe[,c(1,2,3)]
    z_pairs = cbind(rep("Nbr of pairs", length(z_pairs[,1])), z_pairs)
    colnames(z_pairs) = c("method", "chrom", "overlap", "sig")
    lattice_df = rbind(z_pairs, h_sig)
    par.settings.treillis=list(strip.background = list(
            col = c("lightblue", "lightgreen")))

    treillisgraph(lattice_df, par.settings=par.settings.treillis)
}


source.dir        <-  file.path("C:","mv","SMS")

add.sum<-T   # sum SSB, Yield and Recruits
last.y<-"2010"

dirs<-c("SAN-area-1","SAN-area-2","San-area-3")
label<-c("area 1","area 2","area 3")

##################################################

retro.dirs      <-  as.list(file.path(source.dir,dirs))

#Import data
a<-lapply(retro.dirs,function(d) {
  file<-file.path(d,'summary_table_raw.out')
   aa<-read.table(file,header=TRUE)
  aa$SMS<-label[which(d==retro.dirs)]
  aa
})

a<-do.call(rbind,a)
head(a)

a$Recruits<-a$Rec/1000000
a$SSB<-a$SSB/1000
a$MeanF<-a$mean.F
a$Yield<-a$SOP/1000

a<-subset(a,select=c(Year,Recruits,SSB,MeanF,Yield,SMS))

if (add.sum) {
  b<-a
  b<-aggregate(list(SSB=b$SSB, Recruits=b$Recruits, Yield=b$Yield),list(Year=b$Year),sum)

  b$SMS<-"Sum"
  b$MeanF<-NA
  a<-rbind(a,b)
}

graphics.off()
trellis.device(device = "windows",
               color = T, width=10, height=13,pointsize = 12,
               new = T, retain = FALSE)

myCol<- seq(1,length(dirs)+1)
myLty<- c(rep(1,length(dirs)),3)

print(
  xyplot(SSB+Recruits+MeanF+Yield~Year,groups=SMS,data=a, ylab=NULL,  xlab=NULL,  subset=Year>1981,
       layout=c(1,4), between = list(y = c(1, 1),x = c(1, 1)),
       scales = list(x = list( cex=0.8,relation='same'), y= list(cex=0.8),alternating = 1,relation='free'),
      key = list(space = "bottom", points = F, lines = T,cex=1, columns = 3,title='Sandeel',col=myCol, lwd=3,pch=NA,lty=myLty,
               text = list(lab = as.character(unique(a$SMS)),col=myCol )),

        panel = function(x, y, subscripts, groups) {
         panel.superpose(x, y,subscripts=subscripts, groups=groups,type='l',col=myCol,lwd=2,lty=myLty )
        }
))