model_cumulttfrom <- function (calendarMoments_t1 = c('Sowing'),
         calendarCumuls_t1 = c(0.0),
         cumulTT = 8.0){
    #'- Name: CumulTTFrom -Version: 1.0, -Time step: 1
    #'- Description:
    #'            * Title: CumulTTFrom Model
    #'            * Author: Pierre Martre
    #'            * Reference: Modeling development phase in the
    #'                Wheat Simulation Model SiriusQuality.
    #'                See documentation at http://www1.clermont.inra.fr/siriusquality/?page_id=427
    #'            * Institution: INRA Montpellier
    #'            * Abstract: Calculate CumulTT
    #'- inputs:
    #'            * name: calendarMoments_t1
    #'                          ** description : List containing appearance of each stage at previous day
    #'                          ** variablecategory : state
    #'                          ** datatype : STRINGLIST
    #'                          ** default : ['Sowing']
    #'                          ** unit :
    #'                          ** inputtype : variable
    #'            * name: calendarCumuls_t1
    #'                          ** description : list containing for each stage occured its cumulated thermal times at previous day
    #'                          ** variablecategory : state
    #'                          ** datatype : DOUBLELIST
    #'                          ** default : [0.0]
    #'                          ** unit : °C d
    #'                          ** inputtype : variable
    #'            * name: cumulTT
    #'                          ** description : cumul TT at current date
    #'                          ** datatype : DOUBLE
    #'                          ** variablecategory : auxiliary
    #'                          ** min : -200
    #'                          ** max : 10000
    #'                          ** default : 8.0
    #'                          ** unit : °C d
    #'                          ** inputtype : variable
    #'- outputs:
    #'            * name: cumulTTFromZC_65
    #'                          ** description :  cumul TT from Anthesis to current date
    #'                          ** variablecategory : auxiliary
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 5000
    #'                          ** unit : °C d
    #'            * name: cumulTTFromZC_39
    #'                          ** description :  cumul TT from FlagLeafLiguleJustVisible to current date
    #'                          ** variablecategory : auxiliary
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 5000
    #'                          ** unit : °C d
    #'            * name: cumulTTFromZC_91
    #'                          ** description :  cumul TT from EndGrainFilling to current date
    #'                          ** variablecategory : auxiliary
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 5000
    #'                          ** unit : °C d
    cumulTTFromZC_65 <- 0.0
    cumulTTFromZC_39 <- 0.0
    cumulTTFromZC_91 <- 0.0
    if ('Anthesis' %in% calendarMoments_t1)
    {
        cumulTTFromZC_65 <- cumulTT - calendarCumuls_t1[which(calendarMoments_t1 %in% 'Anthesis')]
    }
    if ('FlagLeafLiguleJustVisible' %in% calendarMoments_t1)
    {
        cumulTTFromZC_39 <- cumulTT - calendarCumuls_t1[which(calendarMoments_t1 %in% 'FlagLeafLiguleJustVisible')]
    }
    if ('EndGrainFilling' %in% calendarMoments_t1)
    {
        cumulTTFromZC_91 <- cumulTT - calendarCumuls_t1[which(calendarMoments_t1 %in% 'EndGrainFilling')]
    }
    return (list ("cumulTTFromZC_65" = cumulTTFromZC_65,"cumulTTFromZC_39" = cumulTTFromZC_39,"cumulTTFromZC_91" = cumulTTFromZC_91))
}


## This script replicates Shiller 1981 using a slightly different data set
## The replication code was written by Caraiani and Anghel (forthcoming)

behavioral_equity_market <- function(data_vintage, path_fig) {
    path_data <- here("data/vintages", data_vintage)

    ## First specify the packages of interest
    packages <- c("ggplot2", "readxl", "dplyr", "tidyr", "dplyr", "purrr", "stringr",
                  "Quandl", "viridis", "ggpmisc", "readstata13", "imfr", "writexl", "fredr")

    source(here("src/check_packages.r"))
    check_packages(packages) # Check whether packages are installed already and load them

    ## Collect data on S&P 500 prices and dividends from QUANDL
    ## Note: Quandl does not requires an API unless user exceeds 50 calls a day.

    curr_year   <- as.double(format(Sys.Date(), "%Y")) # don't change this

    first_years <- c(1871, 1947, 1960, 1980)
    # first_years <- c(1871)
    last_years  <- c(1947, 1979, 2009, as.double(format(Sys.Date(), "%Y")))
    # last_years <- c(1979)

    reg_results <- data.frame(matrix(ncol = 5, nrow = 0))
    reg_columns <- c("first year", "last year", "specification", "b", "r2")
    colnames(reg_results) <- reg_columns

    for (f in first_years) {
        first_year <- f

        for (t in last_years) {
            last_year <- t  # can be different from current year
            if (last_year <= first_year) next  # skip iteration if last year is smaller than first year

            #######################################
            ########### DATA PREP #################
            #######################################
            # Read data first
            div   <- readRDS(file = here(path_data, "sp500_dividends.rds"))
            price <- readRDS(file = here(path_data, "sp500_prices.rds"))

            # Prep data: Dividends
            div$month <- as.double(format(div$Date, "%m"))
            div$year <- as.double(format(div$Date, "%Y"))
            div <- subset(div, (year <= last_year & year >= first_year))
            div <- div[order(div$Date), ]
            div <- div %>%  # There may be multiple observations per year - we will choose the first observation of every year
                group_by(year) %>%
                filter(row_number() == 1)
            div$Date <- format(div$Date,"%Y")
            div <- div %>% rename(div = Value) %>% select(-c("month", "year"))

            # Prep data: Prices
            price$month <- as.double(format(price$Date, "%m"))
            price$year <- as.double(format(price$Date, "%Y"))
            price$day <- as.double(format(price$Date, "%d"))
            price <- subset(price, (year <= last_year & year >= first_year))
            price <- price[order(price$Date), ]
            price <- price %>%  # There may be multiple observations per year so we need to choose the last observation of every year
                     group_by(year) %>%
                     filter(row_number() == 1)

            price$Date <- format(price$Date,"%Y")
            price <- price %>% rename(price = Value)
            price <- price %>% select(-c("day", "month", "year"))

            ## Add interest rates from JST database
            rawdata <- read.dta13(here(path_data, "JSTdatasetR5.dta"))
            myvars <- c("year", "country", "cpi",
                        "stir",  # short-term gov bills
                        "ltrate" # long-term gov bonds
            )
            rates <- rawdata[rawdata$country == "USA", myvars]
            rates <- rates %>%
                    rename(Date = year)  %>%
                    mutate(Date = as.character(Date),
                           stir = stir/100,
                           ltrate = ltrate/100)

            ## Merging and indexing
            df <- left_join(div, price, by = "Date")
            df <- left_join(df, rates, by = "Date")

            df$Date <- as.Date(df$Date, format = "%Y")
            df$year <- as.double(format(df$Date, "%Y"))
            df <- df[order(df$year), ] # Note: the data must be ordered in increasing order of time for the Shiller algo to work

            df <- df %>%
                    mutate(div_index   = div/div[year == first_year]*100) %>%
                    mutate(price_index = price/price[year == first_year]*100) %>%
                    mutate(infl = (cpi/lag(cpi, 1) - 1)) %>%             # inflation defined as growth in yearly CPI inflation
                    mutate(stir_real = (1 + stir) / (1 + infl) - 1) %>%  # deflate nominal shot-term interest rates
                    mutate(ltrate_real = (1 + ltrate) / (1 + infl) - 1) %>% # deflate nominal long-term interest rates
                    mutate(price_div_ratio = price / div,
                           div_price_ratio = div / price) %>%
                    mutate(price_lead5 = lead(price, n = 5),
                        annual_return_5yr = (price_lead5/price)^(1/5) - 1,   # real annual returns
                        annual_return_5yr_minus_stir = annual_return_5yr - stir_real,
                        annual_return_5yr_minus_ltrate = annual_return_5yr - ltrate_real,
                        year = as.double(format(Date, "%Y")),
                        ) %>%
                    mutate(return_5yr = (1 + annual_return_5yr)^5 - 1,       # real returns over the whole period
                           return_5yr_minus_stir = (1 + annual_return_5yr_minus_stir)^5 - 1,
                           return_5yr_minus_ltrate = (1 + annual_return_5yr_minus_ltrate)^5 - 1)

            ## Shiller method of detrending data and constructing PV of dividends

            # Compute long run exponential growth rate (last observation - first observation)
            g <- (df[df$year==last_year, "price"] / df[df$year==first_year, "price"]) ^ (1/(nrow(df)-1)) - 1

            # Detrend prices
            df$price_detrended <- df$price
            for (i in 1:nrow(df)) {
                df[i, "price_detrended"] <- df[i, "price_detrended"] / (1+g)^(i-nrow(df))
            }

            # Detrend dividends
            df$div_detrended <- df$div
            for (i in 1:nrow(df)) {
                df[i, "div_detrended"] <- df[i, "div_detrended"] / (1+g)^(i-nrow(df))
            }

            # Compute constant discount rate as the mean yield of stocks over chosen period
            r <- mean(df$div_detrended) / mean(df$price_detrended)

            # Set terminal price
            tp <- mean(df$price_detrended)

            # Compute fundamental values
            df$p_star <- df$price_detrended
            df[nrow(df),"p_star"] <- tp
            for (i in (nrow(df)-1):1) {
                df[i,"p_star"] = (df[i+1,"p_star"] + df[i,"div_detrended"]) / (1 + r)
            }

            #######################################
            ########### PLOTTING ##################
            #######################################

            ## PLOTTING Shiller graph
            # load the formatting and saving settings from another script
            source(here("src/format_plots.r"))
            source(here("src/save_plots.r"))

            if (last_year - first_year <= 60) {
                xticks = 10
            } else {
                xticks = 20
            }
            xlimits <- c(as.double(first_year), as.double(format(Sys.Date(), "%Y")))

            # Wide to long format (better for ggplot)
            df_tmp <- df %>%
                       select(c("price_detrended", "p_star", "year")) %>%
                       pivot_longer(c("price_detrended", "p_star"), names_to = "series", values_to = "value")
            df_tmp$series <- gsub("price_detrended", "Actual", df_tmp$series, fixed = TRUE)
            df_tmp$series <- gsub("p_star", "Expected", df_tmp$series, fixed = TRUE)

            ggplot(data = df_tmp, aes(x = year, y = value)) +
                geom_line(aes(color = series, linetype = series), size = 2) +
                scale_x_continuous(breaks=seq(xlimits[1], xlimits[2], xticks)) +
                labs(y = "", x = "", title = "") +
                format +
                scale_linetype_manual(values = c("solid", "dotdash")) +
                scale_color_manual(values = c(purple, green))

            # Save
            savefolder <- here(path_fig, "behavioral_finance")
            filename <- paste("equity_market_", as.character(first_year), "_", as.character(last_year), ".png", sep = "")
            save_plots(filename, savefolder)

            ## PLOTTING price to dividend against subsequent 7-year return
            df_tmp <- na.omit(df)
            nth_year <- 5 # Label every nth year
            df_tmp <- df_tmp %>% mutate(year = as.character(year))
            df_tmp$year[seq(1, nrow(df_tmp), nth_year)] <- NA
            formula = y ~ x
            xlimits <- c(0, NA)

            ggplot(data = df_tmp, aes(x = div_price_ratio, y = annual_return_5yr)) +
                geom_smooth(method = lm, formula = formula, se = FALSE, color = purple, size = 2) +
                geom_text(aes(label = year), fontface = "bold",  size = 5, color = green,
                         hjust = 0, nudge_x = 0, check_overlap = TRUE) +
                stat_poly_eq(formula = formula,
                            aes(label = paste(stat(rr.label), sep = ", ")),
                            label.x = "right", label.y = "top", size = 6, geom = "label_npc"
                            ) +
                xlim(xlimits) +
                labs(
                    x = "Dividend-to-price ratio",
                    y = "5-year annual return",
                ) +
                format

            filename <- paste("equity_market_pd_return_", as.character(first_year), "_", as.character(last_year), ".png", sep = "")
            save_plots(filename, savefolder)

            ggplot(data = df_tmp, aes(x = div_price_ratio, y = annual_return_5yr_minus_stir)) +
                geom_smooth(method = lm, formula = formula, se = FALSE, color = purple, size = 2) +
                geom_text(aes(label = year), fontface = "bold",  size = 5, color = green,
                         hjust = 0, nudge_x = 0, check_overlap = TRUE) +
                stat_poly_eq(formula = formula,
                            aes(label = paste(stat(rr.label), sep = ", ")),
                            label.x = "right", label.y = "top", size = 6, geom = "label_npc"
                            ) +
                xlim(xlimits) +
                labs(
                    x = "Dividend-to-price ratio",
                    y = "5-year annual return - 3-month Bill yield",
                ) +
                format

            filename <- paste("equity_market_pd_return_stir", as.character(first_year), "_", as.character(last_year), ".png", sep = "")
            save_plots(filename, savefolder)

            ggplot(data = df_tmp, aes(x = div_price_ratio, y = annual_return_5yr_minus_ltrate)) +
                geom_smooth(method = lm, formula = formula, se = FALSE, color = purple, size = 2) +
                geom_text(aes(label = year), fontface = "bold",  size = 5, color = green,
                         hjust = 0, nudge_x = 0, check_overlap = TRUE) +
                stat_poly_eq(formula = formula,
                            aes(label = paste(stat(rr.label), sep = ", ")),
                            label.x = "right", label.y = "top", size = 6, geom = "label_npc"
                            ) +
                xlim(xlimits) +
                labs(
                    x = "Dividend-to-price ratio",
                    y = "5-year annual return - 10-year Bond yield",
                ) +
                format

            filename <- paste("equity_market_pd_return_ltrate", as.character(first_year), "_", as.character(last_year), ".png", sep = "")
            save_plots(filename, savefolder)

            # caption <- paste("Data provider: Quandl. Original data source: Standard & Poor's and Robert Shiller.
            #                 Period: ", as.character(first_year), " - ", as.character(last_year), " Note:
            #                 The 7-year return is calculated as the 7-year growth in prices.", sep ="")

            #################################
            ########## REGRESSIONS ##########
            #################################

            # model<- lm(annual_return_5yr ~ price_div_ratio, data = df_tmp)
            # r1 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(annual_return_5yr_minus_stir ~ price_div_ratio, data = df_tmp)
            # r2 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(annual_return_5yr_minus_ltrate ~ price_div_ratio, data = df_tmp)
            # r3 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(annual_return_5yr ~ div_price_ratio, data = df_tmp)
            # r4 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(annual_return_5yr_minus_stir ~ div_price_ratio, data = df_tmp)
            # r5 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(annual_return_5yr_minus_ltrate ~ div_price_ratio, data = df_tmp)
            # r6 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(return_5yr ~ div_price_ratio, data = df_tmp)
            # r7 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(return_5yr_minus_stir ~ div_price_ratio, data = df_tmp)
            # r8 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # model<- lm(return_5yr_minus_ltrate ~ div_price_ratio, data = df_tmp)
            # r9 <- c(first_year, last_year, as.character(summary(model)$call[2]), summary(model)$coefficients[2,1], summary(model)$r.squared)
            #
            # reg_results <- do.call(rbind, list(reg_results, r1, r2, r3, r4, r5, r6, r7, r8, r9))
            # colnames(reg_results) <- reg_columns
        }

    }

    # write_xlsx(reg_results, "table/reg_results_cochrane_2011.xlsx")
}


## Assorted routines for manipulation of
## Zernike polynomials, interferogram analysis,
## and other forms of optical testing

## Author: M.L. Peck (mlpeck54@gmail.com)
## Language: R (http://www.r-project.org/)
## Copyright (c) 2004-2021, M.L. Peck

## greyscale

grey256 <- gray(seq(0,1,length=256))
gray256 <- grey256

## A rainbow that may be better than R's rainbow() color palette

rygcb <- colorRampPalette(c("red", "yellow", "green", "cyan", "blue"), space = "Lab",
            interpolate="spline")

## Yet another rainbow made up of additive and subtractive primaries

rygcbm <- colorRampPalette(c("red", "yellow", "green", "cyan", "blue", "magenta"),
			space="Lab", interpolate="spline")

## comparison plot of n wavefront estimates

plotn <- function(...,
  labels=NULL, addContours=FALSE,
  wftype="net", col=rygcb(400), qt=c(.01, .99)) {
	fits <- list(...)
	nwf <- length(fits)
	wfget <- paste("wf",wftype,sep=".")
	wfi <- get(wfget, pos=fits[[1]])
	nr <- nrow(wfi)
	nc <- ncol(wfi)
	wfs <- array(0, dim=c(nr,nc,nwf))
	wfs[,,1] <- wfi
	for (i in 2:nwf) wfs[,,i] <- get(wfget,pos=fits[[i]])
	rm(fits)
	if(is.null(labels)) labels=1:nwf

	rdiff <- matrix(0,nwf*(nwf-1)/2,2)
	k <- 1
	for (i in 1:(nwf-1)) {
		for (j in (i+1):nwf) {
			rdiff[k,] <- quantile(wfs[,,i]-wfs[,,j], probs=qt, na.rm=TRUE)
			k <- k+1
		}
	}
	zlim <- range(rdiff)
	if (tolower(.Platform$OS.type)=="windows") {
          windows(width = 5*nwf, height = 5*nwf)
        } else {
          X11(width = 5*nwf,height = 5*nwf)
        }
	par(mar=c(0,0,2,0))
	split.screen(figs=c(nwf,nwf))
	for (i in 1:nwf) {
	  screen(i)
	  wfi <- wfs[,,i]
	  class(wfi) <- "pupil"
	  plot(wfi, col=col, addContours=addContours, eqa=(wftype=="net"), axes=FALSE)
	  title(main=paste(labels[i],"rms=",format(pupilrms(wfi),digits=3)))
	}
	scr <- nwf+1
	for (row in 1:(nwf-1)) {
	  screen(scr)
	  plot(0:1,0:1,type="n",axes=FALSE)
	  text(0.5,0.5,labels[row])
	  scr <- scr+row
	  for (i in (row+1):nwf) {
		screen(scr)
		image(1:nr,1:nc,wfs[,,row]-wfs[,,i],col=grey256,asp=1,axes=FALSE,zlim=zlim,useRaster=TRUE)
		title(main=paste("rms diff =",format(pupilrms(wfs[,,row]-wfs[,,i]), digits=3)))
		scr <- scr+1
	  }
	}
	close.screen(all.screens=T)
}

## plot cross sections through a wavefront map

plotxs <- function(wf, cp, theta0=0, ylim=NULL, N=4, n=101,
	col0="black", col="gray", lty=2) {
  theta0 <- theta0*pi/180
  ixy <- seq(-1,1,length=n)
  theta <- (0:(N-1))*pi/N
  ix <- cp$rx*ixy
  iy <- cp$ry*ixy
  if (is.null(ylim)) ylim <- range(wf, na.rm=TRUE)
  plot(range(ixy),ylim, type="n", xlab="rho", ylab="Height")
  for (i in 1:N) {
	iix <- round(cp$xc+cos(theta[i])*ix)
	iiy <- round(cp$yc+sin(theta[i])*iy)
	points(ixy, diag(wf[iix,iiy]), type="l", col=col, lty=lty)
  }
  iix <- round(cp$xc+cos(theta0)*ix)
  iiy <- round(cp$yc+sin(theta0)*iy)
  points(ixy, diag(wf[iix,iiy]), type="l", col=col0, lty=lty)
}

###########
## Various fun stuff
###########

## Star test simulator & support routines

## calculate phase values for wavefront at wavelength lambda
## replaces NA values with 0

wftophase <- function(X, lambda = 1) {
    phi <- exp(2i*pi*X/lambda)
    phi[is.na(phi)] <- 0
    phi
}

## puts matrix X into corner of npadded x npadded matrix
## padded with zeroes

padmatrix <- function(X, npad, fill=0) {
    nr <- nrow(X)
    nc <- ncol(X)
    xpad <- matrix(fill, npad, npad)
    xpad[1:nr,1:nc] <- X
    xpad
}

## extract a matrix from the center of a larger matrix

submatrix <- function(X,size=255) {
    nr <- nrow(X)
    X[((nr-size)/2+1):((nr+size)/2),((nr-size)/2+1):((nr+size)/2)]
}

## shuffle quadrants of a 2d fft around to display as an image

fftshift <- function(X) {
    nr <- nrow(X)
    XS <- matrix(0,nr,nr)
    XS[1:(nr/2),1:(nr/2)] <- X[(nr/2+1):nr,(nr/2+1):nr]
    XS[(nr/2+1):nr,(nr/2+1):nr] <- X[1:(nr/2),1:(nr/2)]
    XS[(nr/2+1):nr,1:(nr/2)] <- X[1:(nr/2),(nr/2+1):nr]
    XS[1:(nr/2),(nr/2+1):nr] <- X[(nr/2+1):nr,1:(nr/2)]
    XS
}

## computes & displays fraunhofer diffraction pattern
## & mtf for wavefront described in zernike coefficients zcoef

startest <- function(wf=NULL, zcoef=NULL, maxorder=14L, phi=0,
	lambda = 1, defocus=5, cp=NULL,
	obstruct=NULL,
	npad = 4,
	gamma=2, psfmag=2, displaymtf=TRUE, displaywf=FALSE) {

    if (tolower(.Platform$OS.type) == "windows") {
      windows(width=15, height=5)
    } else {
      x11(width=15,height=5)
    }
    screens<- split.screen(c(1,3))

    if (is.null(wf)) {
      wf <- pupil(zcoef=zcoef, maxorder=maxorder, phi=phi, piston=0)
      cp <- cp.default
    } else {
      if (is.null(cp)) {
        stop("must specify cp if wf is non-null")
      }
    }
    nrow <- nrow(wf)
    ncol <- ncol(wf)
    wf.df <- pupil(zcoef=c(0, 0, 0, 1), maxorder=2, nrow=nrow, ncol=ncol, cp=cp)
    if (!is.null(obstruct)) {
      wf[is.na(wf.df)] <- NA
    }
    lx <- round(2*cp$rx)+1
    ly <- round(2*cp$ry)+1
    npad <- npad * .up2(lx,ly)

    phase <- wftophase(wf,lambda)
    up <- Mod(fft(padmatrix(phase,npad)))
    up <- up*up

    otf <- fft(up, inverse=TRUE)/npad^2
    otf <- otf[1:lx,1:ly]
    mtf <- Re(otf)
    mtf <- mtf/max(mtf)
    freqx <- seq(0,1,length=lx)
    freqy <- seq(0,1,length=ly)
    mtfideal <- 2/pi*(acos(freqx)-freqx*sqrt(1-freqx^2))

    nrpsf <- max(lx,ly)
    psf <- submatrix(fftshift(up),floor(nrpsf/psfmag))
    screen(screens[2])
    image(psf^(1/gamma),col=grey256,asp=1,bty='n', axes=FALSE, useRaster=TRUE)
    mtext("0")

    if (defocus >5) nrpsf <- 2*nrpsf
    if (defocus >15) nrpsf <- npad

    phase <- wftophase(wf - defocus/3.46*lambda*wf.df, lambda)
    up <- Mod(fft(padmatrix(phase,npad)))
    up <- up*up
    psf2 <- submatrix(fftshift(up),nrpsf)
    screen(screens[1])
    image(psf2^(1/gamma),col=grey256, asp=1,bty='n', axes=FALSE, useRaster=TRUE)
    mtext(-defocus)

    phase <- wftophase(wf + defocus/3.46*lambda*wf.df, lambda)
    up <- Mod(fft(padmatrix(phase,npad)))
    up <- up*up
    psf2 <- submatrix(fftshift(up),nrpsf)
    screen(screens[3])
    image(psf2^(1/gamma),col=grey256,asp=1,bty='n',axes=FALSE, useRaster=TRUE)
    mtext(defocus)
    close.screen(all.screens=TRUE)

    if (displaywf) {
            if (tolower(.Platform$OS.type) == "windows") windows() else x11()
            plot(wf/lambda)
    }

    if (displaymtf) {
            if (tolower(.Platform$OS.type) == "windows") windows() else x11()
            plot(freqx,mtf[1,],type="l",ylab="mtf",xlab="relative frequency")
            title(main='MTF vs. ideal')
            lines(freqy,mtf[,1])
            lines(freqx,mtfideal, lty=5)
            grid()
    }
    list(psf=psf, otf=otf, mtf=mtf)
}

## Kolmogorov turbulence

turbwf <- function(friedratio=1, zlist=makezlist(2,40), reps=1) {
    require(mvtnorm)
    dimz <- length(zlist$n)
    cova <- matrix(0, nrow=dimz, ncol=dimz)
    c0 <- friedratio^(5/3)*gamma(14/3)*(4.8*gamma(1.2))^(5/6)*(gamma(11/6))^2/(2^(8/3)*pi)
    for (i in 1:dimz) {
        for (j in i:dimz) {
            if ((zlist$m[i] == zlist$m[j]) && (zlist$t[i] == zlist$t[j])) {
                n <- zlist$n[i]
                np <- zlist$n[j]
                cova[i,j] <- (-1)^((n+np-2*zlist$m[i])/2) * sqrt((n+1)*(np+1)) *
                                gamma((n+np-5/3)/2)/(gamma((np-n+17/3)/2)*gamma((n-np+17/3)/2)*
                                gamma((n+np+23/3)/2))
                cova[j,i] <- cova[i,j]
            }
        }
    }
    cova <- cova * c0/(4*pi^2)
    zcoef.turb <- rmvnorm(reps, mean=rep(0,dimz), sigma=cova)
    if (reps==1) zcoef.turb <- as.vector(zcoef.turb)
    list(zcoef.turb=zcoef.turb, V=cova)
}

## Foucault test simulator

foucogram <- function(wf, edgex = 0, phradius = 0, slit=FALSE, pad=4, gamma=1,
                      map =FALSE, lev=0.5) {

    nr <- nrow(wf)
    nc <- ncol(wf)
    npad <- pad*nextn(max(nr,nc))
    phi<-padmatrix(wftophase(wf),npad)
    ca<-fftshift(fft(phi))
    u <- npad/2 + 1 + edgex
    ca[1:(u-phradius-1),]<-0
    ca[u,] <- .5*ca[u,]
    if (phradius>0) {
        if (slit) {
            for (i in 1:phradius) {
                f <- .5*(1+i/phradius)
                ca[(u+i),] <- f*ca[(u+i),]
                ca[(u-i),] <- (1-f)*ca[(u-i),]
            }
        } else {
            for (i in 1:phradius) {
                f <- .5 + i*sqrt(phradius^2-i^2)/(pi*phradius^2) + asin(i/phradius)/pi
                ca[(u+i),] <- f*ca[(u+i),]
                ca[(u-i),] <- (1-f)*ca[(u-i),]
            }
        }
    }
    ike <- Mod(fft(ca,inverse=T)[1:nr,1:nc])^2
    image(ike^(1/gamma),col=grey256,asp=nc/nr,axes=FALSE, useRaster=TRUE)
    if (map) {
        zmin <- floor(min(wf, na.rm=TRUE))
        zmax <- ceiling(max(wf, na.rm=TRUE))
        contour(wf,add=TRUE,levels=seq(zmin,zmax,by=lev),
                axes=FALSE,frame.plot=FALSE,col="red")
    }
    ike/max(ike)
}

## Synthetic interferogram

## notes: must have either non-null wf or zcoef. If wf is NULL
##        and cp is null make a new wavefront with current defaults
##        for matrix size and cp in pupil(), otherwise use the values
##        passed here. If wf is non-null zcoef is ignored
##        and the correct cp must be passed.

synth.interferogram <- function(wf=NULL, zcoef=NULL, maxorder=NULL,
                                nr=nrow(wf), nc=ncol(wf), cp=NULL,
                                phi=0, addzc=rep(0,4), fringescale=1,
                                plots=TRUE) {
  if (is.null(wf)) {
    if (is.null(cp)) {
      wf <- pupil(zcoef=zcoef, maxorder=maxorder, phi=phi, piston=0)+
            pupil(zcoef=addzc, maxorder=2)
      nr <- nrow(wf)
      nc <- ncol(wf)
      cp <- cp.default
    } else {
      wf <- pupil(zcoef=zcoef, maxorder=maxorder, phi=phi, piston=0,
                  nrow=nr, ncol=nc, cp=cp) +
            pupil(zcoef=addzc, maxorder=2, nrow=nr, ncol=nc, cp=cp)
    }
  } else {
    if (is.null(cp)) {
      stop("must specify a value for cp if wf is non-null")
    }
    wf <- wf + pupil(zcoef=addzc, maxorder=2, nrow=nr, ncol=nc, cp=cp)
  }
  igram <- cos(2*pi*wf/fringescale)
  igram[is.na(igram)] <- 0
  class(igram) <- c("pupil", class(igram))
  if (plots) plot(igram, col=grey256, addContours=FALSE)
  igram
}

#########
## general utilities
#########

## crop an array

crop <- function(img, cp, npad=20) {
    nr <- dim(img)[1]
    nc <- dim(img)[2]
    xmin <- max(1, round(cp$xc-cp$rx-npad))
    xmax <- min(nr, round(cp$xc+cp$rx+npad))
    ymin <- max(1, round(cp$yc-cp$ry-npad))
    ymax <- min(nc, round(cp$yc+cp$rx+npad))
    cp.new <- cp
    cp.new$xc <- cp$xc-xmin+1
    cp.new$yc <- cp$yc-ymin+1
    if (length(dim(img)) > 2) img <- img[xmin:xmax,ymin:ymax,]
    else img <- img[xmin:xmax,ymin:ymax]
    list(im=img, cp=cp.new)
}

## general purpose 2D convolution using FFT's.

## kern is the convolution kernel

convolve2d <- function(im, kern) {
	nr <- nrow(im)
	nc <- ncol(im)
	nrp <- nr + nrow(kern) - 1
	ncp <- nc + ncol(kern) - 1
	xs <- nrow(kern) %/% 2 + nrow(kern) %% 2
	ys <- ncol(kern) %/% 2 + ncol(kern) %% 2
	npad <- nextn(max(nrp, ncp))
	kern <- padmatrix(kern, npad=npad)
	im <- padmatrix(im, npad=npad)
	im.f <- Re(fft(fft(kern)*fft(im), inv=T))
	im.f <- im.f[xs:(nr+xs-1),ys:(nc+ys-1)]/(npad^2)
	return(im.f)
}

## Gaussian blur. fw is the standard deviation
## of the gaussian convolution kernel, in pixels.

gblur <- function(X, fw = 0, details=FALSE) {
  if (fw == 0) return(X)
  XP <- X
  XP[is.na(XP)] <- 0
  nr <- nrow(X)
  nc <- ncol(X)
  ksize <- max(ceiling(4 * fw), 3)
  if ((ksize %% 2) == 0) ksize <- ksize+1
  xc <- (ksize %/% 2) + 1
  xs <- ((1:ksize)-xc)/fw
  gkern <- outer(xs, xs, function(x,y) (x^2+y^2))
  gkern <- exp(-gkern/2)
  gkern <- round(gkern/min(gkern))
  npad <- nextn(max(nr,nc)+ksize-1)
  XP <- padmatrix(XP, npad)
  kernp <- padmatrix(gkern, npad)
  XP <- Re(fft(fft(XP)*fft(kernp), inv=TRUE))/(npad^2)/sum(gkern)
  XP <- XP[xc:(nr+xc-1), xc:(nc+xc-1)]
  XP[is.na(X)] <- NA
  if (details)
	  return(list(gkern=gkern, X=XP))
  else return(XP)
}

## Plot a complex matrix.
## Maybe make cmat a class, so this becomes default plot routine?

plot.cmat <- function(X, fn="Mod", col=grey256, cp=NULL, zoom=1, gamma=1, ...) {
	  ## define some possibly useful functions. Note Mod2 produces power spectrum

	logMod <- function(X) log(1+Mod(X))
	Mod2 <- function(X) Mod(X)^2
	logMod2 <- function(X) log(Mod2(X))
	nr <- nrow(X)
	nc <- ncol(X)
	xs <- seq(-nr/2,nr/2-1,length=nr)
	ys <- seq(-nc/2,nc/2-1,length=nc)
	if (!is.null(cp)) {
		xs <- (cp$rx/nr)*xs
		ys <- (cp$ry/nc)*ys
	}
	xsub <- max(1, floor(1+nr/2*(1-1/zoom))) : min(nr, ceiling(nr/2*(1+1/zoom)))
	ysub <- max(1, floor(1+nc/2*(1-1/zoom))) : min(nc, ceiling(nc/2*(1+1/zoom)))
	image(xs[xsub], ys[ysub], (eval(call(fn, X))[xsub,ysub])^(1/gamma),
	  col=col, asp=1, xlab="k_x", ylab="k_y", useRaster=TRUE, ...)
}

pick.sidelobe <- function(imagedata, logm=FALSE, gamma=3) {
	imagedata <- imagedata - mean(imagedata)
	npad <- nextn(max(nrow(imagedata), ncol(imagedata)))
	im <- padmatrix(imagedata, npad=npad, fill=0)
	if (logm) fn <- "logMod2" else fn <- "Mod"
	im.fft <- fftshift(fft(im))
	plot.cmat(im.fft, fn=fn, gamma=gamma)
	cat("Click on the desired sidelobe peak\n")
	peak <- locator(n=1, type="p", col="red")
	cat("Click on background filter size\n")
	edge <- locator(n=1, type="n")
	hw <- round(sqrt((edge$x)^2+(edge$y)^2))
	symbols(0, 0, circles=hw, inches=FALSE, add=TRUE, fg="red")
	return(list(sl=c(round(peak$x), round(peak$y)), filter=hw))
}

## simple utility returns Euclidean length of a vector

hypot <- function(x) sqrt(crossprod(x))

## Zernike moments

zmoments <- function(zcoef, maxorder=14) {
    zcoef <- zcoef[-(1:3)]
    zlist <- makezlist(4, maxorder)
    nz <- length(zlist$n)
    sumstats <- NULL
    i <- 1
    repeat {
        if (i > nz) break
            if (zlist$m[i] == 0) {
                sumstats <- rbind(sumstats, c(zlist$n[i], zlist$m[i], zcoef[i]))
                i <- i+1
            }
            else {
                sumstats <- rbind(sumstats, c(zlist$n[i], zlist$m[i], hypot(zcoef[i:(i+1)])))
                i <- i+2
            }
    }
    sumstats <- data.frame(sumstats)
    names(sumstats) <- c("n", "m", "value")
    return(sumstats)
}

## these two functions provide rudimentary project management. The first adds
## Zernike coefficients and rotation angles from one or more fits to a matrix.
## The second separates polished in from instrumental aberrations (if possible)

addfit <- function(..., th=0, zcm=NULL, theta=numeric(0)) {
    fits <- list(...)
    nt <- length(fits)
    if (length(th)==1 && nt>1) th=rep(th, nt)
    theta <- c(theta, th*pi/180)
    nz <- length(fits[[1]]$zcoef.net)
    for (i in 1:nt) {
        zcm <- cbind(zcm, fits[[i]]$zcoef.net[4:nz])
    }
    list(zcm=zcm, theta=theta)
}

separate.wf <- function(zcm, theta, maxorder=14) {

    nt <- length(theta)
    zlist <- makezlist(4, maxorder)
    nz <- length(zlist$n)
    zcb <- matrix(0, nz, 4)
    sumstats <- NULL
    cx <- c(rep(1,nt),rep(0,nt))
    cy <- c(rep(0,nt),rep(1,nt))
    i <- 1
    repeat {
        if (i > nz) break
        if (zlist$m[i] == 0) {
            zcb[i,1] <- mean(zcm[i,])
            zcb[i,3] <- sd(zcm[i,])/sqrt(nt)
            sumstats <- rbind(sumstats, c(zlist$n[i], zlist$m[i], zcb[i,1],
                                NA, zcb[i, 3], rep(NA, 4)))
            i <- i+1
        } else {
            y <- c(zcm[i,],zcm[i+1,])
            rx <- c(cos(zlist$m[i]*theta),sin(zlist$m[i]*theta))
            ry <- c(-sin(zlist$m[i]*theta),cos(zlist$m[i]*theta))
            lsm <- lm(y ~ -1+rx+ry+cx+cy)
            cc <- coef(summary(lsm))[,1:2]
            zcb[i:(i+1),1] <- cc[1:2,1]
            zcb[i:(i+1),3] <- cc[1:2,2]
            if (nrow(cc)==4) {
                zcb[i:(i+1),2] <- cc[3:4,1]
                zcb[i:(i+1),4] <- cc[3:4,2]
            }
            lsm <- summary(lsm)
            sumstats <- rbind(sumstats,c(zlist$n[i],zlist$m[i],
                                hypot(zcb[i:(i+1),1]), hypot(zcb[i:(i+1),2]),
                                lsm$sigma, lsm$r.squared,lsm$fstatistic))
            i <- i+2
        }
    }
    zcb[is.na(zcb)] <- 0
    colnames(zcb) <- c("zc_mirror", "zc_inst", "se_zc_mirror", "se_zc_inst")
    colnames(sumstats)[1:7] <- c("n","m","rms_mirror", "rms_inst", "sigma","R2", "F")
    wf.mirror <- pupil.arb(zcoef=zcb[,1],zlist=zlist)
    wf.inst <- pupil.arb(zcoef=zcb[,2], zlist=zlist)
    list(zcb=data.frame(zcb),sumstats=data.frame(sumstats),
        wf.mirror=wf.mirror, wf.inst=wf.inst)
}


################################################################################################
# This code uses myExchangeFileH1 from ReadExchangeFileExample_WKRDBEST2 (Author:David Currie) ###
## This data is dummy data for H1 VS, FT, FO, SS, SA, FM and BV.
## This code does not present anything new. It was an opportunity for me to get to use RDBES format and play around with ratio and inclusion probabilities estimation.
## Ana Ribeiro Santos
########################################################################################

library(dplyr)

#load the data
source ("ReadExchangeFileExample_WKRDBEST2.r")
#lst(myExchangeFileH1)

# Because the data did not have sampled weight or total sample weight I had to estimated.
#add species name to FM table
myExchangeFileH1$FM$Species <- myExchangeFileH1$SA$SAspeciesCode[match(myExchangeFileH1$FM$SAid,myExchangeFileH1$SA$SAid)]
#head(myExchangeFileH1$FM)

#dim(myExchangeFileH1$FM)
#For this example we are goin to select just one species = 1015724 = Leuciscus aspius.
myExchangeFileH1$FM <- myExchangeFileH1$FM [myExchangeFileH1$FM$Species =="1015724",]

#calculate the weight per length. a and b paramters from Fisbase
a <- 0.01389
b <- 3

myExchangeFileH1$FM$Wt.at.len_g <- with (myExchangeFileH1$FM, ((as.numeric(FMclass)/10)^ b) * a) # this gives in grams, equation expected CM not MM hence divide by 10

#calculate the weight for each SAid (Sample ID)
sample_wght_SAid <- myExchangeFileH1$FM %>%
  group_by(SAid) %>%
  summarise (., SAsampleWeightLive_est = round (sum (Wt.at.len_g), 0)) %>% as.data.frame()

# merge it with SA table
myExchangeFileH1$SA <- merge (myExchangeFileH1$SA, sample_wght_SAid)
#head(myExchangeFileH1$SA)

#Calculate the totalweightlive based on the ratio_SA
myExchangeFileH1$SA$ratio_SA <- ifelse (is.na (myExchangeFileH1$SA$SAsampleWeightLive), with (myExchangeFileH1$SA, as.numeric (SAnumberTotal) / as.numeric (SAnumberSampled)),
                                        with (myExchangeFileH1$SA, as.numeric (SAtotalWeightLive) / as.numeric (SAsampleWeightLive)))

myExchangeFileH1$SA$SAtotalWeightLive_est <- with (myExchangeFileH1$SA, SAsampleWeightLive_est * ratio_SA)

##########################################################################################################################################################
# uses myExchnageFileH1 data and assumes simple random sampling within stages
# trips within domains (quarter, area and gear) and ages within lengths.
# In the code below, the hierarchy is as follows in order: VS, FT, "FO","SS","SA","FM","BV"

# set up the population totals
popDat <- as.data.frame(matrix(c(16,39,86,54, 32, 29, 1464760,3937180,6831707, 2876490, 3487500, 2345431),byrow=F,nrow=6,ncol=2))
names(popDat) <- c("tripNum","liveWt")
rownames(popDat) <-  c("Q2 27 GND", "Q3 27 DRB", "Q3 27 FPO", "Q4 27 FPO", "Q4 27 FYK", "Q4 27 GNS")

# calculate inclusion probabilities assuming SRS within strata
myExchangeFileH1$VS$VSincProb <- as.numeric (myExchangeFileH1$VS$VSnumberSampled)/ as.numeric (myExchangeFileH1$VS$VSnumberTotal)
myExchangeFileH1$BV$BVincProb <- as.numeric (myExchangeFileH1$BV$BVnumberSampled)/ as.numeric (myExchangeFileH1$BV$BVnumberTotal)
myExchangeFileH1$FM$FMincProb <- 1
myExchangeFileH1$SA$SAincProb <- as.numeric (myExchangeFileH1$SA$SAnumberSampled)/as.numeric (myExchangeFileH1$SA$SAnumberTotal)
myExchangeFileH1$SS$SSincProb <- as.numeric (myExchangeFileH1$SS$SSnumberSampled)/as.numeric (myExchangeFileH1$SS$SSnumberTotal)
myExchangeFileH1$FO$FOincProb <- as.numeric (myExchangeFileH1$FO$FOnumberSampled)/as.numeric (myExchangeFileH1$FO$FOnumberTotal)
myExchangeFileH1$FT$FTincProb <- as.numeric (myExchangeFileH1$FT$FTnumberSampled)/as.numeric (myExchangeFileH1$FT$FTnumberTotal)

# caclulate quarter from date, and then make a domain of quarter combined with area
# allocate domains to FM & BV as well
myExchangeFileH1$SA$FTid <- myExchangeFileH1$SS$FTid[match(myExchangeFileH1$SA$SSid,myExchangeFileH1$SS$SSid)]
myExchangeFileH1$SS$date <- myExchangeFileH1$FO$FOendDate[match(myExchangeFileH1$SS$FOid,myExchangeFileH1$FO$FOid)]
myExchangeFileH1$SA$date <- myExchangeFileH1$SS$date[match(myExchangeFileH1$SA$SSid,myExchangeFileH1$SS$SSid)]
myExchangeFileH1$SA$quarter <- paste("Q",(as.numeric(substr(myExchangeFileH1$SA$date,6,7))-1) %/% 3 + 1,sep="")
myExchangeFileH1$SA$domain <- paste(myExchangeFileH1$SA$quarter,myExchangeFileH1$SA$SAarea, myExchangeFileH1$SA$SAgear)
myExchangeFileH1$FM$domain <- myExchangeFileH1$SA$domain[match(myExchangeFileH1$FM$SAid,myExchangeFileH1$SA$SAid)]
myExchangeFileH1$BV$domain <- myExchangeFileH1$FM$domain[match(myExchangeFileH1$BV$FMid,myExchangeFileH1$FM$FMid)]
myExchangeFileH1$BV$numAtUnit <- 1

# add spp to FM
myExchangeFileH1$FM$spp <- myExchangeFileH1$SA$SAspeciesCode[match(myExchangeFileH1$FM$SAid,myExchangeFileH1$SA$SAid)]

# -----------------------------------------------------------------------
# a function to calculate inclusion probabilities for the units at the final stage
# of sampling given all the inclusion probabilities for the other stages
# -----------------------------------------------------------------------
getIncProb <- function(RDB,stages){
  nStages <- length(stages)
  if (any(stages %in% c("FM"))) {
    RDB[["FM"]][["FMincProb"]] <- 1
  }
  RDB[[stages[[1]]]][["incProb"]] <- RDB[[stages[[1]]]][[paste(stages[[1]],"incProb",sep="")]]
  for (i in 2:(nStages)) {
    indx <- RDB[[stages[[i]]]][[paste(stages[[i-1]],"id",sep="")]]
    indxPrev <- RDB[[stages[[i-1]]]][[paste(stages[[i-1]],"id",sep="")]]
    RDB[[stages[[i]]]][["incProbPrev"]] <- RDB[[stages[[i-1]]]][[paste("incProb",sep="")]][match(indx,indxPrev)]
    RDB[[stages[[i]]]][["incProb"]] <- RDB[[stages[[i]]]][["incProbPrev"]]*RDB[[stages[[i]]]][[paste(stages[[i]],"incProb",sep="")]]
  }
  return(RDB)
}

# -----------------------------------------------------------------------
# Number-at-Length
# -----------------------------------------------------------------------
# set up the stages in the sampling design used in the estimation
#stages <- list("FT","FO", "SS","SA","FM")
stages <- list("VS", "FT","FO", "SS","SA","FM")
#stages <- list("LE","SS","SA","FM")

# calculate the inclusion probabilities by length class given the other inclusion probabilities
# in the higher stages
test <- getIncProb(myExchangeFileH1,stages)

# calculate a Horvitz Thompson estimate for total numbers at length by domain
# assuming srs within the domain - this is valid for the RDBshare data
estL <- tapply(as.numeric(test$FM$FMnumberAtUnit)/test$FM$incProb,list(test$FM$FMclass,test$FM$domain),sum)

#estL <- test$FM %>%
#  group_by(FMclass, domain, spp) %>%
#  summarise (., est.L = sum (as.numeric (FMnumberAtUnit)/incProb))

# calculate a HT estimate for total landed weight using the sampled landed weights
# assuming srs etc as before
estX <- tapply(test$SA$SAtotalWeightLive_est/test$SA$incProbPrev,list(test$SA$domain),sum)/1e3

# get the relevant population totals
popX <- popDat[match(names(estX),rownames(popDat)),"liveWt"]

# calculate the ratio estimates for numbers at length
estLR <- estL*matrix(rep(popX/estX,dim(estL)[1]),byrow=T,ncol=dim(estL)[2])

# compare the estLR with the ratio estimation using the ratio between the sum total sampled weight and the landings in the domain
test <- myExchangeFileH1
sumX <- tapply(myExchangeFileH1$SA$SAtotalWeightLive_est,list(myExchangeFileH1$SA$domain),sum) # sum te total sampled weight for each domain
testLR <- estL*matrix(rep(popX/sumX,dim(estL)[1]),byrow=T,ncol=dim(estL)[2]) # multiply the number at length by the ratio between landings and sampled weight

testLR
estLR
## They are super different!!! So im wondering if Im doing something wrong here!! The fact that Im working with made up data does not help at all!!

########################################################################################################################################
# JUST to test if using a different approach from Liz's code and function we would get the same number at length up to Vessel selection'
#############################################################################################################################################
VSinfo <- myExchangeFileH1$VS %>% transmute(VSid,SDid,
                                            VSincProb)%>%distinct()

FTinfo <-myExchangeFileH1$FT%>%transmute(FTid, VSid,  FTincProb )%>%distinct()

FOinfo<-myExchangeFileH1$FO%>%transmute(FOid,FTid, time=paste("Q",(as.numeric(substr(myExchangeFileH1$FO$FOendDate,6,7))-1) %/% 3 + 1,sep=""),
                                        metier=FOgear, space=FOarea, FOincProb)%>%distinct()
FOinfo$domain <- paste (FOinfo$time, FOinfo$space, FOinfo$metier)

SSinfo<-myExchangeFileH1$SS%>%transmute(SSid, FOid, SSincProb)%>%distinct()
#SAinfo<-myExchangeFileH1$SA%>%transmute(SAid,SSid,spp=SAspeciesCode,
#                              wsamp=SAsampleWeightMeasured,wsamptot=SAtotalWeightMeasured)%>%distinct()

SAinfo<- myExchangeFileH1$SA%>%transmute(SAid,SSid,spp=SAspeciesCode,SAincProb)%>%distinct()

FMinfo<-myExchangeFileH1$FM%>%transmute(FMid,SAid,len=FMclass,n=as.numeric(FMnumberAtUnit), FMincProb)

samp_est_stratum_incPro <-FMinfo%>% left_join (SAinfo)%>%left_join(SSinfo)%>%left_join(FOinfo)%>%left_join(FTinfo)%>% left_join(VSinfo) %>%
  group_by(domain, spp, len, n) %>%
  summarise (., final_prob = FMincProb * SAincProb * SSincProb * FOincProb * FTincProb * VSincProb) %>%
  ungroup() %>%
  group_by(domain, spp, len) %>%
  summarise (., est.n = sum(n / final_prob))


######################
############FRAILTY INDEX CALCULATION

setwd("P:/2019 0970 151 000/User Data/Faraz-Export IDAVE folders")
library(tidyverse)
library(icd)
df <- read.csv("./datasets/New/NACRASCIHI_2_dupRemoved.csv")

tabna <- function(x){table(x, useNA = "ifany")}

ic <- df %>% select(ID, days_to_lefteddate, dx10code1:dx10code10, alc_status)%>%
mutate(IDm = paste(ID, days_to_lefteddate, sep="_"))

codes <- read.csv("./ICD10 codes/frailty_score2.csv")

# codes <- c("F00", "G81","G30", "I69", "R29", "N39", "F05", "W19","S00",
# "R31", "B96", "R41", "R26", "I67", "R56", "R40", "T83", "S06",
# "S42", "E87", "M25", "E86", "R54", "Z50", "F03", "W18", "Z75", "F01",
# "S80", "L03", "H54", "E53", "Z60", "G20", "R55", "S22", "K59", "N17",
# "L89", "Z22", "B95", "L97", "R44", "K26", "I95", "N19", "A41", "Z87",
# "J96", "X59", "", "", "", "", "", "", "", "",
# "", "", "", "", "", "", "", "", "", "")

# points <- c(7.1, 4.4, 4, 3.7, 3.6, 3.2, 3.2, 3.2, 3.2,
# 3.0, 2.9, 2.7, 2.6, 2.6, 2.6, 2.5, 2.4, 2.4,
# 2.3,2.3,2.3,2.3, 2.2, 2.1, 2.1, 2.1, 2.0, 2.0,
# 2.0,2.0,1.9,1.9,1.8,1.8,1.8,1.8,1.8,1.8,
# rep(1.7,3), rep(1.6, 6), 1.5,
# rep(1.5, 4), rep(1.4, 6),

# )

for (i in 1:nrow(codes)){
	id_ls <- ic %>% filter_at(vars(dx10code1:dx10code10),
	 any_vars(substr(.,1,3) == codes$ICD.Code[i])) %>%
	select(IDm , dx10code1:dx10code10) %>% .$IDm

	var_name <- paste0("frail_",as.character(codes$ICD.Code[i])) #dynamic name in mutate

	ic <- ic %>% mutate(!!var_name := ifelse(IDm %in% id_ls, 1, 0))
}

# for (i in 1:nrow(codes)){
	# var_name <- paste0("frail_",as.character(codes$ICD.Code[i])) #dynamic name in mutate

	# ic2 <- ic %>%
	# mutate_at(vars(dx10code1:dx10coSde10), ~substr(.,1,3)) %>%
	# mutate(!!var_name := ifelse(any(vars(dx10code1:dx10code10) %in% codes$ICD.Code[i]), 1, 0))

# }
##Occurences are low usually
ic %>% select(frail_F00:frail_R50) %>% colSums / nrow(ic) * 100 %>% round(digits=2)

mat <- as.matrix(ic %>% select(frail_F00:frail_R50))

scores <- mat %*% codes$coef
####Check Safety
scores[1:5]
ic[1,] %>% select(frail_F00:frail_R50) %>% select_if(.==1)
############
ic$frailty_score <- scores
summary(ic$frailty_score)

ic %>% select(IDm, frailty_score) %>% write.csv("./ICD10 codes/frailty_index.csv", row.names = F)


#required packages
library("openxlsx")
library("stringr")
library(ComplexHeatmap)
library(circlize)
library(colorspace)
library(GetoptLong)
library(ggplot2)
library(gplots)
library(survival)
require("survival")
# install.packages("svglite")
library(survminer)
library(reticulate)

#set the project path
project_path = "./project/"

#load patient data (include survival info)
study_cohort = read.xlsx(paste(project_path, "data/tcga/study_cohort.xlsx",sep=""), sheet = 1, colNames = TRUE)
entire_cohort = read.xlsx(paste(project_path, "data/tcga/entire_cohort.xlsx",sep=""), sheet = 1, colNames = TRUE)

length(rownames(study_cohort))

length(rownames(entire_cohort))

median(entire_cohort$ENSG00000206567.8)

#To build a risk score system based on lncRNA signature, we first estimate the coef of the signature
fit = coxph(Surv(OS, Censor) ~

ENSG00000187185.4+
ENSG00000259641.4+
ENSG00000218510.5 +
ENSG00000257989.1+
ENSG00000206567.8
, data=study_cohort)

# summary(fit)

HR <- exp(coef(fit))
CI <- exp(confint(fit))
P <- round(coef(summary(fit))[,5],8)
COEF <- coef(fit)

colnames(CI) <- c("Lower", "Higher")

table2 <- as.data.frame(cbind(HR, CI,COEF, P))
table2

#calculate the risk score based on expression and coef of the signature
#use coef (estimated in the study cohort) in above table
entire_cohort$risk_lnc = (
    entire_cohort$ENSG00000187185.4   * table2[1,4] +
    entire_cohort$ENSG00000259641.4 * table2[2,4]+
    entire_cohort$ENSG00000218510.5 * table2[3,4] +
    entire_cohort$ENSG00000257989.1 * table2[4,4] +
    entire_cohort$ENSG00000206567.8 * table2[5,4] )

#calculate hazard ratio
fit1 = coxph(Surv(OS, Censor) ~

risk_lnc
, data=entire_cohort)

# summary(fit)

HR <- exp(coef(fit1))
CI <- exp(confint(fit1))
P <- coef(summary(fit1))[,5]
COEF <- coef(fit1)

colnames(CI) <- c("Lower", "Higher")

table22 <- as.data.frame(cbind(HR, CI,COEF, P))
table22

#calculate the time-dependent roc (lncrna signature on entire tcga cohort)
library(timeROC)
ROC<-timeROC(T=entire_cohort$OS,# survival time

delta=entire_cohort$Censor,# results, alive or death

marker=entire_cohort$risk_lnc,# variables

cause=1,

weighting="marginal",

times=c(365*3, 365*5, 365*10),# time, survival 3,5,10 years

ROC = TRUE,

iid = TRUE)
ROC

#the time-dependent roc/auc shown above result and confidence interval shown below
#the lncrna signature performance on entire TCGA dataset
confint(ROC)$CB_AUC

#To build a risk score system based on lncRNA signature + 3 clinical factors,
#we first estimate the coef of the signature and 3 clinical factors

fit = coxph(Surv(OS, Censor) ~ stage + gender + age_at_diagnosis +

ENSG00000187185.4+
ENSG00000259641.4+
ENSG00000218510.5 +
ENSG00000257989.1+
ENSG00000206567.8
, data=study_cohort)

# summary(fit)

HR <- exp(coef(fit))
CI <- exp(confint(fit))
P <- round(coef(summary(fit))[,5],8)
COEF <- coef(fit)

colnames(CI) <- c("Lower", "Higher")

table2 <- as.data.frame(cbind(HR, CI,COEF, P))
table2

#calculate the risk score based on expression and coef of the signature and 3 clinical factors
#use coef (estimated in the study cohort) in above table
entire_cohort$risk_lncClinical = (
    entire_cohort$stage   *         table2[1,4]+
    entire_cohort$gender    *       table2[2,4]+
    entire_cohort$age_at_diagnosis   * table2[3,4] +
    entire_cohort$ENSG00000187185.4   * table2[4,4] +
    entire_cohort$ENSG00000259641.4 * table2[5,4]+
    entire_cohort$ENSG00000218510.5 * table2[6,4] +
    entire_cohort$ENSG00000257989.1 * table2[7,4] +
    entire_cohort$ENSG00000206567.8 * table2[8,4] )

#calculate the time-dependent roc (lncrna signature + 3 clinical factors on entire cohort)
library(timeROC)
ROCAll<-timeROC(T=entire_cohort$OS,# survival time

delta=entire_cohort$Censor,# results, alive or death

marker=entire_cohort$risk_lncClinical,

cause=1,# positive results

weighting="marginal",

times=c(365*3, 365*5, 365*10),# time, survival 3,5,10 years

ROC = TRUE,

iid = TRUE)
ROCAll

#the time-dependent roc/auc is shown above and confidence interval shown below
#the lncrna signature + 3 clinical factors performance on entire dataset
confint(ROCAll)$CB_AUC

title = paste(project_path,"results/tcga_time-dependent-roc", ".pdf", sep="")

# pdf(title)

#(95% CI: 0.6328-0.6835), lncrna, 5 years
#(95% CI: 0.6771-0.7596), lncrna + 3 clinical factors, 10 years

#plot the time-dependent roc
plot(ROC,time=365*3,col = "red", lty=2,add =FALSE, title="")
plot(ROC,time=365*5,col = "red",lwd=2, add =TRUE)
plot(ROC,time=365*10,col = "red",lty=3, add =TRUE)

plot(ROCAll,time=365*3,col = "blue",lty=2,add =TRUE)
plot(ROCAll,time=365*5,col = "blue",lty=3,add =TRUE)
plot(ROCAll,time=365*10,col = "blue",lwd=2,add =TRUE)
legend("bottomright", c("5-lncRNA at 3 years (AUC = 0.64)","5-lncRNA at 5 years (AUC = 0.66)","5-lncRNA at 10 years (AUC = 0.63)", "5-lncRNA + 3 clinical factors at 3 years (AUC= 0.68)","5-lncRNA + 3 clinical factors at 5 years (AUC= 0.67)", "5-lncRNA + 3 clinical factors 10 years (AUC= 0.72)"), lty=c(2,1,3,2,3,1), lwd=c(1,2,1,1,1,2), col = c("red","red","red","blue","blue", "blue"), bty="n")

# dev.off()

#calculate the median risk score (lncrna signature)
mediam_risk_score = median(as.numeric(as.vector(entire_cohort$risk_lnc)))
entire_cohort$risk_cluster <- 1
entire_cohort$risk_cluster[which(entire_cohort$risk_lnc <= mediam_risk_score)] <- 0
mediam_risk_score

#calculate the surviva different of the high- and low-risk score groups (lncrna signature)
survdiff(Surv(OS, Censor) ~risk_cluster, data = entire_cohort)

#plot the survival difference
library(survival)
require("survival")
# install.packages("svglite")
library(survminer)

dif = survdiff(Surv(OS, Censor) ~risk_cluster, data = entire_cohort)
pv = pchisq(dif$chisq, length(dif$n)-1, lower.tail = FALSE)
pv = formatC(pv, format = "e", digits = 2)
ret = c("Log-rank\r\np = ",pv)
pvalue=paste(ret[1],ret[2], sep="")

fit<- survfit(Surv(OS, Censor) ~ risk_cluster, data = entire_cohort)

p <- ggsurvplot(fit, data=entire_cohort,risk.table = TRUE,tables.theme = theme_cleantable(),palette = c("#E7B890", "#2E9FCF"), ggtheme = theme_bw(),surv.median.line = "hv",legend.title = "Risk Score", legend.labs = c("Low", "High"), xlab='Time (days)', conf.int = TRUE,pval = pvalue,pval.method=TRUE,conf.int.style='ribbon', conf.int.alpha=0.1)

p = p + theme_survminer( font.main = c(16, "bold", "darkblue"), font.submain = c(15, "bold", "purple"), font.caption = c(14, "plain", "orange"), font.x = c(14, "bold"), font.y = c(14, "bold"), font.tickslab = c(12, "plain") )

p

#save the results
title = paste(project_path, "results/RiskScore-Survival-TCGA.tiff", sep="")
tiff(title, width=480*7, height=480*7, units="px", res=96*6, compression = "lzw")
print(p, newpage = FALSE)
dev.off()

# let's see the median survival time between high- and low-risk groups
# risk_cluster=0, low-risk score group (low risk), longer survival time
#risk_cluster=1, high-risk score group (high risk), shorter survival time
print(fit)

#we need to drop patients without age_at_diagnosis or OS info
validPatients = entire_cohort[!is.na(entire_cohort$OS),]
validPatients = validPatients[!is.na(validPatients$age_at_diagnosis),]
mediam_risk_score_lncClinical = median(as.numeric(as.vector(validPatients$risk_lncClinical)))
#stratified into two groups based on median risk score
validPatients$risk_cluster1 <- 1
validPatients$risk_cluster1[which(validPatients$risk_lncClinical <= mediam_risk_score_lncClinical)] <- 0
mediam_risk_score_lncClinical

#calculate the survival difference
#risk_cluster1=0, low-risk score group (low risk), longer survival time
#risk_cluster1=1, high-risk score group (high risk), shorter survival time
fit<- survfit(Surv(OS, Censor) ~ risk_cluster1, data = validPatients)
fit

dif = survdiff(Surv(OS, Censor) ~risk_cluster1, data = validPatients)
pv = pchisq(dif$chisq, length(dif$n)-1, lower.tail = FALSE)
pv = formatC(pv, format = "e", digits = 2)
pv # p-value


#!/usr/bin/env Rscript
#	Building a statistical model for predicting sediment reworking from fraction of preserved bar structures.
#	Eric Barefoot
#	Oct 2019

# load packages

library(tidyverse)
library(here)
library(broom)
library(lme4)
if (interactive()) {
    require(colorout)
}

# load data

chamb_data = readRDS(file = here('data','derived_data', 'chamberlin_model_data.rds'))

# for helpfulness later - plot up the data with factors, etc
# factors are that the model results were generated by either:
  # (1) compensational avulsions or
  # (2) random avulsions.
# The other factors shown are sets of experiments with the same input overbank/channel material ratio ("input ratio")
# This is a strong controlling factor, and is also mapped to a frac_input continuous variable.

# chamb_data %>%
# ggplot(., aes(x = frac_elems, y = frac_supply)) +
# geom_point(aes(color = exp_set, shape = model))

# We are interested in finding the best predictive model for these data.
# Let's try out some.

# null
# fixed-effects model with
# fixed effects on % fully preserved structures

fixed_model_null = lm(frac_supply ~ frac_elems, data = chamb_data)

# fixed-effects model with
# fixed effects on % fully preserved structures
# fixed effects on input ratio

fixed_model_input = lm(frac_supply ~ frac_elems + frac_input, data = chamb_data)

# fixed-effects model with
# fixed effects on % fully preserved structures
# fixed effects on input ratio
# also interactions between two predictors.

fixed_model_input_interactions = lm(frac_supply ~ frac_elems*frac_input, data = chamb_data)

# AIC(fixed_model_null, fixed_model_input, fixed_model_input_interactions)

# mixed-effects model with
# fixed effects on % fully preserved structures
# random effects on experiment set (proxy for input ratio) mean

mixed_model_int_expSet = lmer(frac_supply ~ frac_elems + (1|exp_set), data = chamb_data)

# mixed-effects model with
# fixed effects on % fully preserved structures
# random effects on experiment set (proxy for input ratio) mean
# random effects on avulsion type mean

mixed_model_int_expSet_avulType = lmer(frac_supply ~ frac_elems + (1|exp_set) + (1|model), data = chamb_data)

# mixed-effects model with
# fixed effects on % fully preserved structures
# random effects on experiment set (proxy for input ratio) mean and slope

mixed_model_slp_expSet = lmer(frac_supply ~ frac_elems + (1 + frac_elems|exp_set), data = chamb_data, REML = FALSE, control = lmerControl(optimizer = 'Nelder_Mead'))

# mixed-effects model with
# fixed effects on % fully preserved structures
# random effects on experiment set (proxy for input ratio) mean and slope
# random effects on avulsion type mean and slope

mixed_model_slp_expSet_avulType = lmer(frac_supply ~ frac_elems + (1+frac_elems|exp_set) + (1+frac_elems|model), data = chamb_data, REML = FALSE)

# mixed-effects model with
# fixed effects on % fully preserved structures
# fixed effects on input ratio
# random effects on experiment set (proxy for input ratio) mean and slope
# random effects on avulsion type mean and slope

mixed_model_input_slp_expSet_avulType = lmer(frac_supply ~ frac_elems + frac_input + (1+frac_elems|model), data = chamb_data, REML = FALSE)

# So after doing all of those models, we need to evaluate them.

models_aic = AIC(fixed_model_null, fixed_model_input, mixed_model_int_expSet, mixed_model_int_expSet_avulType, mixed_model_slp_expSet, mixed_model_slp_expSet_avulType, mixed_model_input_slp_expSet_avulType)

# based on this, it looks like the best model is really the simplest fixed-effects model with the input ratio included.
# That effect seems to be pretty important! Since we have NO way of constraining that in the ancient, the best route forward is determine the sensitivity of the analysis to overbank vs channel deposition and work that in somehow.

# for fun, plot this model on the data.

# chamb_data_pred = chamb_data %>% mutate(pred_supply = predict(fixed_model_input_interactions))
#
# ggplot(data = chamb_data_pred) +
# geom_point(aes(x = frac_elems, y = frac_supply, color = frac_input, shape = model)) +
# geom_point(aes(x = frac_elems, y = pred_supply), color = 'red3')

# export the best model as a derived data product.

saveRDS(fixed_model_input_interactions, file = here('data','derived_data', 'chamberlin_stat_model.rds'))


#' Choose initial parameters for direct effects on X and Y
#'
#' @param nsnp_x nsnp_x
#' @param var_gx.x var_gx.x
#' @param var_x.y var_x.y
#' @param var_gx.y var_gx.y=0
#' @param nsnp_y nsnp_y=0
#' @param var_gy.y var_gy.y=0
#' @param mu_gx.y mu_gx.y=0
#' @param var_gy.x var_gy.x=0
#' @param mu_gy.x mu_gy.x=0
#' @param prop_gy.x prop_gy.x=1
#' @param prop_gx.y prop_gx.y=1
#'
#' @export
#' @return List of model parameters
init_parameters <- function(nsnp_x, var_gx.x, var_x.y, var_gx.y=0, nsnp_y=0, var_gy.y=0, mu_gx.y=0, var_gy.x=0, mu_gy.x=0, prop_gy.x=1, prop_gx.y=1)
{
	parameters <- list(
		# Causal effect
		var_x.y = var_x.y,

		# Direct effects on x
		nsnp_x = nsnp_x,
		var_gx.x = var_gx.x,
		var_gx.y = var_gx.y,
		mu_gx.y = mu_gx.y,
		prop_gx.y = prop_gx.y,

		nsnp_y = nsnp_y,
		var_gy.y = var_gy.y,
		var_gy.x = var_gy.x,
		mu_gy.x = mu_gy.x,
		prop_gy.x = prop_gy.x,
		u = list()
	)
	return(parameters)
}

#' Add confounder variables and their instruments
#'
#' @param parameters Output from \code{init_parameters}
#' @param nsnp_u nsnp_u
#' @param var_u.x var_u.x
#' @param var_u.y var_u.y
#' @param var_gu.u var_gu.u
#'
#' @export
#' @return List of model parameters
add_u <- function(parameters, nsnp_u, var_u.x, var_u.y, var_gu.u)
{
	nom <- "u"
	if(var_u.x == 0 & var_u.y != 0) nom <- "b"
	if(var_u.x != 0 & var_u.y == 0) nom <- "a"
	i <- length(parameters$u) + 1
	parameters$u[[i]] <- list(
		nsnp_u = nsnp_u,
		var_u.x = var_u.x,
		var_u.y = var_u.y,
		var_gu.u = var_gu.u,
		name_u = nom
	)
	return(parameters)
}

#' Sample the actual effects based on initial parameters
#'
#' @param parameters Output from \code{init_parameters} or \code{add_u}
#'
#' @export
#' @return List of effect parameters
sample_system_effects <- function(parameters)
{
	nu <- length(parameters$u)
	if(nu > 0)
	{
		for(i in 1:nu)
		{
			parameters$u[[i]]$eff_gu.u <- choose_effects(parameters$u[[i]]$nsnp_u, parameters$u[[i]]$var_gu.u)
			parameters$u[[i]]$eff_u.x <- choose_effects(1, parameters$u[[i]]$var_u.x)
			parameters$u[[i]]$eff_u.y <- choose_effects(1, parameters$u[[i]]$var_u.y)
		}
	}

	# Make genetic effects for x instruments
	parameters$eff_gx.x <- choose_effects(parameters$nsnp_x, parameters$var_gx.x)

	# Make genetic effects for y instruments
	parameters$eff_gy.y <- choose_effects(parameters$nsnp_y, parameters$var_gy.y)

	# Create pleiotropic effects for some proportion of the effects
	# X-Y
	nchoose <- round(parameters$nsnp_x * parameters$prop_gx.y)
	ind <- sample(1:parameters$nsnp_x, nchoose, replace=FALSE)
	parameters$eff_gx.y <- rep(0, parameters$nsnp_x)
	if(nchoose > 0)
	{
		parameters$eff_gx.y[ind] <- choose_effects(nchoose, parameters$var_gx.y, mua=parameters$mu_gx.y)
	}

	# Create pleiotropic effects for some proportion of the effects
	# Y-X
	nchoose <- round(parameters$nsnp_y * parameters$prop_gy.x)
	ind <- sample(1:parameters$nsnp_y, nchoose, replace=FALSE)
	parameters$eff_gy.x <- rep(0, parameters$nsnp_y)
	if(nchoose > 0)
	{
		parameters$eff_gy.x[ind] <- choose_effects(nchoose, parameters$var_gy.x, mua=parameters$mu_gy.x)
	}

	parameters$eff_u.x <- sapply(parameters$u, function(x) choose_effects(1, x$var_u.x))
	parameters$eff_u.y <- sapply(parameters$u, function(x) choose_effects(1, x$var_u.y))
	parameters$eff_x.y <- sqrt(parameters$var_x.y)

	return(parameters)
}

#' Simulate individual level data from initial parameters
#'
#' @param parameters Output from \code{init_parameters} or \code{add_u}
#' @param nid Sample size to generate
#'
#' @export
#' @return List of matrices and vectors that represent individual level data
simulate_population <- function(parameters, nid)
{
	Gx <- matrix(rbinom(parameters$nsnp_x * nid, 2, 0.5), nid, parameters$nsnp_x)

	Gy <- matrix(rbinom(parameters$nsnp_y * nid, 2, 0.5), nid, parameters$nsnp_y)

	U <- lapply(parameters$u, function(param)
	{
		G <- matrix(rbinom(param$nsnp_u * nid, 2, 0.5), nid, param$nsnp_u)
		u <- make_phen(param$eff_gu.u, G)
		return(list(p=u, G=G, nom=param$name_u))
	})

	bx <- parameters$eff_gx.x
	by <- parameters$eff_gx.y

	Gt <- Gx
	if(parameters$nsnp_y > 0)
	{
		bx <- c(bx, parameters$eff_gy.x)
		by <- c(by, parameters$eff_gy.y)
		Gt <- cbind(Gt, Gy)
	}

	if(length(parameters$u) > 0)
	{
		bx <- c(bx, sapply(parameters$u, function(x) x$eff_u.x))
		by <- c(by, sapply(parameters$u, function(x) x$eff_u.y))
		Gt <- cbind(Gt, do.call(cbind, lapply(U, function(x) x$p)))
	}
	by <- c(by, parameters$eff_x.y)

	x <- make_phen(bx, Gt)
	y <- make_phen(by, cbind(Gt, x))

	return(list(
		y=y,
		x=x,
		Gx=Gx,
		Gy=Gy,
		U=U
	))

}

#' Estimate the effects of all SNPs on all phenotypes
#'
#' @param sim Output from \code{simulate_population}
#'
#' @export
#' @return Lists of SNP-trait effect estimates
estimate_system_effects <- function(sim)
{
	# Get effects of all X SNPs

	l <- list()

	gx.x <- gwas(sim$x, sim$Gx)
	gx.x$inst <- "x"
	gx.x$snp <- paste0(1:nrow(gx.x), "x")

	gx.y <- gwas(sim$y, sim$Gx)
	gx.y$inst <- "x"
	gx.y$snp <- paste0(1:nrow(gx.y), "x")

	gx <- gx.x
	gy <- gx.y

	# Get effects of all SNPs on Y

	if(ncol(sim$Gy) > 0)
	{
		gy.x <- gwas(sim$x, sim$Gy)
		gy.x$inst <- "y"
		gy.x$snp <- paste0(1:nrow(gy.x), "y")
		gx <- rbind(gx, gy.x)

		gy.y <- gwas(sim$y, sim$Gy)
		gy.y$inst <- "y"
		gy.y$snp <- paste0(1:nrow(gy.y), "y")
		gy <- rbind(gy, gy.y)
	}

	nconf <- length(sim$U)
	if(nconf > 0)
	{
		gu.x <- list()
		gu.y <- list()
		gx.u <- list()
		gy.u <- list()
		gu.u <- list()
		for(i in 1:nconf)
		{
			gu.u[[i]] <- gwas(sim$U[[i]]$p, sim$U[[i]]$G)
			gu.u[[i]]$inst <- paste0("u", i)
			gu.u[[i]]$snp <- paste0(1:nrow(gu.u[[i]]), gu.u[[i]]$inst)

			gu.x[[i]] <- gwas(sim$x, sim$U[[i]]$G)
			gu.x[[i]]$inst <- paste0("u", i)
			gu.x[[i]]$snp <- paste0(1:nrow(gu.x[[i]]), gu.x[[i]]$inst)

			gx.u[[i]] <- gwas(sim$U[[i]]$p, sim$Gx)
			gx.u[[i]]$inst <- "x"
			gx.u[[i]]$snp <- paste0(1:nrow(gx.u[[i]]), gx.u[[i]]$inst)

			gu.y[[i]] <- gwas(sim$y, sim$U[[i]]$G)
			gu.y[[i]]$inst <- paste0("u", i)
			gu.y[[i]]$snp <- paste0(1:nrow(gu.y[[i]]), gu.y[[i]]$inst)

			if(ncol(sim$Gy) > 0)
			{
				gy.u[[i]] <- gwas(sim$U[[i]]$p, sim$Gy)
				gy.u[[i]]$inst <- "y"
				gy.u[[i]]$snp <- paste0(1:nrow(gy.u[[i]]), gy.u[[i]]$inst)
			}
		}
		gx <- rbind(gx, dplyr::bind_rows(gu.x))
		gy <- rbind(gy, dplyr::bind_rows(gu.y))
		gu <- list()
		for(i in 1:nconf)
		{
			if(ncol(sim$Gy) > 0)
			{
				gu[[i]] <- rbind(gx.u[[i]], gy.u[[i]], gu.u[[i]])
			} else {
				gu[[i]] <- rbind(gx.u[[i]], gu.u[[i]])
			}
		}
		names(gu) <- paste0("u", 1:nconf)
		l <- gu
	}
	l$x <- gx
	l$y <- gy
	return(l)
}

#' Wrapper for simulation pipeline
#'
#' Based on the parameters specified this function will call \code{init_parameters}, \code{add_u}, \code{sample_system_effects}, \code{simulate_population} and \code{estimate_system_effects}. A separate population is generated for each phenotype (x, y and each of u) to allow 2SMR and PRS analyses
#'
#' @param nidx number of individuals for x
#' @param nidy number of individuals for y
#' @param nidu Default=0 If 0 then don't simulate separate populations for the u variables
#' @param nu Default=0 Number of confounders influencing x and y
#' @param na Default=0 Number of traits upstream of x
#' @param nb Default=0 Number of traits upstream of y
#' @param var_x.y var in y explained by x
#' @param nsnp_x nsnp influencing x
#' @param var_gx.x variance in x explained by instruments for x
#' @param var_gx.y Default=0
#' @param mu_gx.y Default=0
#' @param prop_gx.y Default=1
#' @param nsnp_y Default=0
#' @param var_gy.y Default=0
#' @param var_gy.x Default=0
#' @param mu_gy.x Default=0
#' @param prop_gy.x Default=1
#'
#' @export
#' @return List of effects across system
create_system <- function(nidx, nidy, nidu=0, nu=0, na=0, nb=0, var_x.y, nsnp_x, var_gx.x, var_gx.y=0, mu_gx.y=0, prop_gx.y=1, nsnp_y=0, var_gy.y=0, var_gy.x=0, mu_gy.x=0, prop_gy.x=1)
{
	parameters <- init_parameters(nsnp_x=nsnp_x, nsnp_y=nsnp_y, var_gx.x=var_gx.x, var_gy.y=var_gy.y, var_x.y=var_x.y, var_gx.y=var_gx.y, mu_gx.y=mu_gx.y, prop_gx.y=prop_gx.y, var_gy.x=var_gy.x, mu_gy.x=mu_gy.x, prop_gy.x=prop_gy.x)
	if(nu > 0)
	{
		for(i in 1:nu)
		{
			parameters <- add_u(
				parameters,
				nsnp_u=sample(5:30, 1),
				var_u.x=stats::runif(1, min=0.01, max=0.1),
				var_u.y=stats::runif(1, min=0.01, max=0.1),
				var_gu.u=stats::runif(1, min=0.02, 0.2)
			)
		}
	}
	if(na > 0)
	{
		for(i in 1:na)
		{
			parameters <- add_u(
				parameters,
				nsnp_u=sample(5:30, 1),
				var_u.x=stats::runif(1, min=0.01, max=0.1),
				var_u.y=0,
				var_gu.u=stats::runif(1, min=0.02, 0.2)
			)
		}
	}
	if(nb > 0)
	{
		for(i in 1:nb)
		{
			parameters <- add_u(
				parameters,
				nsnp_u=sample(5:30, 1),
				var_u.x=0,
				var_u.y=stats::runif(1, min=0.01, max=0.1),
				var_gu.u=stats::runif(1, min=0.02, 0.2)
			)
		}
	}

	parameters <- sample_system_effects(parameters)

	message("X")
	pop <- simulate_population(parameters, nidx)
	x <- estimate_system_effects(pop)

	message("Y")
	pop <- simulate_population(parameters, nidy)
	y <- estimate_system_effects(pop)

	out <- list(x=x, y=y, parameters=parameters)

	if(nidu > 0)
	{
		u <- list()
		if(nu > 0)
		{
			for(i in 1:nu)
			{
				message("U: ", i, " of ", nu)
				pop <- simulate_population(parameters, nidu)
				u[[i]] <- estimate_system_effects(pop)
			}
		}

		a <- list()
		if(na > 0)
		{
			for(i in 1:na)
			{
				message("A: ", i, " of ", na)
				pop <- simulate_population(parameters, nidu)
				a[[i]] <- estimate_system_effects(pop)
			}
		}

		b <- list()
		if(nb > 0)
		{
			for(i in 1:nb)
			{
				message("B: ", i, " of ", nb)
				pop <- simulate_population(parameters, nidu)
				b[[i]] <- estimate_system_effects(pop)
			}
		}
		u <- c(u, a, b)
		names(u) <- paste0("u", 1:length(u))
		out$u <- u
	}
	return(out)
}

#' Apply MR tests to system
#'
#'
#' @param ss Output from create_syste,
#' @param id string denoting simulation ID
#'
#' @export
#' @return List
test_system <- function(ss, id="test")
{
	out <- list()

	dx <- merge_exp_out(ss$x$x, ss$y$y)
	dx$exposure <- paste0("X:", id)
	dx$outcome <- paste0("Y:", id)
	dx$id.exposure <- paste0("X:", id)
	dx$id.outcome <- paste0("Y:", id)

	dy <- merge_exp_out(ss$y$y, ss$x$x)
	dy$exposure <- paste0("Y:", id)
	dy$outcome <- paste0("X:", id)
	dy$id.exposure <- paste0("Y:", id)
	dy$id.outcome <- paste0("X:", id)

	# Oracle
	ox <- subset(dx, grepl("x", SNP))
	out$ox <- try(TwoSampleMR::mr_wrapper(ox)[[1]])
	oy <- subset(dy, grepl("y", SNP))
	out$oy <- try(TwoSampleMR::mr_wrapper(oy)[[1]])

	# Empirical
	ex <- subset(dx, pval.exposure < 5e-8)
	out$ex <- try(TwoSampleMR::mr_wrapper(ex)[[1]])
	ey <- subset(dy, pval.exposure < 5e-8)
	out$ey <- try(TwoSampleMR::mr_wrapper(ey)[[1]])

	param <- expand.grid(
		hypothesis = c("x", "y"),
		selection = c("e", "o"),
		type = c("valid", "reverse", "confounder")
	)
	o <- list()
	for(i in 1:nrow(param))
	{
		o[[i]] <- dplyr::tibble(
			hypothesis = param$hypothesis[i],
			selection = param$selection[i],
			type = param$type[i],
			measure = c("nofilter", "outlier", "steiger", "both"),
			counts = get_counts(
				param$type[i],
				param$hypothesis[i],
				get(paste0(param$selection[i], param$hypothesis[i])),
				out[[paste0(param$selection[i], param$hypothesis[i])]]
			)
		)
	}

	out$instrument_validity <- dplyr::bind_rows(o) %>% dplyr::group_by(hypothesis, selection, type) %>% dplyr::mutate(n=dplyr::first(counts), pcounts=counts/n)
	out$instrument_validity$id <- id

	# Best model
	out$ex$estimates <- best_model(out$ex$estimates, ss$parameters$eff_x.y)
	out$ox$estimates <- best_model(out$ox$estimates, 0)

	return(out)
}

get_counts <- function(type, hyp, dat, res)
{
	if(type == "valid")
	{
		node = hyp
	} else if(type == "reverse") {
		node = ifelse(hyp=="x", "y", "x")
	} else {
		node = "u"
	}
	c(
		nrow(res$snps_retained %>% subset(grepl(node, SNP))),
		nrow(subset(res$snps_retained, outlier) %>% subset(grepl(node, SNP))),
		nrow(subset(res$snps_retained, steiger) %>% subset(grepl(node, SNP))),
		nrow(subset(res$snps_retained, both) %>% subset(grepl(node, SNP)))
	)
}

best_model <- function(res, bxy)
{
	res$beta_correct <- res$ci_low <= bxy & res$ci_upp >= bxy
	res$beta_best <- FALSE
	res$beta_best[which.min(abs(res$b - bxy))] <- TRUE
	res$pval_sig <- res$pval < 0.05
	res$pval_lowest <- FALSE
	res$pval_lowest[which.min(res$pval)] <- TRUE
	res$pval_highest <- FALSE
	res$pval_highest[which.max(res$pval)] <- TRUE
	return(res)
}


###########################################################################################
### Table 3 simulation (linear)
###########################################################################################
eachCase <- function(case,eff=0)
    sapply(c('sh','ik','cft'),
           function(method) c(bias=case[paste0(method,'.est')]-eff,
                              cover=case[paste0(method,'.CIl')]<=eff &
                                  case[paste0(method,'.CIh')]>=eff,
                              width=case[paste0(method,'.CIh')]-
                                  case[paste0(method,'.CIl')])
        )

levTabCI <- function(res,eff=0){
    out <- NULL
     for(n in c(50,250,2500))
         for(err in c('norm','t')){
             run <- res[[paste(n,eff,err,sep='_')]]
            row <- NULL
            for(meth in c('cft','sh','ik'))
                row <- c(row,
                         if(is.list(run))
                             mean(sapply(run,function(x) x[paste0(meth,'.p')])<0.05,na.rm=TRUE)
                         else mean(run[paste0(meth,'.p'),]<0.05))

             out <- rbind(out,row)
             rownames(out)[nrow(out)] <- paste(err,n)
         }
    colnames(out) <- c('cft','sh','ik')
    out
}

displayCIsimHet <- function(res,tau=0,caption='',label=''){
  ## omit cases where one of the methods (ours?) didn't converge, gave NA
  res <- lapply(res,function(x) t(x[,apply(x,2,function(cc) !any(is.na(cc)))]))

  res <- res[grep(paste0('_tau',tau,'W'),names(res),fixed=TRUE)]
  res <- res[grepl('Wnone|Wt',names(res))]

  tab0 <- lapply(res,function(x)
    apply(sapply(1:nrow(x),function(i) eachCase(x[i,],eff=tau),simplify='array'),
      1:2,mean))

  tab0 <- lapply(tab0,function(xx){
    out <- rbind(sprintf("%.2f",round(xx[1,],2)),sprintf("%i",as.integer(round(xx[2,]*100))),sprintf("%.2f",round(xx[3,],2)))
    dimnames(out) <- dimnames(xx)
    out})

  cat('
\\begin{table}
\\footnotesize
\\begin{tabular}{ccc|ccc|ccc|ccc}
\\hline

&&& \\multicolumn{ 3 }{c}{Permutation}&\\multicolumn{ 3 }{c}{\`\`Limitless\'\'}&\\multicolumn{ 3 }{c}{Local OLS}\\\\
$n$& Effect& Error &', paste(rep(c('Bias','Cover.','Width'),3),collapse='&'),'\\\\
\\hline \n')
  for(n in c(50,250,2500))
#    for(err in c('norm','t')){
    for(rr in c('n0','t0','tt')){
      row <- NULL
      err <- ifelse(rr=='n0','norm','t')
      eff <- ifelse(rr=='tt','$t_3$',0)
      for(meth in c('cft','sh','ik'))
        row <- c(row,tab0[[paste0(n,'_err',err,'_tau',tau,'W',ifelse(rr=='tt','t','none'))]][,meth])
      if(err=='norm'){
        cat('\\hline \n')
        cat('\\multirow{3}{*}{',n,'} &0& $\\mathcal{N}(0,1)$ &')
      } else cat(' & ',eff,'& $t_3$ &')
      cat(paste(row,collapse='&'),'\\\\ \n')
    }
  cat('\\hline
\\end{tabular}
  \\caption{',caption,'}
  \\label{',label,'}\n',sep='')
  cat('\\end{table}\n')
}

dispAllSimp <- function(res){
  res <- sapply(res,function(x) t(x[,apply(x,2,function(cc) !any(is.na(cc)))]),simplify=FALSE)
  tab0 <- lapply(names(res),function(nn){
    x <- res[[nn]]
    tau <- ifelse(grepl('tau0.2',nn,fixed=TRUE),0.2,0)
    apply(sapply(1:nrow(x),function(i) eachCase(x[i,],eff=tau),simplify='array'),
      1:2,mean)
  })

  cond <- strsplit(names(res),'_')
  cond <- lapply(cond,function(x) gsub('err|tau','',x))
  cond <- lapply(cond,function(x) c(x[1],x[2],strsplit(x[3],'W')[[1]]))
  cond <- do.call('rbind',cond)
  out <- data.frame(tab0[[1]],stringsAsFactors=FALSE)
  out <- cbind(cond[rep(1,3),],c('Bias','Coverage','Width'),out)

  for(i in 2:length(tab0))
    out <- rbind(out,cbind(cond[rep(i,3),],
      c('Bias','Coverage','Width'),
      data.frame(tab0[[i]],stringsAsFactors=FALSE)))

  names(out) <- c('n','error','ATE','TE randomness','meas','Limitless','Local-OLS','Permutation')

  out
}

###########################################################################################
### Table 4 simulation (polynomial)
###########################################################################################

#' @export
## dgms <- function(tp){

##     ### the DGMs
##     curve(0.5*x,from=-1,to=1,ylim=c(-2,2),xlab='r',ylab='$\\EE[Y|R=r]$',main='Linear')
##     abline(v=0,lty=2)
##     curve(ifelse(abs(x)>0.5, 3*x+sign(x)*(0.5-3)*0.5,0.5*x),from=-1,to=1,ylim=c(-2,2),xlab='r',ylab='$\\EE[Y|R=r]$',main='Anti-Symmetric')
##     abline(v=0,lty=2)
##   #curve(ifelse(x>0.5,3*x+(0.5-3)*0.5,0.5*x),from=-1,to=1,ylim=c(-2,2),xlab='r',ylab='$\\EE[Y|R=r]$',main='One-Sided')
##   curve(mu4,from=-1,to=1,ylim=c(-2,2),xlab='r',ylab='$\\EE[Y|R=r]$',main='One-Sided')
##     abline(v=0,lty=2)
## }

dgms <- function(){
  lin <- function(x) 0.5*x
  as <- function(x) ifelse(abs(x)>0.5,3*x+sign(x)*(0.5-3)*0.5,lin(x))
  mu4 <- function(x) sin(3*x)
  p <- ggplot(data = data.frame(x = 0), mapping = aes(x = x))+xlim(-1,1)+
    theme(axis.line=element_blank(),
      axis.text.x=element_blank(),
      axis.text.y=element_blank(),
      axis.ticks=element_blank(),
      #axis.title.x='$r$',
      #axis.title.y='$\\EE[Y|R=r]$',
      legend.position="none",
      ## panel.background=element_blank(),
      ## panel.border=element_blank(),
      ## panel.grid.major=element_blank(),
      ## panel.grid.minor=element_blank(),
      ## plot.background=element_blank()
    )+xlab('$r$')+ylab('$\\EE[Y|R=r]$')

  gridExtra::grid.arrange(p+stat_function(fun=lin)+ggtitle('Linear'),
    p+stat_function(fun=as)+ggtitle('Anti-Symmetric'),
    p+stat_function(fun=mu4)+ggtitle('Sine'),nrow=1)
}

resTab <- function(run,full=FALSE){
    llr <- run[c(nrow(run)-1,nrow(run)),]
    run <- run[-c(nrow(run)-1,nrow(run)),]

    pvals <- run[seq(1,nrow(run)-1,2),]
    ests <- run[seq(2,nrow(run),2),]

    runPsh <- pvals[seq(1,nrow(pvals),2),]
    runPik <- pvals[seq(2,nrow(pvals),2),]

    runEstsh <- ests[seq(1,nrow(ests),2),]
    runEstik <- ests[seq(2,nrow(ests),2),]

    tabFun <- function(runP,runEst,full){
        tab <- rbind(
            level=apply(runP,1,function(x) mean(x<0.05,na.rm=TRUE)),
            RMSE=apply(runEst,1,function(x) sqrt(mean(x^2,na.rm=TRUE))))
        if(full) tab <- rbind(tab,
                              bias=rowMeans(runEst,na.rm=TRUE),
                              sd=apply(runEst,1,sd,na.rm=TRUE))
        tab
    }
    list(tabSH=tabFun(runPsh,runEstsh,full=full),
         tabIK=tabFun(runPik,runEstik,full=full),
         tabLLR=tabFun(rbind(llr[1,]),rbind(llr[2,]),full=full))
}

resTab <- function(run,full=FALSE){
  run <- run[,apply(run,2,function(x) !any(is.na(x)))]
  tabFun <- function(ests)
      rbind(
        bias=rowMeans(rbind(ests)),
        RMSE=sqrt(rowMeans(rbind(ests^2))))

  tabFunP <- function(ps)
    apply(rbind(ps),1,function(x) mean(x<0.05,na.rm=TRUE))

  out <- sapply(c('sh','ik','ll'),function(mm) tabFun(run[grep(paste0(mm,'.est'),rownames(run)),]),
    simplify=FALSE)

  if(full) out <- sapply(c('sh','ik','ll'),
    function(mm) rbind(out[[mm]],level= tabFunP(run[grep(paste0(mm,'.p'),rownames(run)),])),
    simplify=FALSE)

  out
}

#' @export
prntTab <- function(totalPoly,maxDeg=4,full=TRUE,md=FALSE){
    tab <- NULL

    if(maxDeg>(nrow(totalPoly[[1]])-2)/4){
        maxDeg <- (nrow(totalPoly[[1]])-2)/4
        if(!full)
            warning(paste('There don\'t seem to be that many degrees in the simulation.\n Setting maxDeg=',maxDeg))
    }

    ctab <- function(runname,full){
        run <- totalPoly[[runname]]
        res <- resTab(run,full=full)

        cbind(res[['sh']][,1:maxDeg],
              res[['ik']][,1:maxDeg],
              res[['ll']])
    }
    for(dgm in c('lin','antiSym','wass')){
        tab <- rbind(tab,ctab(paste0(dgm,'_t'),full))
        if(full) if(paste0(dgm,'_norm')%in%names(totalPoly)) tab <- rbind(tab,ctab(paste0(dgm,'_norm'),full))
    }
    if(md){
        colnames(tab) <- c(paste('LRD, deg=',1:maxDeg),paste('OLS, deg=',1:maxDeg),'Loc.Lin')
        rownames(tab) <- paste(rep(c('lin','antiSym','sine'),
          each=nrow(tab)/sum(rownames(tab)=='bias')),#,times=2),
          #rep(c('t err','norm err'),each=nrow(tab)/2),
          rownames(tab))
    }
    return(tab)
}

#' @export
polyLatex <- function(tab,full,caption='',label='tab:poly'){
    if(NCOL(tab)!=9) stop('This only works with polynomial degree=1,...,4')
    cat('
        \\begin{table}[ht]
\\centering
\\begin{tabular}{cr|llll|llll|l',ifelse(full,'|llll|llll|l}','}'),'
  \\hline \n')
    if(full) cat('&&\\multicolumn{9}{c|}{$t_3$ Errors} &\\multicolumn{9}{c|}{$\\mathcal{N}(0,1)$ Errors} \\\\ \n')

 cat('&& \\multicolumn{4}{c|}{Limitless} &  \\multicolumn{4}{c|}{OLS} &\\makecell[c]{Local\\\\Linear}',
        ifelse(full,'\\multicolumn{4}{c|}{Limitless} &  \\multicolumn{4}{c|}{OLS} &\\makecell[c]{Local\\\\Linear}',''),'\\\\
 \\multicolumn{2}{r|}{\\makecell[r]{Polynomial\\\\Degree}}&1&2&3&4&1&2&3&4&',ifelse(full,'&1&2&3&4&1&2&3&4&n/a','n/a'),' \\\\
')
    for(rr in 1:nrow(tab)){
        if(rr==1) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}Linear\\end{sideways}','Linear'),'}')
        if(rr==3) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}Anti-Sym\\end{sideways}','Anti-Sym'),'}')
        if(rr==5) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}One-Side\\end{sideways}','One-Side'),'}')
        cat('&',rownames(tab)[rr],'&')
        cat(paste(sprintf("%.1f", round(tab[rr,],1)),collapse='&'))
        cat('\\\\ \n')
    }
    cat('
 \\hline
\\end{tabular}
\\caption{',caption,'}
\\label{',label,'}
\\end{table}\n',sep='')

}

#' @export
polyLatex5 <- function(tab,full,caption='',label='tab:poly'){
  if(NCOL(tab)!=11) stop('This only works with polynomial degree=1,...,5')
  tab2 <- tab
  for(i in 1:nrow(tab)) for(j in 1:ncol(tab))
                          tab2[i,j] <- ifelse(tab[i,j]>10,
                            sprintf("%i",as.integer(round(tab[i,j]))),
                            sprintf("%.1f",round(tab[i,j],1)))

    cat('
        \\begin{table}[ht]
\\centering
\\begin{tabular}{ll|ccccc|ccccc|c',ifelse(full,'|llll|llll|l}','}'),'
  \\hline \n')
    if(full) cat('&&\\multicolumn{11}{c|}{$t_3$ Errors} &\\multicolumn{11}{c|}{$\\mathcal{N}(0,1)$ Errors} \\\\ \n')

 cat('&& \\multicolumn{5}{c|}{Limitless} &  \\multicolumn{5}{c|}{OLS} &Local',
        ifelse(full,'\\multicolumn{5}{c|}{Limitless} &  \\multicolumn{5}{c|}{OLS} &Local',''),'\\\\
 && \\multicolumn{5}{c|}{Polynomial Degree}&\\multicolumn{5}{c|}{Polynomial Degree}&Linear',
 ifelse(full,'\\multicolumn{5}{c|}{Polynomial Degree}&\\multicolumn{5}{c|}{Polynomial Degree}&Linear',''),'\\\\
 DGM&Measure&1&2&3&4&5&1&2&3&4&5&',ifelse(full,'&1&2&3&4&1&2&3&4&5&',''),' \\\\
')
    for(rr in 1:nrow(tab)){
        if(rr==1) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}Linear\\end{sideways}','Linear'),'}',sep='')
        if(rr==3) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}Anti-Sym\\end{sideways}','\\makecell[c]{Anti-\\\\Sym}'),'}',sep='')
        if(rr==5) cat('\\hline\n\\hline\n\\multirow{',ifelse(full,4,2),'}{*}{',ifelse(full,'\\begin{sideways}One-Side\\end{sideways}','Sine'),'}',sep='')
        cat('&',rownames(tab)[rr],'&')
        cat(paste(tab2[rr,],collapse='&'))
        cat('\\\\ \n')
    }
    cat('
 \\hline
\\end{tabular}
\\caption{',caption,'}
\\label{',label,'}
\\end{table}\n',sep='')

}


=begin
= Changes
== 0.71(2005.07.26)
* MPolynomial#derivate
* MPolynomial#map_to

== 0.70(2005.02.02)
* id => object_id

== 0.69(2004.10.15)
* MRationalFunctionField
* MPolynomial#gcd
* SquareMatrix#inverse depends on ufd?

== 0.68(2004.07.28)
* fix Mpolynomial#to_s (for tex)

== 0.67(2004.06.10)
* need_paren_in_coeff?
* Matrix: entries are regulated.
* Matrix: type*type is due to wedge.
* AlgebraicSystem: wedge, superior?
* Polynomial#sub, MPolynomial#sub
* SquareMatrix.determinant
* fix MPolynomial#to_s
* fix Factors#normalize!
* (Co)Vector#norm2, inner_product

== 0.66(2004.06.05)
* fix _hensel_lift (add param 'where')
* sample-geometry01.rb
* title of docs

== 0.65(2004.06.04)
* Polynomial#factor over Polynomial
* reverse var's order in value_on & convert_to in polynomial-convert.rb

== 0.64(2004.03.01)
* return value of SquareMatrix#diagonalize is Strut

== 0.63(2004.01.08)
* fix matrix[i, j] = x

== 0.62(2003.11.13)
* allow f.gcd_coeff_all()

== 0.61(2003.11.06)
* corrected Enumerable#any? in finite-set.rb

== 0.60(2003.07.28)
* AlgebraMatrix(cofactor, cofactor_matrix, adjoint)

== 0.59(2003.07.28)
* adapted to the version 1.8.0(preview 4)
* include installer

== 0.58(2002.04.00)
* AlgebraicExtensionField#[n] is lift[n]
* abolish the field of 'poly_exps' of the splitting field' and add 'proots'
* MatrixAlgebra#rank
* Raise Exception for the inverse of the non invertible square matrix
* Use "import-module" for monomial order (0.57.18.01)

== 0.57 (2002.03.16)
* New class
  * Set
  * Map
  * Group
  * PermuationGroup
  * AlgebraicExtensionField
* New methods
  * Galois#group
  * Polynomial#splitting_field
* Rename MinimalDecompositionField to decompose
* Move *.rb to lib/
* Define Polynomial.to_ary and MPolynomial.to_ary

== 0.56 (2002.02.06)
* MPolynomial#with_ord
* Orderings belong to MPolynomial.
* MIndex is not a subclass of Array, now.
* "Integer < Numeric" ==> "Integer"

== 0.55 (2001.11.15)
* Improve Algorithm of Factorization over algebraic fields

== 0.54 (2001.11.05)
* Elementary Divisor
* Jordan Form
* Minimal Decomposition Field

== 0.53 (2001.10.03)

* move "MatrixAlgebra" to "Algebra::MatrixAlgebra"
* move "MatrixAlgebra::Vector" to "Algebra::Vector"
* move "MatrixAlgebra::Covector" to "Algebra::Covector",
* move "MatrixAlgebra::SquareMatrix" to "Algebra::SquareMatrix"

== 0.52 (2001.09.22)

* One-variate polynomial
  * Fundamental operations (addition, multiplication, quotient/remainder, ...)
  * factorization
* Multi-variate polynomial
  * Fundamental operations (addition, multiplication, ...)
  * factorization
  * Creating Groebner-basis, quotient/remainder by Groebner-basis.
* Algebraic systems
  * Creating quotient fields
  * Creating residue class fields
  * Operating matrices.
* Linear Algebra
  * Diagonalization of Symmetrix Matrix

== 0.40 (2001.03.03)

* Initial Version

=end


#' ---
#' title: "Prior probabilities in the interpretation of 'some': analysis of uniform prior wonky world model predictions"
#' author: "Judith Degen"
#' date: "January 26, 2014"
#' ---

library(ggplot2)
theme_set(theme_bw(18))
setwd("/Users/titlis/cogsci/conferences_talks/_2015/5_berkeley/")
source("rscripts/helpers.r")

# get prior expectations
priorexpectations = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/12_sinking-marbles-prior15/results/data/expectations.txt",sep="\t", header=T, quote="")
row.names(priorexpectations) = paste(priorexpectations$effect, priorexpectations$object)

# histogram of expectations
exps = ggplot(priorexpectations, aes(x=expectation_corr)) +
  geom_histogram() +
  scale_x_continuous(name="Expected value of prior distribution",breaks=seq(1,15, by=2)) +
  scale_y_continuous(name="Number of cases",breaks=seq(0,8, by=2))
ggsave("priorexpectations-histogram.pdf",width=5,height=3.7)

# get prior allstate-probs
priorprobs = read.table(file="~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/12_sinking-marbles-prior15/results/data/smoothed_15marbles_priors_withnames.txt",sep="\t", header=T, quote="")
row.names(priorprobs) = priorprobs$Item
head(priorprobs)

# histogram of allstate-probs
allprobs = ggplot(priorprobs, aes(x=X15)) +
  geom_histogram() +
  scale_x_continuous(name="Prior all-state probability") +
  scale_y_continuous(name="Number of cases")
ggsave("priorallprobs-histogram.pdf")

pdf("priordistributions.pdf",width=10,height=4)
grid.arrange(exps, allprobs, nrow=1,widths=unit.c(unit(.5, "npc"), unit(.5, "npc")))
dev.off()

#####################################
# plot model predictions: expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

library(dplyr)
library(plyr)
# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5))
nrow(toplot)
pexpectations = ddply(toplot, .(Item), summarise, PosteriorExpectation_predicted=sum(State*PosteriorProbability)/15)
head(pexpectations)
some = pexpectations#droplevels(subset(pexpectations, Quantifier == "some"))

toplot = some
nrow(toplot)
head(toplot)
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode

ggplot(toplot, aes(x=PriorExpectation_smoothed, y=PosteriorExpectation_predicted)) +
  geom_point(color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth(color="#00B0F6") +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior expectation") +
  scale_y_continuous(limits=c(0,1), name="Model predicted posterior expectation")

load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5 & State == 15))
nrow(toplot)

# adjust speaker optimality at will
toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)

toplot_w = toplot
# get rRSA predictions for qud=how-many, alts=0_basic, spopt=2
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/complex_prior/smoothed_unbinned15/results/data/mp.RData")
summary(mp)
toplot = droplevels(subset(mp, QUD == "how-many" & Alternatives == "0_basic" &  State == 15))
nrow(toplot)

# adjust speaker optimality at will
toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)

toplot_r = toplot
toplot_r$Model = "RSA"
toplot_w$Model = "wRSA"

# plot both rRSA and uniform wRSA expectation predictions in same plot
toplot = merge(toplot_r,toplot_w, all=T)
head(toplot)
nrow(toplot)
p_probs = ggplot(toplot, aes(x=PriorProbability, y=PosteriorProbability, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred",wes_palette("Darjeeling")[1])) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior mean all-state probability") +
  scale_y_continuous(limits=c(0,1), name="Predicted posterior all-state probability")
p_probs
ggsave("pred_probs.pdf")

p_probs = ggplot(toplot_r, aes(x=PriorProbability, y=PosteriorProbability, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,1), name="Prior all-state probability") +
  scale_y_continuous(limits=c(0,1), name="Posterior all-state probability")
p_probs
ggsave("pred_prob_rsa.pdf")

# plot model predictions: expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/mp-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(mp, Quantifier == "some" & WonkyWorldPrior == .5))
nrow(toplot)

pexpectations = ddply(toplot, .(Item, SpeakerOptimality,PriorExpectation), summarise, PosteriorExpectation_predicted=sum(State*PosteriorProbability)/15)
head(pexpectations)
some = pexpectations#droplevels(subset(pexpectations, Quantifier == "some"))

toplot = droplevels(subset(some, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)
toplot_w = toplot

# get rRSA predictions for qud=how-many, alts=0_basic, spopt=2
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/complex_prior/smoothed_unbinned15/results/data/toplot-expectations.RData")
toplot_r = toplot
head(toplot_w)
summary(toplot_r)
toplot_r$Model = "RSA"
toplot_r$PriorExpectation = toplot_r$PriorExpectation_smoothed*15
toplot_w$Model = "wRSA"

# plot both rRSA and uniform wRSA expectation predictions in same plot
toplot = merge(toplot_r,toplot_w, all=T)
head(toplot)
nrow(toplot)

p_exps = ggplot(toplot, aes(x=PriorExpectation, y=PosteriorExpectation_predicted*15, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred",wes_palette("Darjeeling")[1])) +#values=c("#007fb1", "#4ecdff")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Prior mean number of objects") +
  scale_y_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Predicted posterior number of objects")
p_exps
ggsave("pred_exps.pdf")

p_exps = ggplot(toplot_r, aes(x=PriorExpectation, y=PosteriorExpectation_predicted*15, color=Model)) +
  geom_point() + #color="#00B0F6") + #values=c("#F8766D", "#A3A500", "#00BF7D", "#E76BF3", "#00B0F6")
  geom_smooth() + #color="#00B0F6") +
  scale_color_manual(values=c("darkred")) +#values=c("#007fb1", "#4ecdff")) +
  #  geom_abline(intercept=0,slope=1,color="gray50") +
  scale_x_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Prior expectation") +
  scale_y_continuous(limits=c(0,15), breaks=seq(1,15,by=2), name="Posterior expectation")
p_exps
ggsave("pred_exps_rsa.pdf")

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_exps + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_exps_nolegend = p_exps + theme(legend.position="none")
p_probs_nolegend = p_probs + theme(legend.position="none")

pdf("rsa-predictions.pdf",width=12,height=4.9)
grid.arrange(p_exps_nolegend, p_probs_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

# expectations
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/13_sinking-marbles-priordv-15/results/data/r.RData")

agr = aggregate(ProportionResponse ~ PriorExpectationProportion + quantifier + Item,data=r,FUN=mean)

min(agr[agr$quantifier == "Some",]$ProportionResponse)
max(agr[agr$quantifier == "Some",]$ProportionResponse)
agr$Quantifier = factor(x=agr$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eexps = ggplot(agr, aes(x=PriorExpectationProportion*15, y=ProportionResponse*15, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  #geom_errorbar(aes(ymin=YMin,ymax=YMax)) +
  geom_smooth(method="lm") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(breaks=seq(1,15,by=2),name="Posterior mean number of objects") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects")
p_eexps
ggsave(file="empirical_exps.pdf")#,width=5,height=3.7)

# allstate-probs
load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/16_sinking-marbles-sliders-certain/results/data/r.RData")

# exclude people who are doing some sort of bullshit and not responding reasonably to all/none (see subject-variability.pdf for behavior on zero-slider)
tmp = subset(r,!workerid %in% c(0,22,43,98,100,103,117,118))
agrr = aggregate(normresponse ~ AllPriorProbability + Proportion + quantifier + Item,data=tmp,FUN=mean)
agrr = aggregate(normresponse ~ AllPriorProbability + Proportion + quantifier + Item,data=r,FUN=mean)
ub = subset(agrr, Proportion == "100")
ub = droplevels(ub)
ub$Quantifier = factor(x=ub$quantifier, levels=c("Some","All","None","long_filler","short_filler"))

p_eprobs = ggplot(ub, aes(x=AllPriorProbability, y = normresponse, color=Quantifier, shape=Quantifier)) +
  geom_point() +
  geom_smooth(method="lm") +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +#values=c("#F8766D", "black", "#00BF7D", "gray30", "#00B0F6"),breaks=levels(agr$quantifier),labels=c("all","long filler","none","short filler","some")) +
  scale_y_continuous(limits=c(0,1),name="Posterior probability of all-state ") +
  #  geom_abline(intercept=0,slope=1,color="gray70") +
  scale_x_continuous(limits=c(0,1),name="Prior probability of all-state")
p_eprobs
ggsave(file="empirical-allprobs.pdf")#,width=5,height=3.7)

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_eexps + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_eexps_nolegend = p_eexps + theme(legend.position="none")
p_eprobs_nolegend = p_eprobs + theme(legend.position="none")

pdf("empirical-results.pdf",width=12,height=4.9)
grid.arrange(p_eexps_nolegend, p_eprobs_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()

load("~/cogsci/projects/stanford/projects/thegricean_sinking-marbles/models/wonky_world/results/data/wr-uniform.RData")

# plot expectations for best basic model:
toplot = droplevels(subset(wr, WonkyWorldPrior == .5 & Wonky == "true"))
nrow(toplot)
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode

toplot = droplevels(subset(toplot, SpeakerOptimality == 2))
nrow(toplot)
head(toplot)
#toplot$quantifier = capitalize(as.character(toplot$Quantifier))
toplot$Quantifier = factor(toplot$Quantifier, levels=c("some","all","none"))

p_wmodel = ggplot(toplot, aes(x=PriorExpectation, y=PosteriorProbability,color=Quantifier,shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +
  #scale_color_manual(values=c("#F8766D","#00BF7D","#00B0F6")) +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects") +
  scale_y_continuous(breaks=seq(0,1,by=.25),name="Predicted posterior wonkiness probability") +
  theme(plot.margin=unit(c(0,0,0,0),units="cm"))
p_wmodel
ggsave("model-wonkiness-uniform.pdf",width=6.8,height=5)#,width=30,height=10)

############
# empirical wonkiness posteriors
############
load("/Users/titlis/cogsci/projects/stanford/projects/thegricean_sinking-marbles/experiments/17_sinking-marbles-normal-sliders/results/data/r.RData")
head(r)
nrow(r)

toplot = aggregate(response ~ quantifier + Item + PriorExpectation, FUN="mean", data=r)
toplot = droplevels(subset(toplot, quantifier %in% c("All","Some","None")))
toplot$Quantifier = factor(tolower(toplot$quantifier), levels=c("all", "none","some"))
toplot$Mode = priormodes[as.character(toplot$Item),]$Mode
toplot$quantifier = capitalize(as.character(toplot$Quantifier))
toplot$Quantifier = factor(toplot$quantifier, levels=c("Some","All","None"))

p_wempirical = ggplot(toplot, aes(x=PriorExpectation, y=response, color=Quantifier,shape=Quantifier)) +
  geom_point() +
  geom_smooth() +
  scale_color_manual(values=c(wes_palette("Darjeeling")[1:3],"black","gray40")) +
  scale_x_continuous(breaks=seq(1,15,by=2),name="Prior mean number of objects") +
  scale_y_continuous(breaks=seq(0,1,by=.25),name="Mean empirical wonkiness probability")  +
  theme(plot.margin=unit(c(0,0,0,0),units="cm"))
p_wempirical
ggsave(file="empirical-wonkiness.pdf",width=6.8,height=5)

library(gridExtra)
#  share a legend between multiple plots
g <- ggplotGrob(p_wmodel + theme(legend.position="right"))$grobs
legend <- g[[which(sapply(g, function(x) x$name) == "guide-box")]]
p_wmodel_nolegend = p_wmodel + theme(legend.position="none")
p_wempirical_nolegend = p_wempirical + theme(legend.position="none")

#pdf("rsa-predictions-uniform.pdf",width=10,height=4)
pdf("wonkiness-fullplot.pdf",width=12,height=5)
grid.arrange(p_wmodel_nolegend, p_wempirical_nolegend, legend,nrow=1,widths=unit.c(unit(.45, "npc"), unit(.45, "npc"), unit(.1, "npc")))
dev.off()