# packages needed
source("project_paths.r")

library(tidyverse)
library(plm)
library(lmtest)
library(broom)
# library(flextable)
# library(officer)
library(knitr)

source(paste(PATH_IN_ANALYSIS, "unconditionalModels.R", sep="/"))
source(paste(PATH_IN_ANALYSIS, "conditionalModels.R", sep="/"))

outcomes <- read_rds(paste(PATH_OUT_DATA, "outcomes.rds", sep="/"))
events_by_time <-  read_rds(paste(PATH_OUT_DATA, "events_by_time.rds", sep="/"))
demographics <-read_rds(paste(PATH_OUT_DATA, "demographics.rds", sep="/"))

knitr::opts_chunk$set(include = FALSE,
                      eval = TRUE,
                      echo = FALSE,
                      message = FALSE,
                      warning = FALSE,
                      fig.path = "figures/")

# 21 life events included

lifevents <- c('lefnw', 'lefni', 'lefrd', 'lejob', 'leprm', 'lertr', 'lemar',
               'leprg', 'lebth', 'ledsc', 'lesep', 'lemvd', 'leins', 'lepcm',
               'lejls', 'lejlf', 'ledfr', 'ledrl', 'leinf', 'lercl', 'levio')

# 10 time points included
timelevels = c('pre36', 'pre24', 'pre12', 'post03', 'post06', 'post09',
               'post12', 'post24', 'post36', 'post48')

# name of events
renamevents <- list(
    Widowed = "ledsc",
    Separated = "lesep",
    Money_lost = "lefnw",
    Jailed = "lejls",
    Attacked = "levio",
    Health_shock = "leins",
    Reconciled = "lercl",
    Fired = "lefrd",
    Family_harmed = "leinf",
    Robbed = "lepcm",
    Friend_died = "ledfr",
    Relative_died = "ledrl",
    Relative_jailed = "lejlf",
    Moved = "lemvd",
    Hired = "lejob",
    Promoted = "leprm",
    Retired = "lertr",
    Money_gained = "lefni",
    Pregnant = "leprg",
    Childbirth = "lebth",
    Married = "lemar"
    )

# good event names
good_events <- c(
  "Married", "Childbirth", "Money_gained",
  "Retired", "Pregnant", "Promoted",
  "Hired", "Moved", "Reconciled"
  )

# bad event names
bad_events <- c(
  "Widowed", "Separated", "Money_lost",
  "Health_shock", "Attacked", "Fired",
  "Family_harmed", "Friend_died", "Robbed"
  )

#Table 1. List of 22 life events and their description*

tribble(
  ~Event, ~Description,
  "Widowed", "death of a spouse or a child",
  "Separated", "separated or divorced from a spouse or long-term partner",
  "Money lost", "major worsening in financial situation (e.g. bankruptcy)",
  "Jailed", "detained in jail / correctional facility",
  "Attacked", "victim of physical violence (e.g. assault)",
  "Health shock", "serious personal injury or illness (e.g., disability)",
  "Reconciled", "reconciled with spouse/long-term partner after separation",
  "Fired", "fired or made redundant by an employer",
  "Family harmed", "serious injury or illness to a close family member",
  "Robbed", "victim of property crime (e.g. theft, house breaking)",
  "Friend died", "death of a close friend",
  "Relative died", "death of a close family member (e.g. parent or sibling)",
  "Relative jailed", "jail for a close friend or relative",
  "Moved", "changed residence",
  "Hired", "changed jobs (i.e. employer)",
  "Promoted", "promoted at work",
  "Retired", "retired from workforce",
  "Money gained", "major gain in finances (e.g., lottery win, inheritance)",
  "Pregnant", "you (or your partner) got pregnant",
  "Married", "got married",
  "Childbirth","birth (or adoption) of a child"
  ) %>%
  kable()

outcomes %>%
  filter(code == 'losatyh') %>%
  mutate(
    Y = case_when(
      val < 0 ~ NA_integer_,
      TRUE ~ as.integer(val)
    )
  ) %>%
  select(-code, -val) -> cognitive.df

outcomes %>%
  filter(code == 'gh9_sum') %>%
  mutate(
    Y = case_when(
      val < 0 ~ NA_integer_, # no negatives in gh9_sum
      TRUE ~ as.integer(val)
    )
  ) %>%
  select(-code, -val) -> affect.df

outcomes %>%
  filter(code == 'losatsf') %>%
  mutate(
    Y = case_when(
      val < 0 ~ NA_integer_, # no negatives in gh9_sum
      TRUE ~ as.integer(val)
    )
  ) %>%
  select(-code, -val) -> safety.df

outcomes %>%
  filter(code == 'losatfs') %>%
  mutate(
    Y = case_when(
      val < 0 ~ NA_integer_, # no negatives in gh9_sum
      TRUE ~ as.integer(val)
    )
  ) %>%
  select(-code, -val) -> finance.df

outcomes %>%
  filter(code == 'firisk') %>%
  mutate(
    Y = case_when(
      val < 0 ~ NA_integer_, # no negatives in gh9_sum
      TRUE ~ as.integer(val)
    )
  ) %>%
  select(-code, -val) -> risk_aversion.df

# outcomes %>%
#   filter(code == 'mcs') %>%
#   select(-code, Y = val) -> affect.df

# load the list of life event data (for preprocessing see Appendix)
#events_by_time <- read_rds('src/data/events_by_time.rds')

# These models include all available observations of each event as a single
# independent variable (+ wave) in each model, thus representing the uncondit-
# ional effect of life events, and ignoring any mediating effects of other
# events as well as the risk of contamination from unobserved events.
#

# Our helpful wrapper returns the coefficients, residuals and covariance matrix

cognitive.results <- unconditionalModels(cognitive.df, events_by_time, timelevels, lifevents)
affective.results <- unconditionalModels(affect.df, events_by_time, timelevels, lifevents)

cognitive.resids <- cognitive.results$resids
affective.resids <- affective.results$resids
cognitive.resids$outcome <- 'health'
affective.resids$outcome <- 'emotional'
resid.results <- bind_rows(cognitive.resids, affective.resids)
write_rds(resid.results,paste(PATH_OUT_DATA, "unconditional_resid.rds", sep="/"))

covar.results <- list(cognitive = cognitive.results$covar,
                      affective = affective.results$covar)

write_rds(covar.results,paste(PATH_OUT_DATA, "unconditional_covar.rds", sep="/"))

# Get the coefficients and combine them into a table of ordered time points
cognitive.coefs <- cognitive.results$coefs
affective.coefs <- affective.results$coefs
cognitive.coefs$outcome <- 'health'
affective.coefs$outcome <- 'emotional'
bind_rows(cognitive.coefs, affective.coefs) -> results

# Calculate confidence intervals on clustered errors
results$upper <- results$estimate -
                      results$std.error*qt(p = 0.1, df = results$n-1)
results$lower <- results$estimate +
                      results$std.error*qt(p = 0.1, df = results$n-1)

# Save results
write_csv(results,paste(PATH_OUT_DATA, "unconditional_results.csv", sep="/"))

safety_effect.results <- unconditionalModels(safety.df, events_by_time, timelevels, lifevents)
finance_effect.results <- unconditionalModels(finance.df, events_by_time, timelevels, lifevents)

# Get the coefficients and combine them into a table of ordered time points
safety_effect.coefs <- safety_effect.results$coefs
finance_effect.coefs <- finance_effect.results$coefs
safety_effect.coefs$outcome <- 'safety'
finance_effect.coefs$outcome <- 'finance'
bind_rows(safety_effect.coefs, finance_effect.coefs) -> results

# Calculate confidence intervals on clustered errors
results$upper <- results$estimate -
                      results$std.error*qt(p = 0.1, df = results$n-1)
results$lower <- results$estimate +
                      results$std.error*qt(p = 0.1, df = results$n-1)
write_csv(results,paste(PATH_OUT_DATA, "unconditional_results_finance_and_safety.csv", sep="/"))

##########

risk_results <- unconditionalModels(risk_aversion.df, events_by_time, timelevels, lifevents)

# Get the coefficients and combine them into a table of ordered time points
risk_effect.coefs <-risk_results$coefs
finance_effect.coefs <- finance_effect.results$coefs
risk_effect.coefs$outcome <- 'risk'
finance_effect.coefs$outcome <- 'finance'
bind_rows(risk_effect.coefs, finance_effect.coefs) -> results

# Calculate confidence intervals on clustered errors
results$upper <- results$estimate -
                      results$std.error*qt(p = 0.1, df = results$n-1)
results$lower <- results$estimate +
                      results$std.error*qt(p = 0.1, df = results$n-1)
write_csv(results,paste(PATH_OUT_DATA, "unconditional_results_finance_and_risk.csv", sep="/"))


require(nza)
require(nzr)
require(xgboost)
require(caTools)
require(caret)
require(pROC)

nzConnectDSN('DBM_SANDBOX')

if (nzExistTable('NZA_TEST')) nzDeleteTable('NZA_TEST')

nzQuery("CREATE TABLE DBM_SANDBOX..NZA_TEST AS
                SELECT *
                FROM (
                 SELECT DETAIL.LAST_NAME,
                 DETAIL.ADDRESS_LINE_1,
                 DETAIL.CITY_NAME,
                 DETAIL.STATE_CODE,
                 DETAIL.ZIP,
                 MAX(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-(365) AND CURRENT_DATE::DATE THEN 1 ELSE 0 END) AS PURCH_6MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-(365) AND CURRENT_DATE::DATE THEN DETAIL.SALES ELSE 0 END) AS SLS_6MO,
                 MAX(CASE WHEN DETAIL.REWARDS_ACTIVE_CODE IN ('A','Y') THEN 1 ELSE 0 END) AS REWARDS,
                 MAX(CASE WHEN DETAIL.VISA_STATUS = 'Y' THEN 1 ELSE 0 END) AS VISA,
                 MIN(DETAIL.DATE_VALUE) AS FIRST_DATE,
                 MAX(DETAIL.DATE_VALUE) AS LAST_DATE,
                 MAX(DETAIL.DATE_VALUE)-MIN(DETAIL.DATE_VALUE) AS TIME_AS_CUSTOMER,
                 COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY) AS ORDERS_LT,
                 (MAX(DETAIL.DATE_VALUE) - MIN(DETAIL.DATE_VALUE)) / NULLIF(COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY), 0)::DOUBLE AS AVG_DAYS_BTW_PURCH,
                 NVL(SUM(DETAIL.SALES)/NULLIF(COUNT(DISTINCT DETAIL.INBOUND_INTERACTION_KEY),0)::DOUBLE,0) AS AVG_TICKET_LT,
                 NVL(SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE THEN DETAIL.SALES ELSE 0 END) /
                 NULLIF(COUNT(DISTINCT CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE THEN DETAIL.INBOUND_INTERACTION_KEY ELSE NULL END),0),0) AS AVG_TICKET_24MO,
                 NVL(SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE THEN DETAIL.SALES ELSE 0 END) /
                 NULLIF(COUNT(DISTINCT CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE THEN DETAIL.INBOUND_INTERACTION_KEY ELSE NULL END),0),0) AS AVG_TICKET_48MO,
                 SUM(CASE WHEN DETAIL.CHANNEL_GROUP = 'RETAIL' THEN DETAIL.SALES ELSE 0 END)/ NULLIF(SUM(DETAIL.SALES),0) AS PERCENT_RETAIL_SALES,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE THEN DETAIL.SALES ELSE 0 END) AS SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-365 AND CURRENT_DATE::DATE THEN DETAIL.SALES ELSE 0 END) AS SALES_12MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER IN (100, 151, 200, 400, 450) THEN DETAIL.SALES ELSE 0 END) AS CORE_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 100 THEN DETAIL.SALES ELSE 0 END) AS D100_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 151 THEN DETAIL.SALES ELSE 0 END) AS D151_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 175 THEN DETAIL.SALES ELSE 0 END) AS D175_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 200 THEN DETAIL.SALES ELSE 0 END) AS D200_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 300 THEN DETAIL.SALES ELSE 0 END) AS D300_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 350 THEN DETAIL.SALES ELSE 0 END) AS D350_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 400 THEN DETAIL.SALES ELSE 0 END) AS D400_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 450 THEN DETAIL.SALES ELSE 0 END) AS D450_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 475 THEN DETAIL.SALES ELSE 0 END) AS D475_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 500 THEN DETAIL.SALES ELSE 0 END) AS D500_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 600 THEN DETAIL.SALES ELSE 0 END) AS D600_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 650 THEN DETAIL.SALES ELSE 0 END) AS D650_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-729 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 675 THEN DETAIL.SALES ELSE 0 END) AS D675_SALES_24MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER IN (100, 151, 200, 400, 450) THEN DETAIL.SALES ELSE 0 END) AS CORE_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 100 THEN DETAIL.SALES ELSE 0 END) AS D100_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 151 THEN DETAIL.SALES ELSE 0 END) AS D151_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 175 THEN DETAIL.SALES ELSE 0 END) AS D175_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 200 THEN DETAIL.SALES ELSE 0 END) AS D200_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 300 THEN DETAIL.SALES ELSE 0 END) AS D300_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 350 THEN DETAIL.SALES ELSE 0 END) AS D350_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 400 THEN DETAIL.SALES ELSE 0 END) AS D400_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 450 THEN DETAIL.SALES ELSE 0 END) AS D450_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 475 THEN DETAIL.SALES ELSE 0 END) AS D475_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 500 THEN DETAIL.SALES ELSE 0 END) AS D500_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 600 THEN DETAIL.SALES ELSE 0 END) AS D600_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 650 THEN DETAIL.SALES ELSE 0 END) AS D650_SALES_48MO,
                 SUM(CASE WHEN DETAIL.DATE_VALUE BETWEEN CURRENT_DATE::DATE-1458 AND CURRENT_DATE::DATE AND DETAIL.DEPARTMENT_MEMBER_NUMBER = 675 THEN DETAIL.SALES ELSE 0 END) AS D675_SALES_48MO
                 FROM (
                 SELECT IND.LAST_NAME,
                 ADDR.ADDRESS_LINE_1,
                 ADDR.CITY_NAME,
                 ADDR.STATE_CODE,
                 ADDR.ZIP,
                 IND.REWARDS_ACTIVE_CODE,
                 IND.VISA_STATUS,
                 DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY,
                 ST.CHANNEL_GROUP,
                 PD.DEPARTMENT_MEMBER_NUMBER,
                 PD.SUB_DEPARTMENT_MEMBER_NUMBER,
                 SUM(SD.SALE_QUANTITY) AS UNITS,
                 SUM(SD.SALES_PRICE) AS SALES
                 FROM EDW_SPOKE..FACT_BPS_SALES_HEADER SH
                 JOIN EDW_SPOKE..FACT_BPS_SALES_DETAIL SD ON SH.INBOUND_INTERACTION_KEY = SD.SH_INBOUND_INTERACTION_KEY
                 JOIN EDW_SPOKE..DIM_STORE ST ON SD.STORE_MEMBER_KEY = ST.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_DATE DT ON SD.SALE_DATE_KEY = DT.DATE_KEY
                 JOIN EDW_SPOKE..DIM_PRODUCT_BPS PD ON SD.SALES_PRODUCT_MEMBER_KEY = PD.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_PERSONA PERS ON SH.CUSTOMER_MEMBER_KEY = PERS.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_INDIVIDUAL IND ON PERS.INDIVIDUAL_MEMBER_KEY = IND.MEMBER_KEY
                 JOIN EDW_SPOKE..DIM_ADDRESS ADDR ON IND.ADDRESS_KEY = ADDR.MEMBER_KEY
                 WHERE ST.CHANNEL_GROUP IN ('DIRECT','RETAIL')
                 AND DT.DATE_VALUE BETWEEN CURRENT_DATE::DATE - (365*5) AND CURRENT_DATE::DATE
                 AND SD.ACTUAL_SALE_FLAG = 'Y'
                 AND SD.RETURN_FLAG <> 'Y'
                 AND PD.DEPARTMENT_MEMBER_NUMBER NOT IN (875,999)
                 GROUP BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP, IND.REWARDS_ACTIVE_CODE, IND.VISA_STATUS, DT.DATE_VALUE,
                 SH.INBOUND_INTERACTION_KEY, ST.CHANNEL_GROUP, PD.DEPARTMENT_MEMBER_NUMBER, PD.SUB_DEPARTMENT_MEMBER_NUMBER
                 ORDER BY IND.LAST_NAME, ADDR.ADDRESS_LINE_1, ADDR.CITY_NAME, ADDR.STATE_CODE, ADDR.ZIP, DT.DATE_VALUE, SH.INBOUND_INTERACTION_KEY,
                 PD.DEPARTMENT_MEMBER_NUMBER, PD.SUB_DEPARTMENT_MEMBER_NUMBER
                 ) AS DETAIL
                 GROUP BY DETAIL.LAST_NAME, DETAIL.ADDRESS_LINE_1, DETAIL.CITY_NAME, DETAIL.STATE_CODE, DETAIL.ZIP
                ) AS A
                WHERE A.AVG_TICKET_48MO > 0")

if (nzExistTable('CHURNMODEL')) nzDeleteTable('CHURNMODEL')

nzQuery('CREATE TABLE DBM_SANDBOX..CHURNMODEL AS
         SELECT LAST_NAME,
                ADDRESS_LINE_1,
                CITY_NAME,
                STATE_CODE,
                ZIP,
                (ORDERS_LT - AVG(ORDERS_LT) OVER()) / STDDEV(ORDERS_LT) OVER() AS ORDERS_LT,
                (AVG_DAYS_BTW_PURCH - AVG(AVG_DAYS_BTW_PURCH) OVER()) / STDDEV(AVG_DAYS_BTW_PURCH) OVER() AS AVG_DAYS_BTW_PURCH,
                (AVG_TICKET_LT - AVG(AVG_TICKET_LT) OVER()) / STDDEV(AVG_TICKET_LT) OVER() AS AVG_TICKET_LT,
                (TIME_AS_CUSTOMER - AVG(TIME_AS_CUSTOMER) OVER()) / STDDEV(TIME_AS_CUSTOMER) OVER() AS TIME_AS_CUSTOMER,
                (SALES_12MO - AVG(SALES_12MO) OVER()) / STDDEV(SALES_12MO) OVER() AS SALES_12MO,
                (CORE_SALES_24MO - AVG(CORE_SALES_24MO) OVER()) / STDDEV(CORE_SALES_24MO) OVER() AS CORE_SALES_24MO,
                (D200_SALES_24MO - AVG(D200_SALES_24MO) OVER()) / STDDEV(D200_SALES_24MO) OVER() AS D200_SALES_24MO,
                (D350_SALES_24MO - AVG(D350_SALES_24MO) OVER()) / STDDEV(D350_SALES_24MO) OVER() AS D350_SALES_24MO,
                (D400_SALES_24MO - AVG(D400_SALES_24MO) OVER()) / STDDEV(D400_SALES_24MO) OVER() AS D400_SALES_24MO,
                (D450_SALES_24MO - AVG(D450_SALES_24MO) OVER()) / STDDEV(D450_SALES_24MO) OVER() AS D450_SALES_24MO,
                PURCH_6MO,
                REWARDS,
                VISA
        FROM DBM_SANDBOX..NZA_TEST
        limit 1000')

target <-c('TARGET_RESPONSE')
finalvars <- c('ORDERS_LT','AVG_DAYS_BTW_PURCH','AVG_TICKET_LT', 'TIME_AS_CUSTOMER',
               'SALES_12MO','CORE_SALES_24MO', 'D200_SALES_24MO', 'D350_SALES_24MO',
               'D400_SALES_24MO','D450_SALES_24MO','PURCH_6MO','REWARDS', 'VISA')

nzChurn <- nz.data.frame('CHURNMODEL')
head(nzChurn)

bst <- xgb.load('C:\\Users\\pairwin\\Documents\\GitHub\\RAnalyses\\xgboost_churn.model')

load("C:\\Users\\pairwin\\Documents\\GitHub\\RAnalyses\\CHURN_LOGISTIC.rda") #logmod

pred <- function(x, model1) {
 predval <- predict(model1, data.frame(LAST_NAME=as.character(x[[1]]),
                             ADDRESS_LINE_1=as.character(x[[2]]),
                             CITY_NAME=as.character(x[[3]]),
                             STATE_CODE=as.character(x[[4]]),
                             ZIP=as.character(x[[5]]),
                          ORDERS_LT=as.numeric(x[[6]]),
                          AVG_DAYS_BTW_PURCH=as.numeric(x[[7]]),
                          AVG_TICKET_LT=as.numeric(x[[8]]),
                          TIME_AS_CUSTOMER=as.numeric(x[[9]]),
                          SALES_12MO=as.numeric(x[[10]]),
                          CORE_SALES_24MO=as.numeric(x[[11]]),
                          D200_SALES_24MO=as.numeric(x[[12]]),
                          D350_SALES_24MO=as.numeric(x[[13]]),
                          D400_SALES_24MO=as.numeric(x[[14]]),
                          D450_SALES_24MO=as.numeric(x[[15]]),
                          PURCH_6MO=factor(x[[16]],levels=c(1,0)),
                          REWARDS=factor(x[[17]], levels=c(1,0)),
                          VISA=factor(x[[18]], levels=c(1,0))))
 return(predval)
}

nzApply(nzChurn,
        FUN=pred,
        model1=logmod,
        output.name='CHURNSCORED',
        clear.existing=TRUE,
        # output.signature = list(LAST_NAME=list(NZ.NATIONAL_VARIABLE,30),
        #                         ADDRESS_LINE_1=list(NZ.NATIONAL_VARIABLE,30),
        #                         CITY_NAME=list(NZ.NATIONAL_VARIABLE,30),
        #                         STATE_CODE=list(NZ.NATIONAL_VARIABLE,30),
        #                         ZIP=list(NZ.NATIONAL_VARIABLE,30),
        #                         ORDERS_LT=NZ.DOUBLE,
        #                         AVG_DAYS_BTW_PURCH=NZ.DOUBLE,
        #                         AVG_TICKET_LT=NZ.DOUBLE,
        #                         TIME_AS_CUSTOMER=NZ.DOUBLE,
        #                         SALES_12MO=NZ.DOUBLE,
        #                         CORE_SALES_24MO=NZ.DOUBLE,
        #                         D200_SALES_24MO=NZ.DOUBLE,
        #                         D350_SALES_24MO=NZ.DOUBLE,
        #                         D400_SALES_24MO=NZ.DOUBLE,
        #                         D450_SALES_24MO=NZ.DOUBLE,
        #                         PURCH_6MO=NZ.INT64,
        #                         REWARDS=NZ.INT64,
        #                         VISA=NZ.INT64,
        #                         SCORE=NZ.DOUBLE))
        output.signature = list(SCORE=NZ.DOUBLE))

nzDisconnect()


require(data.table)
require(rstan)
require(EnvStats)
require(covid19AgeModel)
pkg.dir <- system.file(package = "covid19AgeModel" )

## command line parsing if any
args <- list(
  indir.results= '/rds/general/project/ratmann_covid19/live/age_renewal_usa/base_age_fsq_mobility_201015e8_cmdstanv-4states_AZCTFKNYC_Sep20_Levin/base_age_fsq_mobility_201015e8_cmdstanv-4states_AZCTFKNYC_Sep20_Levin-2474173[1].pbs',
  location.index= 1,
  with.flow=1,
  forecast.period= 0, # days
  school.reopen= 1,
  multiplier_cntct_school_opening=1,
  school_level="K5",
  shield = 1,
  return_age_band='20-49',
  shielded_age_band='60-100',
  shielded_contacts=0.5,
  shielded_date='2020-06-20'
)

args_line <-  as.list(commandArgs(trailingOnly=TRUE))
if(length(args_line) > 0)
{
  stopifnot(args_line[[1]]=='-indir.results')
  stopifnot(args_line[[3]]=='-location.index')
  stopifnot(args_line[[5]]=='-with.flow')
  stopifnot(args_line[[7]]=='-forecast.period')
  stopifnot(args_line[[9]]=='-school.reopen')
  stopifnot(args_line[[11]]=='-multiplier_cntct_school_opening')
  stopifnot(args_line[[13]]=='-school_level')
  stopifnot(args_line[[15]]=='-shield')
  stopifnot(args_line[[17]]=='-return_age_band')
  stopifnot(args_line[[19]]=='-shielded_age_band')
  stopifnot(args_line[[21]]=='-shielded_contacts')
  stopifnot(args_line[[23]]=='-shielded_date')
  args <- list()
  args[['indir.results']] <- args_line[[2]]
  args[['location.index']] <- args_line[[4]]
  args[['with.flow']] <- args_line[[6]]
  args[['forecast.period']] <- args_line[[8]]
  args[['school.reopen']] <- args_line[[10]]
  args[['multiplier_cntct_school_opening']] <- args_line[[12]]
  args[['school_level']] <- args_line[[14]]
  args[['shield']] <- args_line[[16]]
  args[['return_age_band']] <- args_line[[18]]
  args[['shielded_age_band']] <- args_line[[20]]
  args[['shielded_contacts']] <- args_line[[22]]
  args[['shielded_date']] <- args_line[[24]]
}

#	print args
str(args)

#	read args
indir.results <- args$indir.results
location.index <- as.integer(args$location.index)
with.flow <- args$with.flow
forecast.period <- as.numeric(args$forecast.period)
multiplier_cntct_school_opening <- as.numeric(args$multiplier_cntct_school_opening)
school.reopen <- args$school.reopen == '1'
school_level <- as.character(args$school_level)
shield=as.character(args$shield)
return_age_band=as.character(args$return_age_band)
shielded_age_band=as.character(args$shielded_age_band)
shielded_contacts=as.numeric(args$shielded_contacts)
shielded_date=as.Date(args$shielded_date)

#	read Stan input data and add location.index to parallelise computations
cat('\nReading Stan input data...')
infile.stanin <- list.files(indir.results, pattern='.*_stanin.RData$', recursive=TRUE)
stopifnot(length(infile.stanin)>0)
stopifnot(length(infile.stanin)<=1)
tmp <- load(file.path(indir.results, infile.stanin))
stopifnot(c('args','stan_data')%in%tmp)
stan_data$LOCATION_PROCESSING_IDX <- location.index

# reset args
forecast <- forecast.period
args$file_contact_intensities_outbreak_China <- file.path(pkg.dir, "data", "estimate_contact_intensities_outbreak_China.rds")

if(is.null(args$with_EduWeek_data)) args$with_EduWeek_data = 0
if(is.null(args$effect_weeks)) args$effect_weeks <- 2

# determine which levels have closed/reopened
args$with_elementary_schools_status = 0
args$with_elementary_high_schools_status = 0
if(school_level == "K5"){
  args$with_elementary_schools_status = 1
}
if(school_level == "K12"){
  args$with_elementary_high_schools_status = 1
}

num_days_sim <- (max(death_data$date) - min(death_data$date) + 1 + forecast)[[1]]

# reprocess data
processed_data <- make_stan_data_core(states = args$states,
                                      death_data = death_data,
                                      deathByAge_data = deathByAge_data,
                                      ifr.by.age = ifr.by.age,
                                      serial_interval = serial_interval,
                                      pop_info = pop_info,
                                      dcontact= dcontact,
                                      num_days_sim = num_days_sim,
                                      forecast = forecast)

if(args$with_google_mobility &&
   args$decouple_mobility_trends_into_baseline_plus_rest_parts==0 &&
   args$decouple_mobility_trends_into_decline_plus_upswing_parts==0 &&
   args$decouple_mobility_trends_into_baseline_plus_decline_plus_upswing_parts==0)
{
  cat('\nAdding google mobility covariates to stan_data ...')
  processed_data <- stan_data_add_google_mobility(processed_data, mobility_data, args$with_avg_mobility_data)
}
if(args$with_fsq_mobility &&
   args$decouple_mobility_trends_into_baseline_plus_rest_parts==0 &&
   args$decouple_mobility_trends_into_decline_plus_upswing_parts==0 &&
   args$decouple_mobility_trends_into_baseline_plus_decline_plus_upswing_parts==0)
{
  cat('\nAdding Foursquare mobility covariates to stan_data ...')
  processed_data <- stan_data_add_mobility_trends_fsq(processed_data, mobility_data, log=TRUE)
}
if( args$with_google_mobility &&
    args$decouple_mobility_trends_into_decline_plus_upswing_parts==1 )
{
  cat('\nAdding 2 part decoupled Google mobility trends to stan_data ... ')
  processed_data <- stan_data_add_google_mobility_covariates_after_breakpoint(processed_data)
}
if( args$with_fsq_mobility==1 &&
    args$decouple_mobility_trends_into_decline_plus_upswing_parts==1 )
{
  cat('\nAdding decoupled FSQ mobility trends (decline + upswing) to stan_data ... ')
  processed_data <- stan_data_add_decoupled_mobility_trends_decline_plus_upswing_fsq(processed_data, mobility_data, log=TRUE)
}
if( args$with_fsq_mobility==1 &&
    args$decouple_mobility_trends_into_baseline_plus_rest_parts==1 )
{
  cat('\nAdding decoupled FSQ mobility trends (baseline + rest) to stan_data ... ')
  processed_data <- stan_data_add_decoupled_mobility_trends_decline_plus_rest_fsq(processed_data, mobility_data, log=TRUE)
}
if( args$with_fsq_mobility==1 &&
    args$decouple_mobility_trends_into_baseline_plus_decline_plus_upswing_parts==1 )
{
  cat('\nAdding decoupled FSQ mobility trends (baseline + decline + upswing) to stan_data ... ')
  processed_data <- stan_data_add_decoupled_mobility_trends_3_parts_fsq(processed_data, mobility_data, log=TRUE)
}
if( args$with_estimated_pairwise_mobility_trends)
{
  # process_chi_emodo_mobilty(mobility_data,pop_info,infile_emodo='~/contacts_and_mobility/contacts_by_age_20200729.csv')
  cat('\n Add pairwise mobility trend ... ')
  processed_data <- stan_data_add_pairwise_mobility_trends(processed_data,file.path(pkg.dir,'data','chi_emodo_mobility_trends.rds'))
}
if( args$with_biweekly_upswing_time_effect )
{
  cat('\nAdding upswing time effects ... ')
  processed_data <- stan_data_add_upswing_time_effect(
	processed_data,
	mobility_data,
	effect_weeks=args$effect_weeks)
}
if( args$with_contact_intensities_zhang && !args$with_EduWeek_data )
{
  reopening_date <- as.Date("2020-08-24")
  cat('\nAdding school status ... ')
  processed_data <- stan_data_add_school_status(
    processed_data,
    args$file_school_Ox_NPI_data,
    school.reopen,
    reopening_Date= reopening_date
  )
}

if( args$with_contact_intensities_zhang && args$with_EduWeek_data)
{
  default_school_reopening_date <- as.Date("2020-08-24")
  cat('\nAdding school status with Education Weekly data ... ')
  processed_data <- stan_data_add_school_status_EduWeek(
    processed_data = processed_data,
    path_to_file_school_Ox_data = args$file_school_Ox_NPI_data,
    path_to_file_school_EduWeek_data = args$file_school_EduWeek_data,
    default_school_reopening_date = default_school_reopening_date,
    counterfactual_scenario = as.numeric(school.reopen == 0),
    forecast_with_schools_reopened = school.reopen
  )
}

if( args$with_contact_intensities_zhang && args$with_elementary_schools_status ){
  cat('\nAdding contact intensities during elementary school closure from Zhang et al. and mutliplier... ')
  processed_data <- stan_data_add_contact_school_closure(
    processed_data,
    dcontact,
    args$file_contact_intensities_outbreak_China,
    args$multiplier_cntct_school_closure
  )
}
if( args$with_contact_intensities_zhang && args$with_elementary_high_schools_status ){
  cat('\nAdding contact intensities during elementary and high school closure from Zhang et al. and mutliplier... ')
  processed_data <- stan_data_add_contact_elementary_high_school_closure(
    processed_data,
    dcontact,
    args$file_contact_intensities_outbreak_China,
    args$multiplier_cntct_school_closure
  )
}
if( args$with_contact_intensities_zhang && !args$with_EduWeek_data && args$with_elementary_schools_status )
{
  cat('\nAdding contact intensities after school re-opening from Zhang et al. and mutliplier... ')
  processed_data <- stan_data_add_contact_school_opening(
    processed_data,
    dcontact,
    args$file_contact_intensities_outbreak_China,
    args$multiplier_cntct_school_closure,
    multiplier_cntct_school_opening=multiplier_cntct_school_opening,
    min_pc_contacts=0.05,
    reopening_date= reopening_date
  )
}
if(args$with_contact_intensities_zhang && args$with_EduWeek_data && args$with_elementary_schools_status){
  cat('\nAdding contact intensities after school re-opening from Zhang et al. and mutliplier... ')
  processed_data <- stan_data_add_contact_school_opening_EduWeek(
    processed_data,
    dcontact,
    path_to_file_contact_intensities_outbreak_China=args$file_contact_intensities_outbreak_China,
    args$multiplier_cntct_school_closure,
    multiplier_cntct_school_opening=multiplier_cntct_school_opening,
    min_pc_contacts=0.05
  )
}

if(args$with_contact_intensities_zhang && args$with_EduWeek_data && args$with_elementary_high_schools_status){
  cat('\nAdding contact intensities after elementary and high school re-opening from Zhang et al. and mutliplier... ')
  processed_data <- stan_data_add_contact_elementary_high_school_opening_EduWeek(
    processed_data,
    dcontact,
    args$file_contact_intensities_outbreak_China,
    args$multiplier_cntct_school_closure,
    multiplier_cntct_school_opening=multiplier_cntct_school_opening,
    min_pc_contacts=0.05
  )
}

# reset stan_data
for (name in names(processed_data$stan_data)) {
  stan_data[[name]] <- processed_data$stan_data[[name]]
}

if(with.flow=='1')
{
  stan_data <- gqs_add_stan_data_for_flows(stan_data,dates)
}

if(is.null(stan_data$rev_iar_daysSinceInfection))
{
  stan_data <- gqs_add_stan_data_for_iar(stan_data)
}
if(shield=='1')
{
  stan_data <- gqs_add_stan_data_for_shield(stan_data, return_age_band,shielded_age_band,shielded_contacts,shielded_date, dates)
}

# reset stan data
stan_data$upswing_timeeff_map = stan_data$upswing_timeeff_map[seq_len(stan_data$N2),]

#	read stanfits
cat('\nReading Stanfit ...')
infile.stanfits <- list.files(indir.results, pattern='.*_stanout.RData$', recursive=TRUE)
stopifnot(length(infile.stanfits)>0)
stopifnot(length(infile.stanfits)<=1)
tmp <- load(file.path(indir.results, infile.stanfits[1]))

#	find stan model and stan model for generating quantities
file_stanModel <- gsub('.*covid19AgeModel/(.*)','\\1',args$file_stanModel)
# file_stanModel <- file.path(pkg.dir, file_stanModel)
file_stanModel <- file.path(pkg.dir, "inst",file_stanModel)
file_stanModel_gqs <- gsub('cmdstanv','gqs_shielding',file_stanModel)
stopifnot( file.exists(file_stanModel_gqs) )

cat('\nCompiling gqs model file ...')
m2 <- rstan::stan_model(file_stanModel_gqs)

cat('\nGenerating quantities ...')
draws <- as.matrix(fit)
draws <- draws[,!grepl('Rt|RtByAge|rho|lp__|E_deaths|E_deathsByAge|E_casesByAge', colnames(draws)) ]
fit2 <- rstan::gqs(m2, data=stan_data, draws=draws)
fit.gqs <- rstan::extract(fit2)

#	save stan_data
if(location.index==1)
{
  multiplier_name <- (multiplier_cntct_school_opening)*100
  tmp <- file.path(indir.results, paste0(basename(args$job_dir),'_stan_data', '_sensitivity_shielding_',as.integer(shield),'_school_reopen_', as.integer(school.reopen),'_multiplier_', multiplier_name, '_level_', school_level,'.RDS'))
  if(as.integer(school.reopen) == 0){
    tmp <- file.path(indir.results, paste0(basename(args$job_dir),'_stan_data', '_sensitivity_shielding_',as.integer(shield),'_school_reopen_', as.integer(school.reopen), '.RDS'))
  }
  saveRDS(stan_data, file=tmp)
}

#	save fit.gqs, try up to 5 times
multiplier_name <- (multiplier_cntct_school_opening)*100
tmp <- file.path(indir.results, paste0(basename(args$job_dir), '_location', location.index,'_stangqs', '_sensitivity_shielding_',as.integer(shield),'_school_reopen_', as.integer(school.reopen),'_multiplier_', multiplier_name, '_level_', school_level, '.RDS'))
if(as.integer(school.reopen) == 0){
  tmp <- file.path(indir.results, paste0(basename(args$job_dir), '_location', location.index,'_stangqs', '_sensitivity_shielding_',as.integer(shield),'_school_reopen_', as.integer(school.reopen), '.RDS'))
}
tmp2 <- 20
repeat
{
  cat('\nSave quantities to file ',tmp,'...')
  saveRDS(fit.gqs, file = tmp)
  check_if_saved <- try(readRDS(file=tmp))
  tmp2 <- tmp2-1
  if(!'try-error'%in%class(check_if_saved))
    break
  if(tmp2<=0)
  {
    stop('Failed to save 20 times')
    break
  }

}

cat('\nFinished generate-shielding-forecast.r ...')


findSigGenes <- function(Expr, Label, Method = "tTest", Directed = TRUE, FdrCut = 0.01, FDCut = 1){
## for a given Expression profile, use statistic approach to evaluate the significance of difference between case and control samples;
## Expr: a data frame, the expression profile to be dealed with, the rownames should be the IDs of genes;
## Method: a string, should be one of the "tTest" or "foldChange";
## Label: a vector of 0/1s, indicating the class of samples in the expression profile, 0 represents case,1 represents control;
## Directed: logical, if the the up or down regulated set should be distinguished;
## FdrCut: the fdr cutoff for T test;
## FDCut: the cutoff for fold change;

        if(Method == "tTest"){
                Tscore <- calT(Expr, Label);
                P_Val <- pt(abs(Tscore$T), lower.tail = FALSE, df = Tscore$df);
                Fdr <- p.adjust(P_Val);
                Result <- as.data.frame(cbind(rownames(Expr), Tscore$T, P_Val, Fdr));
                names(Result) <- c("ID", "t_statistic", "P", "FDR");
                save(Result, file = "tTestTable.rda");
                if(Directed == TRUE){
                        UpGeneSet <- Result[["ID"]][intersect(which(as.numeric(Result[["FDR"]]) < FdrCut), which(as.numeric(Result[["t_statistic"]]) > 0))];
                        LoGeneSet <- Result[["ID"]][intersect(which(as.numeric(Result[["FDR"]]) < FdrCut), which(as.numeric(Result[["t_statistic"]]) < 0))];
                        SigGenes <- list(UpGeneSet, LoGeneSet);
                        names(SigGenes) <- c("Up", "Low");
                        return(SigGenes);
                }else{
                        SigGenes <- Result[["ID"]][which(as.numeric(Result[["FDR"]]) < FdrCut)];
                        return(SigGenes);
                }
        }else if(Method == "foldChange"){
                Result <- calFD(Expr, Label);
                if(Directed == TRUE){
                        UpGeneSet <- Result[["ID"]][which(as.numeric(Result[["Log2FoldChange"]]) > FDCut)];
                        LoGeneSet <- Result[["ID"]][which(as.numeric(Result[["Log2FoldChange"]]) < (-FDCut))];
                        SigGenes <- list(UpGeneSet, LoGeneSet);
                        names(SigGenes) <- c("Up", "Low");
                        return(SigGenes);

                }else{
                        SigGenes <- Result[["ID"]][which(abs(as.numeric(Result[["Log2FoldChange"]])) > FDCut)];
                        return(SigGenes);
                }

        }
}


# A set of functions for preprocessing of MS2s
library(intervals)
library(CAMERA)
library(MSnbase)
library(xcms)
library(plyr)

ppmCal<-function(run,ppm)
{
  return((run*ppm)/1000000)
}

merge.spectra.group <- function(spectra, eppm, eabs, int.threshold) {
  if(is.na(spectra) || length(spectra) == 0) return(NULL)
  max.precursor.int <- -1
  rt.at.max.int <- NA
  start.rt <- attributes(spectra[[1]])$rt
  end.rt <- attributes(spectra[[1]])$rt
  all_unique_files<-paste(sapply(spectra,function(x){attributes(x)$fileName}),collapse = ";")
  mean.precursor.mz <- 0
  max.tic <- -1
  tandemms <- lapply(1:length(spectra), function(index) {
    # print(index)
    x<-spectra[[index]]
    # assign local variables
    cur.rt <- attributes(x)$rt
    cur.precursor.int <- attributes(x)$precursorIntensity
    cur.tic <- attributes(x)$tic
    cur.precursor.mz <- attributes(x)$precursorMz
    # assign global variables
    mean.precursor.mz <<- mean.precursor.mz + cur.precursor.mz
    max.tic <<- if(max.tic < cur.tic) cur.tic else max.tic
    if(max.precursor.int < cur.precursor.int) {
      max.precursor.int <<- cur.precursor.int
      rt.at.max.int <<- cur.rt
    }
    start.rt <<- if(start.rt > cur.rt) cur.rt else start.rt
    end.rt <<- if(end.rt < cur.rt) cur.rt else end.rt
    # start merging ms2 data
    a <- cbind(attributes(x)$mz, attributes(x)$intensity)
    colnames(a) <- c("mz", "intensity")
    return(a)
  })
  mean.precursor.mz <- mean.precursor.mz / length(spectra)

  # Filter intensity
  tandemms <- lapply(tandemms, function(spec) spec[spec[,"intensity"]>=int.threshold,])
  tandemrm <- sapply(1:length(tandemms), function(x) { if ( (dim(tandemms[[x]])[1] == 0) || is.null(dim(tandemms[[x]])[1]) ) x <- TRUE else x <- FALSE } )

  if ( (all(as.logical(tandemrm)) == TRUE) | (is.na(all(as.logical(tandemrm)))) )
    return(NULL)
  else
    tandemms <- tandemms[!tandemrm]

  # Process spectra
  peaks <- do.call(rbind, tandemms)
  g <- xcms:::mzClust_hclust(peaks[,"mz"], eppm=eppm, eabs=eabs)

  mz <- tapply (peaks[,"mz"], as.factor(g), mean)
  intensity <- tapply (peaks[,"intensity"], as.factor(g), max)

  if(length(mz) == 0)return(NULL);

  intensity <- intensity[order(mz)]
  mz <- mz[order(mz)]

  res <- .Call("Spectrum2_constructor_mz_sorted",
               attributes(spectra[[1]])$msLevel, length(mz), rt.at.max.int,
               attributes(spectra[[1]])$acquisitionNum,
               attributes(spectra[[1]])$scanIndex, max.tic, mz,
               intensity, attributes(spectra[[1]])$fromFile, attributes(spectra[[1]])$centroided,
               attributes(spectra[[1]])$smoothed, attributes(spectra[[1]])$polarity,
               attributes(spectra[[1]])$merged, attributes(spectra[[1]])$precScanNum,
               mean.precursor.mz, max.precursor.int, attributes(spectra[[1]])$precursorCharge, attributes(spectra[[1]])$collisionEnergy,
               TRUE, attributes(spectra[[1]])$.__classVersion__,
               PACKAGE = "MSnbase")
  attributes(res)$start.rt <- start.rt
  attributes(res)$end.rt <- end.rt
  attributes(res)$fileName <- all_unique_files

  return(res)
}

get.ppm.from.abs <- function(peak, ppm) {
  return ((peak / 1000000.0) * ppm)
}

collect.spectra.lists <- function(spectra, mzabs, mzppm, rtabs) {
  a<-t(sapply(1:length(spectra), function(spectrum.id) {
    return(c(attributes(spectra[[spectrum.id]])$rt, attributes(spectra[[spectrum.id]])$precursorMz, spectrum.id))
  }))
  a <- a[order(a[,2], a[,1]),]
  index <- 1
  grouped.spectra.tmp <- list()
  while(index <= dim(a)[1]) {
    spectra.list <- list()
    spectra.list[[1]] <- spectra[[a[index,3]]]
    cur.abs <- get.ppm.from.abs(a[index, 2], mzppm * 0.5) + mzabs
    index <- index + 1
    while(index <= dim(a)[1]) {
      diff.mz <- abs(a[index - 1, 2] - a[index, 2])
      if(diff.mz <= cur.abs) {
        spectra.list[[length(spectra.list) + 1]] <- spectra[[a[index,3]]]
        index <- index + 1
      } else break
    }
    grouped.spectra.tmp[[length(grouped.spectra.tmp) + 1]] <- spectra.list
  }
  grouped.spectra <- list()
  for(spectrum.group.index in 1:length(grouped.spectra.tmp)) {
    spectrum.group <- grouped.spectra.tmp[[spectrum.group.index]]
    a<-t(sapply(1:length(spectrum.group), function(spectrum.index) {
      c(attributes(spectrum.group[[spectrum.index]])$rt, attributes(spectrum.group[[spectrum.index]])$precursorMz, spectrum.index)
    }))
    a <- matrix(a[order(a[,1], a[,2]),], ncol=3)
    index <- 1
    while(index <= dim(a)[1]) {
      spectra.list <- list()
      spectra.list[[1]] <- spectrum.group[[a[index,3]]]
      index <- index + 1
      while(index <= dim(a)[1]) {
        diff.rt <- abs(a[index - 1, 1] - a[index, 1])
        if(diff.rt <= rtabs) {
          spectra.list[[length(spectra.list) + 1]] <- spectrum.group[[a[index,3]]]
          index <- index + 1
        } else break
      }
      grouped.spectra[[length(grouped.spectra) + 1]] <- spectra.list
    }
  }
  return(grouped.spectra)
}

preprocess_msms<-function(mapped_msms=NA,centroid=F,merge=TRUE,centroid_after_merge=F,ppm=10,ppm_precursor=10,abs_mz=0.01,abs_mz_precursor=0.01,rt=10,centroid_onlymapped=FALSE,merge_onlymapped=FALSE,int_threshold=0,verbose=FALSE)
{
  if(!is.logical(verbose)){
    stop("verbose must be TRUE or FALSE!")
  }

  if(verbose)cat("Checking inputs ...","\n")
  if(!is.list(mapped_msms))
  {
    stop("mapped_msms has to be a list from map_features")
  }else{
    if(!all(names(mapped_msms)==c("mapped","unmapped"))){
      stop("mapped_msms does not appear to be map_features! mapped_msms has to be a list from map_features!")
    }
    if(all(sapply(mapped_msms, length)<1)){
      stop("No MS2s were found in mapped_msms")
    }
  }

  if(!is.logical(centroid)){
    stop("centroid must be TRUE or FALSE!")
  }

  if(!is.logical(merge)){
    stop("merge must be TRUE or FALSE!")
  }

  if(!is.logical(centroid_after_merge)){
    stop("centroid_after_merge must be TRUE or FALSE!")
  }

  if(!is.logical(merge_onlymapped)){
    stop("merge_onlymapped must be TRUE or FALSE!")
  }

  if(!is.logical(centroid_onlymapped)){
    stop("centroid_onlymapped must be TRUE or FALSE!")
  }

  if(any(is.na(ppm)) | any(is.null(ppm)))
  {
    stop("No ppm input have been provided!")
  }

  if(!is.numeric(rt))
  {
    stop("rt must be numberic")
  }

  if(!is.numeric(int_threshold))
  {
    stop("int_threshold must be numberic")
  }

  if(!is.numeric(ppm_precursor))
  {
    stop("ppm_precursor must be numberic")
  }

  if(!is.numeric(abs_mz_precursor))
  {
    stop("abs_mz_precursor must be numberic")
  }

  if(abs_mz_precursor<0)
  {
    stop("abs_mz_precursor must be numberic and higher than 0")
  }

  if(!is.numeric(abs_mz))
  {
    stop("abs_mz must be numberic")
  }

  if(!is.numeric(ppm))
  {
    stop("ppm must be numberic")
  }

  if(ppm<0)
  {
    stop("ppm must be numberic and higher than 0")
  }

  if(rt<0)
  {
    stop("rt must be numberic and higher than 0")
  }

  if(abs_mz<0)
  {
    stop("abs_mz must be numberic and higher than 0")
  }
  if(ppm_precursor<0)
  {
    stop("ppm_precursor must be numberic and higher than 0")
  }

  if(centroid==TRUE){
    if(verbose)cat("Centroiding ...","\n")
    for(i in 1:length(mapped_msms$mapped))
    {
      for(j in 1:length(mapped_msms$mapped[[i]]))
      {
        MSMS<-mapped_msms$mapped[[i]][[j]]
        MSMS@centroided<-F
        MSMS@polarity<-as.integer(1)
        MSMS@smoothed<-F
        MSMS<-MSnbase::pickPeaks(MSMS)
        mapped_msms$mapped[[i]][[j]]<-MSMS
      }
    }

    if(!centroid_onlymapped){
      for(i in 1:length(mapped_msms$unmapped))
      {

          MSMS<-mapped_msms$unmapped[[i]]
          MSMS@centroided<-F
          MSMS@polarity<-as.integer(1)
          MSMS@smoothed<-F
          MSMS<-MSnbase::pickPeaks(MSMS)
          mapped_msms$unmapped[[i]]<-MSMS

      }
    }
  }

  if(merge==TRUE)
  {
    if(verbose)cat("Creating hyperMS2 for mapped ions...","\n")
    for(i in 1:length(mapped_msms$mapped))
    {
      if(length(mapped_msms$mapped[[i]])>1)
      {
        merge_msms<-merge.spectra.group(mapped_msms$mapped[[i]],eppm = ppm* 10e-6,eabs = abs_mz_precursor,int.threshold = int_threshold)
        if(centroid_after_merge==TRUE)
        {
          MSMS<-merge_msms
          MSMS@centroided<-F
          MSMS@polarity<-as.integer(1)
          MSMS@smoothed<-F
          MSMS<-MSnbase::pickPeaks(MSMS)
          merge_msms<-MSMS
        }
        mapped_msms$mapped[[i]]<-list(merge_msms)
      }
    }

    if(merge_onlymapped==FALSE)
    {
      a <- collect.spectra.lists(mapped_msms$unmapped, abs_mz, ppm_precursor, rt)
      merged.spectra <- sapply(1:length(a), function(x) {
        merge.spectra.group(a[[x]], ppm * 10e-6, abs_mz, int_threshold)
      })

      if(centroid_onlymapped==FALSE & centroid_after_merge==TRUE)
      {
        list_of_unmapped<-c()
        for(i in 1:length(merged.spectra))
        {

          MSMS<-merged.spectra[[i]]
          MSMS@centroided<-F
          MSMS@polarity<-as.integer(1)
          MSMS@smoothed<-F
          MSMS<-MSnbase::pickPeaks(MSMS)
          list_of_unmapped[[i]]<-MSMS

        }
        mapped_msms$unmapped<-list_of_unmapped
      }

    }

  }

  return(mapped_msms)

}


# -*- coding: utf-8 -*-

##################################################################################
# Ugalde et al a la Bonica
# This code runs a Bonica-like algorithm to provide a dynamic view of Ugalde's
# IFE.  The run has already been saved ("DynUgaldeBonica.RData")
# March 19, 2013: Add party-based priors for councilors coming in from outside
# March 21, 2013: Added more precise priors for new entrants:
#	PRI-sponsored=N( 2,4)
#	PAN-sponsored=N(-1,4)
#	PRD-sponsored=N(-2,4)
##################################################################################

library(arm)
library (MCMCpack)
library (foreign)
library (car)
library (gtools)
#library (multicore)
library (R2jags)
library (mcmcplots)
library (sm)
library(lubridate)

rm(list = ls())
workdir <- c("/home/eric/Dropbox/data/rollcall/ife_cg/ife-update/data/")
setwd(workdir)

# Define colors and plotting names
# OJO: en tenure term==10 es a, term==11 es b etc.
ids <- matrix(c("Ugalde",           "ugalde",      "PRI",  "4",
                "Albo",             "albo",        "PAN",  "456",
                "Andrade",          "andrade",     "PRI",  "4567",
                "Gmz. Alcántar",    "alcantar",    "PVEM", "4567",
                "Glez. Luna",       "glezluna",    "PAN",  "456",
                "Latapí",           "latapi",      "PRI",  "45",
                "López Flores",     "lopezflores", "PRI",  "456",
                "Morales",          "morales",     "PAN",  "45",
                "Sánchez",          "sanchez",     "PAN",  "4567c",
                "Valdés",           "valdes",      "PRD",    "6789a",
                "Baños",            "banos",       "PRI",    "6789abcde",
                "Nacif",            "nacif",       "PAN",    "6789abcde",
                "Elizondo",         "elizondo",    "PAN",     "789a",
                "Figueroa",         "figueroa",    "PRD",     "789a",
                "Guerrero",         "guerrero",    "PRI",     "789a",
                "Córdova",          "cordova",     "PRD",       "9abcdef",
                "García Rmz.",      "garcia",      "PRI",       "9"  ,
                "Marván",           "marvan",      "PAN",       "9ab",
                "E. Andrade",       "andrade2",    "",             "cde",
                "Favela",           "favela",      "",             "cdef",
                "Santiago",         "santiago",    "",             "c",
                "Galindo",          "galindo",     "",             "c",
                "Murayama",         "murayama",    "",             "cdef",
                "Ruiz Saldaña",     "ruiz",        "",             "cdef",
                "San Martín",       "snmartin",    "",             "cde",
                "Santiago",         "santiago",    "",             "c",
                "Ravel",            "ravel",       "",              "def",
                "J. Rivera",        "rivera2",     "",              "def",
                "Zavala",           "zavala",      "",              "def",
                "De la Cruz",       "magana",      "",                "f",
                "Faz",              "faz",         "",                "f",
                "Humphrey",         "humphrey",    "",                "f",
                "Kib Espadas",      "kib",         "",                "f"),
              ncol = 4,
              byrow = TRUE)
#
ids <- as.data.frame(ids, stringsAsFactors = FALSE)
colnames(ids) <- c("name", "column", "pty", "tenure")
#ids$tenure <- as.numeric(ids$tenure)
ids <- within(ids, party <- ifelse (pty=="PRI", 1,
                            ifelse (pty=="PAN", 2,
                            ifelse (pty=="PRD", 3,
                            ifelse(pty=="PVEM", 4, 5)))))
ids <- within(ids, color <- ifelse (pty=="PRI", "red",
                            ifelse (pty=="PAN", "blue",
                            ifelse (pty=="PRD", "gold",
                            ifelse(pty=="PVEM", "green", "orangered4")))))

# select terms 4-8, more or less
tees <- 6:7
T <- length(tees)

sel    <- grep(pattern = "[67]", ids$tenure) #sel    <- grep(pattern = "[456789ab]", ids$tenure)
name   <- ids$name[sel]
party  <- ids$party[sel]
color  <- ids$color[sel]
column <- ids$column[sel]

## rgb.23 <- c(length=11)
## rgb.23[c(1,6,8,10)] <- rgb(1,       0, 0, 0.6) #red
## rgb.23[c(2:4,9)]    <- rgb(1, 215/255, 0, 0.6) #gold
## rgb.23[c(5,7,11)]   <- rgb(0,       0, 1, 0.6) #blue

###############################################################################
## Read votes (includes informative votes only, exported by code/data-prep.r ##
###############################################################################
vot <-read.csv("v456789ab.csv",  header=TRUE)
#
# subset to chosen periods
sel.r <- which(vot$term %in% tees)
drop.c <- ids$column[sel] # column names not in terms selected
drop.c <- setdiff(ids$column, drop.c)
drop.c <- which(colnames(vot) %in% drop.c)
if (length(drop.c)>0) vot <- vot[sel.r, -drop.c]
colnames(vot)
# total members
J <- length(name); J

########################
## recode vote values ##
########################
vs <- vot[,1:J]
#table(v$albo, useNA = "always")
vs[vs==0] <- NA    ## Version probit requiere 0s y 1s
vs[vs>2] <- NA
vs[vs==2] <- 0

# format dates
vot$date <- ymd(vot$date)
# summarize then drop uncontested votes
table(factor(vot$dunan, labels = c("contested","not")), vot$term, useNA = "ifany")
table(factor(vot$dunan, labels = c("contested","not")), useNA = "ifany")
sel <- which(vot$dunan==1)
vot <- vot[-sel,] # drop uncontested votes
vs  <- vs [-sel,] # drop uncontested votes

#############################
###     UGALDE ET AL      ###
#############################

## MODEL
model1Dj.irt <- function() {
	for (j in 1:J){                ## loop over respondents
		for (i in 1:I){              ## loop over items
			v[j,i] ~ dbern(p[j,i]);                                 ## voting rule
			probit(p[j,i]) <- mu[j,i];                              ## sets 0<p<1 as function of mu
			mu[j,i] <- signal[i]*x[j] - difficulty[i];              ## utility differential
		}
	}
	## priors ################
	for (j in 1:J){
		x[j] ~ dnorm (x.mean[j], x.tau[j]);
	}
	for (i in 1:I){
		signal[i] ~ dnorm(0, 0.1);
		difficulty[i] ~ dnorm(0, 0.25);
	}
	for (p in 1:4){ # need 5 when morena also considered
		partyPos[p] <- mean (x[sponsors[p]]); # 4mar21: should be median, unknown function in bugs?
	}
}
#end model##############

I <- nrow (vot)
## # Center on vote (for date), extend windows to both sides
## item <- 1:I
## inicio <- item-15; inicio[inicio<0] <- 1
## final  <- item+15; final[final>I] <- I
## item.date <- vot$date #ymd(vot$yr*10000+vot$mo*100+vot$dy)
## S <- length(inicio)

# We need a matrix showing whether each councilor is actually in IFE the moment the vote takes place --- will be subset to J below
IsCouncilor <- matrix (1, ncol=18, nrow=I)
IsCouncilor[ vot$term < 4 | vot$term >  4,1 ] <- NA  #      ugalde
IsCouncilor[ vot$term < 4 | vot$term >  6,2 ] <- NA  #        albo
IsCouncilor[ vot$term < 4 | vot$term >  7,3 ] <- NA  #     andrade
IsCouncilor[ vot$term < 4 | vot$term >  7,4 ] <- NA  #    alcantar
IsCouncilor[ vot$term < 4 | vot$term >  6,5 ] <- NA  #    glezluna
IsCouncilor[ vot$term < 4 | vot$term >  5,6 ] <- NA  #      latapi
IsCouncilor[ vot$term < 4 | vot$term >  6,7 ] <- NA  # lopezflores
IsCouncilor[ vot$term < 4 | vot$term >  5,8 ] <- NA  #     morales
IsCouncilor[ vot$term < 4 |(vot$term >  7 & vot$term!=12),9 ] <- NA  # sanchez
IsCouncilor[ vot$term < 6 | vot$term > 10,10] <- NA  #      valdes
IsCouncilor[ vot$term < 6 | vot$term > 14,11] <- NA  #       banos
IsCouncilor[ vot$term < 6 | vot$term > 14,12] <- NA  #       nacif
IsCouncilor[ vot$term < 7 | vot$term > 10,13] <- NA  #    elizondo
IsCouncilor[ vot$term < 7 | vot$term > 10,14] <- NA  #    figueroa
IsCouncilor[ vot$term < 7 | vot$term > 10,15] <- NA  #    guerrero
IsCouncilor[ vot$term < 9 | vot$term > 15,16] <- NA  #     cordova
IsCouncilor[ vot$term < 9 | vot$term >  9,17] <- NA  #  garcia rmz
IsCouncilor[ vot$term < 9 | vot$term > 11,18] <- NA  #      marvan
####################################################
## vector to select members present in estimation ##
####################################################
sel.members <- which(ids$name %in% name)
IsCouncilor <- IsCouncilor[, sel.members]

# Round 1 ideal points to anchor ideological space
# (later entrants at party mean)
#                 1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
#                 u  a  a  a  g  l  l  m  s  v  b  n  e  f  g  c  g  m  a  f  g  m  r  s  s
#                 g  l  n  l  l  a  p  o  a  a  a  a  l  i  u  o  a  a  n  a  a  u  u  n  a
#                 a  b  d  c  z  t  z  r  n  l  ñ  c  i  g  e  r  r  r  d  v  l  r  i  m  n
#                 l  o  r  a  l  a  f  a  c  d  o  i  z  u  r  d  c  v  r  e  i  a  z  a  t
x.location <-   c(1, 0, 0, 2,-2, 0, 0, 2,-2, 0, 2,-1,-2, 0, 2,-2, 2,-2, 0, 0, 0, 2, 0, 0, 0)[sel.members] # sel = members in estimation
x.precision  <- c(4, 1, 1, 4, 4, 4, 1, 4, 4, 1, 4, 4, 4, 1, 4, 4, 4, 4, 1, 1, 1, 4, 1, 1, 1)[sel.members]
# x.location and x.precision are amipulated by loop, keep a version to use as prior for new entrants instead of party mean
x.prior.location <-  x.location
x.prior.precision <- x.precision
#
term <- vot$term
term.results <- list () ## WILL ADD TERM'S RESULTS TO OBJECT HOLDING ALL RESULTS
partyPlacement <- rep (NA,J)
x.mean <- numeric ()
x.tau  <- numeric ()

## Save overall totals for use later (I J redefined to window s totals in next loop)
J.all <- J; I.all <- I
for (t in tees){        # <= BIG FUNCTION STARTS (loop over terms with constant membership)
#
	# Added March 19: We include councilors (and their party IDs) only if they were actual councilors for at least one vote
	# This means that the length of estimated ideal points is either
	# 9 (for most votes) or 11 (when there is some overlap: two councilors are leaving , two are coming in)
	councilor.in <- apply (IsCouncilor[term==t,], 2, invalid)
	councilors   <- name [councilor.in==FALSE]
	sponsors     <- party[councilor.in==FALSE]
#
	for (c in 1:J.all){
# 5mar21: this uses partyPlacement for new entrants
		x.mean[c] <- ifelse (!is.na(x.location [c]), x.location [c],  partyPlacement[sponsors[c]])
		x.tau[c]  <- ifelse (!is.na(x.precision[c]), x.precision[c], 4)
# 5mar21: this uses fixed x0 prior for new entrants
#		x.mean[c] <- ifelse (!is.na(x.location [c]), x.location [c], x.prior.location [c])
#		x.tau[c]  <- ifelse (!is.na(x.precision[c]), x.precision[c], x.prior.precision[c])
	}
#
	v <- vs[term==t, 1:J.all][, councilor.in==FALSE]; ## EXTRACT ALL VOTES FROM TERM t
	v <- t(v)                       ## ROLL CALLS NEED ITEMS IN COLUMNS, LEGISLATORS IN ROWS
	J <- nrow(v); I <- ncol(v)      ## SESSION TOTALS
#
	ife.data <- list ("J", "I", "v", "x.mean", "x.tau", "sponsors") # emm 9mar21: no debería ir sponsors en vez de party?
	ife.inits <- function (){
		list (
			x=rnorm(J),
			signal=rnorm(I),
			difficulty=rnorm(I)
		)
	}
	ife.parameters <- c("x", "signal", "difficulty", "partyPos")
#
	print(cat("Term no.", t,", with", I, "votes \n"))
#
	#full JAGS run
	start.time <- proc.time()
#
	# Use dual core capabilities
	results <-
#            mclapply(1:2, function(x) {
#		model.jags.re <- try(
                                 jags (data=ife.data, inits=ife.inits, ife.parameters,
#					model.file=model1Dj.irt, n.chains=1,
					model.file=model1Dj.irt, n.chains=2,
#					n.iter=500, n.burnin=250, n.thin=25)
					n.iter=50000, n.burnin=30000, n.thin=200)
#		)
#		if(inherits(model.jags.re,"try-error")) {return()}
#		return(model.jags.re)
#	}, mc.cores=2 )
	time.elapsed <- round(((proc.time()-start.time)[3])/60,2); rm(start.time)
	print(cat("\tTime elapsed in estimation:", time.elapsed, "minutes", "\n")); rm(time.elapsed)
#
	# Quick check on convergence of ideal point chains
#	GHconv <- gelman.diag(mcmc.list(list (as.mcmc (results[[2]]$BUGSoutput$sims.list$x), as.mcmc (results[[1]]$BUGSoutput$sims.list$x))))[[2]]
#	print (cat ("Gelman-Rubin R-hat:", GHconv, "\n"))
#
	# ADD COUNCILOR NAMES AND VOTE INFO TO RESULTS OBJECT
        results <- c(results, councilors=list(councilors)); # should be faster than results[[length(results)+1]] <- councilors;
        results <- c(results, term=t); # add term
        term.results <- c(term.results, list(results)); ## ADD TERM'S RESULTS TO OBJECT HOLDING ALL RESULTS
#        results[[4]] <- GHconv; rm (GHconv)
#
	# Update location of ideal point at time s, to be used as location prior at time s+1
	x.location  <- rep (NA, J.all)
	x.precision <- rep (100, J.all)
#	locs <- apply( rbind (results[[1]]$BUGSoutput$sims.list$x, results[[2]]$BUGSoutput$sims.list$x), 2, median)
#	partyPlacement <- apply( rbind (results[[1]]$BUGSoutput$sims.list$partyPos, results[[2]]$BUGSoutput$sims.list$partyPos), 2, median)
	locs <- apply( results$BUGSoutput$sims.list$x, 2, median)
	partyPlacement <- apply( results$BUGSoutput$sims.list$partyPos, 2, median)
	for (n in 1:J.all){
		if (length( which(councilors==name[n]) )==0) {               # if councilor not member current round
			x.location[n] <- NA                                  # then prior for next round set to NA
			x.precision[n] <- NA                                 # (and line above sets it to party placement)
		}
		else { x.location[n] <-  locs[which (councilors==name[n])] } # councilor's prior for next round is current x
	}
	# Precision prior is always constant at 100, implying standard deviation = sqrt (1/100) = 0.1
}  # <---   END OF LOOP OVER WINDOWS
#
## Restore overall totals
J <- J.all; I <- I.all; rm(J.all, I.all)

# plot results
ideal.points <- as.data.frame(matrix (NA, nrow=T, ncol=J))
colnames(ideal.points) <- ids$column[sel.members]
name <- ids$name[sel.members]
column <- ids$column[sel.members]
#color <- ids$color[sel.members] # party colors
library(RColorBrewer)
color <- brewer.pal(n = J, name = "Paired")
for (i in 1:T){
#	i <- 2 # debug
	sel <- which(name %in% term.results[[i]]$councilors) # select columns with member councilors in window sessions
#	ideal.points[i,sel] <- apply (rbind (window.results[[i]][[1]]$BUGSoutput$sims.list$x, window.results[[i]][[2]]$BUGSoutput$sims.list$x), 2, median) # merge both chains, report median
	ideal.points[i,sel] <- apply (term.results[[i]]$BUGSoutput$sims.list$x, 2, median) # merge both chains, report median
}
sel <- which(!is.na(ideal.points[1,]))
plot(rep(6,J)[sel], ideal.points[1,sel], col = color[sel], pch = 19)
sel <- which(!is.na(ideal.points[2,]))
points(rep(7,J)[sel], ideal.points[2,sel], col = color[sel], pch = 19)

# rename object with posterior sims
summary(term.results[[1]])
term.results.67 <- term.results
rm(term.results)

# clean
ls()
rm(c, s, n, v, sel, ife.inits, ife.parameters, ife.data)
rm(councilors, sponsors, inicio, final, councilor.in)
rm(x.location, x.mean, x.precision, x.tau, item, results, item.date)
rm(color, column, locs, name, party, partyPlacement)

# save
summary(term.results.67)
#summary(window.results[[232]]) # 11 members, overlap
save.image(file = "posterior-samples/vald67-term-results-compress.RData", compress = "xz")
#save(window.results.23, file = "posterior-samples/wold23-window-results-compress.RData")
x

# Save window.results, containing all chains from all runs
# save (window.results, file="DynUgaldeBonicaMarch21.RData") # With more anchors
# save (window.results, file="DynUgaldeBonicaMarch19.RData")
# save (window.results, file="DynUgaldeBonica.RData")

source("http://rtm.wustl.edu/code/sendEmail.R")
#sendEmail (subject="Ugalde et al esta listo", text="", address="rosas.guillermo@gmail.com")
sendEmail (subject="Ugalde et al esta listo", text="", address="emagar@gmail.com")

# RData file with runs carried out in Mexico, early March
# load ("DynUgaldeBonica.RData")

# RData file with runs carried out in Wash U, March 19
# These runs omit non-sitting Councilors and party precisions for new Councilors
load ("DynUgaldeBonicaMarch21.RData")

S <- length (semester.results)
multiGelman.hat <- numeric ()
for (i in 1:S){
	chainsConv <- mcmc.list(list (as.mcmc (semester.results[[i]][[2]]$BUGSoutput$sims.list$x), as.mcmc (semester.results[[i]][[1]]$BUGSoutput$sims.list$x)))
	tmp <- gelman.diag (chainsConv)[[2]]
	multiGelman.hat <- c(multiGelman.hat, tmp)
}
rm (tmp, chainsConv)

nonConverged <- ifelse (multiGelman.hat > 2, 1, 0)  # Should be 0

CouncilorIn <- matrix (1, nrow=J, ncol=S)
CouncilorIn[1,  all45678901$date > 20071217] <- NA
CouncilorIn[2,  all45678901$date > 20080814] <- NA
CouncilorIn[3,  all45678901$date > 20101027] <- NA
CouncilorIn[4,  all45678901$date > 20101027] <- NA
CouncilorIn[5,  all45678901$date > 20080814] <- NA
CouncilorIn[6,  all45678901$date > 20071217] <- NA
CouncilorIn[7,  all45678901$date > 20080814] <- NA
CouncilorIn[8,  all45678901$date > 20071217] <- NA
CouncilorIn[9,  all45678901$date > 20101027] <- NA
CouncilorIn[10,  all45678901$date < 20080215] <- NA
CouncilorIn[11,  all45678901$date < 20080215] <- NA
CouncilorIn[12,  all45678901$date < 20080215] <- NA
CouncilorIn[13,  all45678901$date < 20080829] <- NA
CouncilorIn[14,  all45678901$date < 20080829] <- NA
CouncilorIn[15,  all45678901$date < 20080829] <- NA
#CouncilorIn[16,  all45678901$date < 20111215] <- NA
#CouncilorIn[17,  all45678901$date < 20111215] <- NA
#CouncilorIn[18,  all45678901$date < 20111215] <- NA

# If using DynUgaldeBonica.RData, use the following code to extract ideal points
ideal.points <- matrix (NA, nrow=S, ncol=J)
for (i in 1:S){
	ideal.points[i,] <- apply (rbind (semester.results[[i]][[1]]$BUGSoutput$sims.list$x, semester.results[[i]][[2]]$BUGSoutput$sims.list$x), 2, median)
}

# If using DynUgaldeBonicaMarch19.RData or DynUgaldeBonicaMarch21.RData, use the following code to extract ideal points
ideal.points <- matrix (NA, nrow=S, ncol=J)
ideal.points.var <- matrix (NA, nrow=S, ncol=J)
for (i in 1:S){
	for (j in 1:J){
		councilor <- names45678901[j]
		num <- which (semester.results[[i]][[3]]==councilor)
		if ( length (num)==0 ) {
			ideal.points[i,j] <- 1
			ideal.points.var[i,j]  <- 0
		} else {
				ideal.points[i,j] <- median (c (semester.results[[i]][[1]]$BUGSoutput$sims.list$x[,num], semester.results[[i]][[2]]$BUGSoutput$sims.list$x[,num]))
				ideal.points.var[i,j]  <- var (c (semester.results[[i]][[1]]$BUGSoutput$sims.list$x[,num], semester.results[[i]][[2]]$BUGSoutput$sims.list$x[,num]))
		}
	}
}

# Plot different individual councilor paths:
# There is definitely a problem at different spots of the estimation:
# Between i=200 and i=500 and after i=600
# The problem really obtains when we get new guys in
# Valdes, Banos, Nacif, Elizondo, Figueroa, Guerrero come in in quick succession and destroy everything
# Maybe the best solution is to break down estimation of this very long period in two or three chunks
# Or maybe we need to combine priors on ideal points with priors on item parameters
k <-  12
plot (ideal.points[ideal.points.var[,k] != 0, k], type="l")

# Get SDs of estimates, the width should be useful to plot thickness of data points
ideal.points.var <- sqrt (ideal.points.var)

# Non-smoothed ideal point time-paths
setwd ("/Users/grosas/Dropbox/ifesharedge/graphs")
pdf ("UgaldeEtAlNonSmooth.pdf", h=7, w=9)
plot(c(1:S), ideal.points[1:S,1], main="", ylim=c(-3,3), type="n", xlab="", ylab="Ideal points")
for (j in 1:J){
# 	lines(ideal.points[1:S,j], lwd=3, col=color45678901[j])
	lines(CouncilorIn[j,1:S] * ideal.points[1:S,j], lwd=3, col=rgb.45678901[j])
}
dev.off()

# Smoothed ideal point time-paths
# pdf ("UgaldeEtAlSmooth.pdf", h=7, w=9)
pdf ("UgaldeEtAlSmoothMarch21.pdf", h=7, w=9)
plot(c(1:S), ideal.points[1:S,1], main="", ylim=c(-3,3), type="n", xlab="", ylab="Ideal points")
for (j in 1:J){
	lines(smooth.spline(c(1:S)[!is.na(CouncilorIn[j,1:S])], ideal.points[!is.na(CouncilorIn[j,1:S]),j], df=10), lwd=3, col=rgb.45678901[j])
# 	lines(smooth.spline(c(1:S), ideal.points[,j], df=10), lwd=3, col=color45678901[j])
}
dev.off()

setwd ("/Users/grosas/Dropbox/ifesharedge/graphs/animBits/ugalde/")

# Smoothed ideal point time-paths, daily change movie
Smooth <- list ()
for (j in 1:J){
	Smooth[[j]] <- smooth.spline(c(1:S), ideal.points[,j], df=10)
}

for (j in 1:J){
	Smooth[[j]]$y[is.na(CouncilorIn[j,1:S])] <- NA
}

snapshot <- seq (1, 1091, by=10)
for (s in snapshot){
	which.s <- which (snapshot==s)
	if (which.s < 10) {name = paste('Ugalde00', which.s,'plot.png',sep='')}
	if (which.s >= 10 && which.s < 100) {name = paste('Ugalde0', which.s,'plot.png', sep='')}
	if (which.s >= 100) {name = paste('Ugalde', which.s,'plot.png', sep='')}

	jpeg (name, quality=100, height=500, width=500)

	plot(c(1:S), ideal.points[1:S,1], main="", ylim=c(-3,3), type="n", xlab="", ylab="Ideal points")
	for (j in 1:J){
		lines ( Smooth[[j]]$x[1:s], Smooth[[j]]$y[1:s], lwd=6, col=rgb.45678901[j])
	}
	# Eric: Creo que tú tienes más información para hacer los cambios en las siguientes tres líneas
	legend ("topright", bty="n", legend=paste ("vote", s, sep=" ")) #Change the legend for the date of the vote
	# Add vertical lines for elections here, but only for some realizations of S
	# Add names of councilors here as well, but only for some realizations of S
	dev.off()
}

# Make a short film
# system("convert -loop 1 -delay 40 *.png IFEwolden.gif")
system("convert -loop 1 -delay 20 *.png IFEugaldeSofis.gif")
# The number after loop controls the number of automatic replays (0 stands for infinite loop)
# The number after delay determines length of transition between slides)


HPC<-TRUE   # use High Performance Computer (on Unix only)

############################
## North Sea or Baltic Sea runs
my.area<-c('North Sea','Baltic Sea')[1]
NS.Key.2014<-T  # use the NS key run from 2014 (more species in the 2014 run, with more species for output figures)
do.MSFD.indicators<-FALSE
use.stochastic.recruitement<-T
doHessian<-F
do.sensitivity<-FALSE

#######################
source(file.path(prog.path.func,"hcr_op_batch_optimize_function.r"))

  #                     Cod     Whiting     Haddock      Saithe     Herring      NorthSandeel  southSan     Nor. pout       Sprat      Plaice        Sole
  HCR1<-            c(    1,          1,          1,          1,          1,          1,           1,          1,          1,          1,          1)
  HCR2<-            c(    2,          2,          2,          2,          2,         22,          22,         22,         22,          2,          2)
  my.FBlim.adjust<- c(    0,          0,          0,          0,          0,          1,           1,          1,          1,          0,          0)  # for specifying T1 in HCR (0: F=0 at SSB at 0, 1: F=0 at SSB at Blim)

  my.Fpa.mult1<-     c(    1,          1,          1,          1,          1,          1,          1,          1,          1,          1,          1)  # Maximum F relative to Fpa
  my.Fpa.mult2<-     c(    1,          1,          1,          1.5,        1,          1,          1,          1,          1,          1,          1)  # to let Saithe F be greater than Fpa (0.4)
  my.Fpa.mult3<-     c(    3,          3,          3,          1,          3,          3,          3,          3,          3,          3,          3)

my.Fpa.mult.all<-matrix(c(my.Fpa.mult1,my.Fpa.mult2,my.Fpa.mult3),byrow=T,nrow=3)
HCRs<-matrix(c(HCR1,HCR2),byrow=TRUE,nrow=2)

if (do.sensitivity) {  # Sensistivity of the stock size of othet predators
  #
  #Other species multiplier

  mult<-0.8
  my.other<-c(
  #1    2  3     4     5     6     7    8     9   10
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Fulmar
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Guillemot
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Her. Gull
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Kittiwake
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #GBB. Gull
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Gannet
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Puffin
  1, mult, 1,    1,    1,    1,    1,   1,    1,   mult,   #Razorbill
  1, 1,    mult, 1,    1,    1,    1,   1,    1,   mult,   #R. radiata
  1, 1,    1,    mult, 1,    1,    1,   1,    1,   mult,   #G. gurnards
  1, 1,    1,    1,    mult, 1,    1,   1,    1,   mult,   #W. mackerel
  1, 1,    1,    1,    mult, 1,    1,   1,    1,   mult,   #N. mackerel
  1, 1,    1,    1,    1,    mult, 1,   1,    1,   mult,   #W. horse mac
  1, 1,    1,    1,    1,    mult, 1,   1,    1,   mult,   #N. horse mac
  1, 1,    1,    1,    1,    1,    mult,1,    1,   mult,   #Grey seal
  1, 1,    1,    1,    1,    1,    1,   mult, 1,   mult,   #H. porpoise
  1, 1,    1,    1,    1,    1,    1,   1,    mult,mult)   #Hake

  fac.other.first.year<-2015  # first year for change (must be larger than first year in predeiction)
  fac.other.last.year<-2015  # last year for change

  my.other<-matrix(my.other,nrow=first.VPA-1,byrow=TRUE);
  dimnames(my.other)[[1]]<-sp.names[1:(first.VPA-1)]
  my.other
  #my.other<-my.other[,1:3]
  my.other<-my.other^(1/(fac.other.last.year-fac.other.first.year+1))
  n.oth<-(1:dim(my.other)[[2]])
} else {
  fac.other.first.year<- -1  # first year for change (must be larger than first year in predeiction)
  fac.other.last.year<- -1  # last year for change
  my.other<-matrix(1,nrow=first.VPA-1,ncol=2);
  n.oth<-1
}

###############  user input

just.batch.file<-FALSE   # make a batch file for a later run
just.plot.results<-FALSE   # just plot results from a previous run

nsp.VPA<-sum(SMS.control@species.info[,'predator'] %in% c(0,1))
recruit.adjust.factor<-rep(1,nsp.VPA)  # default no adjustment factor

if (use.stochastic.recruitement) {
  my.first.year<-2024
  my.last.year<-2124
  stochastic.recruitment<-1  # Stochastic recruitment?.
                          # 1=Yes   S-R parameters and variance from op_config.dat
                          # 2=Yes,  S-R parameters from op_config.dat and variance-covariance matrix from file covariance_Rec.in
                          # 3=Yes,  S-R parameters from op_config.dat and residuals from op_ssb_rec_residuals.in
                           # 21=Yes, S-R parameters from EquiSim results, file op_eqsim_stoch.in
  my.recruit.adjust.CV<-0    # adjust recruitment with half of the variance (factor exp((CV^2)/2) option adjust.recruit.CV
                             # 0=no adjustment
                             # 1=stochastic recruitment, do adjustment such that average recruitment becomes equal to the median (downward adjustment). Close to meaningless
                             # 2=determenistic recruitment, do ajustment such that recruitment becomes equal to the mean of the log-normal distribution (upward adjustment)
} else {   # deterministic
  my.first.year<-2034 # years for optimization
  my.last.year<-2043
  #recruit.adjust.factor[3]<-0.75  #addjust haddock recruitment
  stochastic.recruitment<-20          # Stochastic recruitment?.
                # 0=No,   S-R parameters from op_config.dat
                # 20=No,  S-R parameters from EquiSim results, file op_eqsim.in

  my.recruit.adjust.CV<-c(2) #adjust recruitment with half of the variance (factor exp((CV^2)/2) option adjust.recruit.CV
                            # 0=no adjustment
                            # 1=stochastic recruitment, do adjustment such that average recruitment becomes equal to the median (downward adjustment). Close to meaningless
                            # 2=determenistic recruitment, do ajustment such that recruitment becomes equal to the mean of the log-normal distribution (upward adjustment)
}

# FPAmult is multiplier for Fpa to define maximum F

for (FPAmult in (c(1))) for (adjRec in my.recruit.adjust.CV) for (my.HCR in c(1)) for (my.oth in n.oth) {
  my.Fpa.mult<-as.vector(my.Fpa.mult.all[FPAmult,])    # upper limit for Fmsy relative to Fpa
  HCR.used<-HCRs[ifelse(my.HCR==3,2,my.HCR),]  # do not change
  if (my.HCR==2)  FBlim.adjust<-my.FBlim.adjust else FBlim.adjust<-1

   #  scenario        x.a     x.b     x.c   x.d    x.e    x.f      where x=HCR
   #----------------------------------------------------------
   #  FMSY<=Fpa       yes    yes     yes    yes    yes     yes
   #  ssb penalty     no     Blim    Bpa    no     Blim    Bpa
   #  price at age    yes    yes     yes    no     no      no
   #----------------------------------------------------------

  #value
  #OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"a",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=T,stochastic.recruitment=stochastic.recruitment,use.penalty=F,HCR=HCR.used,penalty='Blim',penalty.type=0,Hessian=doHessian, Fpa.Mult=my.Fpa.mult, recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results,fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))
  #OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"b",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=T,stochastic.recruitment=stochastic.recruitment,use.penalty=T,HCR=HCR.used,penalty='Blim',penalty.type=1,Hessian=doHessian, Fpa.Mult=my.Fpa.mult, recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results,fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))
  #OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"c",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=T,stochastic.recruitment=stochastic.recruitment,use.penalty=T,HCR=HCR.used,penalty='Bpa', penalty.type=1,Hessian=doHessian, Fpa.Mult=my.Fpa.mult, recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results,fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))

  #weight
  #OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"d",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=F,stochastic.recruitment=stochastic.recruitment,use.penalty=F,HCR=HCR.used,penalty='Blim',penalty.type=0,Hessian=doHessian,Fpa.Mult=my.Fpa.mult,recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results, fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))
  OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"e",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=F,stochastic.recruitment=stochastic.recruitment,use.penalty=T,HCR=HCR.used,penalty='Blim',penalty.type=1,Hessian=doHessian,Fpa.Mult=my.Fpa.mult,recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results, fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))
  #OP.opti(scenario.basis=paste(ifelse(use.stochastic.recruitement,'S','D'),my.HCR,"f",sep=''),my.first.year=my.first.year,adjust.ini.N=F,my.last.year=my.last.year,use.weigting=F,use.price.at.age=F,stochastic.recruitment=stochastic.recruitment,use.penalty=T,HCR=HCR.used,penalty='Bpa' ,penalty.type=1,Hessian=doHessian,Fpa.Mult=my.Fpa.mult,recruit.adjust.CV=adjRec, recruit.adjust.factor=recruit.adjust.factor, FBlim.adjust=FBlim.adjust, just.batch.file=just.batch.file, just.plot.results=just.plot.results,fac.other.first.year=fac.other.first.year,fac.other.last.year=fac.other.last.year,fac.other=my.other[,my.oth],lab.other=ifelse(!do.sensitivity,'',as.character(my.oth)))
}


library(ggplot2)
library(reshape2)
library(rtracklayer)
library(ggpubr)
library(purrr)
library(GenomicRanges)
library(pheatmap)

source('./code/bwutils.r')
source('./code/heatmap_utils.r')

chromhmm.file <- './data/bed/ChromHMM17.bed'
bw.dir <- './data/bw/'
out.dir <- './final_figures/02/'

# Keep the order
levels.order <- c('1_Txn_Elongation',
                  '2_Weak_Txn',
                  '3_Txn_Transition',
                  '4_Poised_Enhancer',
                  '5_Active_Promoter',
                  '6_Strong_Enhancer',
                  '7_Active_Promoter',
                  '8_Strong_Enhancer',
                  '9_Strong_Enhancer',
                  '10_Poised_Promoter',
                  '11_Repressed',
                  '12_Heterochrom',
                  '13_Heterochrom',
                  '14_Heterochrom',
                  '15_Insulator',
                  '16_H3K9_Hetero',
                  '17_H3K9_H33_Hetero')

# Panel 1: ESC H3.3 SNAP time log2 fold change normalized to input
files <- paste0(bw.dir, c(
  'Deaton_2016/ES_H33_00h_rep1.bw',
  'Deaton_2016/ES_H33_03h_rep1.bw',
  'Deaton_2016/ES_H33_06h_rep1.bw',
  'Deaton_2016/ES_H33_12h_rep1.bw'
))

files.inp <- paste0(bw.dir, c(
  'Deaton_2016/ES_H33_inp_rep1.bw'
))

summaries <- map(files, summarize.by.bed, bed.file <- chromhmm.file, drop.name.col=F)
df.fg <- Reduce(function(...) merge(..., all=TRUE), summaries)
rownames(df.fg) <- df.fg$name
df.fg$name <- NULL

input <- summarize.by.bed(files.inp[[1]], chromhmm.file)
df.logfc <- log2(sweep(df.fg, 1, input$ES_H33_inp_rep1.bw, `/`))

colnames(df.logfc) <- c('0h','3h', '6h', '12h')
out.file <- paste0(out.dir, 'ES_SNAP_logfc_cov_enrichment.pdf')
summarized.heatmap(t(df.logfc[levels.order, ]), out.file, "ESC H3.3-SNAP Time-ChIP")

# Panel 2: ES_H33HA_tet_ON
files <- paste0(bw.dir, c(
  'Yildirim_2014/2014_Rando_H33_P3.mm9.bw',
  'Yildirim_2014/2014_Rando_H33_P6.mm9.bw'
))

summaries <- map(files, summarize.by.bed, bed.file <- chromhmm.file, drop.name.col=F)
df.fg <- Reduce(function(...) merge(..., by='name', all=TRUE), summaries)
rownames(df.fg) <- df.fg$name
df.fg$name <- NULL
out.file <- paste0(out.dir, 'ES_H33HA_tet_ON_cov.pdf')
summarized.heatmap(t(df.fg[levels.order, ]), out.file, "ESC H3.3-HA Tet-ON")

# Panel 3: H3K9me3
files <- paste0(bw.dir, c(
   'Shi_2019/2019_Shi_ESC1_WT_H3K9me3.bw'
))

out.file <- paste0(out.dir, 'ESC_H3K9me3.pdf')
summaries <- summarize.by.bed(files, chromhmm.file)
summarized.heatmap(t(summaries[levels.order, ]), out.file, "ESC H3K9me3")


#' dmulti_ddirch_forecast.r
#' NOTE: this assumes the predictor "map" is mean annual precipitation in mm, and natural log transforms.
#'
#' @param mod                        #model output from ddirch fitting function.
#' @param cov_mu                     #dataframe of site-level covariates. Must include all covariates used in models.
#' @param cov_sd                     #associated sd's if present.
#' @param glob.covs                  #global level covariates and standard deviations to fill in missing data.
#' @param n.samp                     #number of times to sample covariates/parameters for forecast. Default 1000.
#' @param zero_parameter_uncertainty #turns off drawing from parameter distributions, keeps parameters fixed at means.
#' @param zero_predictor_uncertainty #turns off drawing from covariate distributions, keeps covaraites fixed at means.
#' @param zero_process_uncertainty   #turns off process draw from rdirichlet. Basically just makes predictive interval = credible interval.
#'
#' @return              #returns a list of forecasts (mean, 95% CI, 95% PI) same length as model list.
#' @export
#'
#' @examples
source('NEFI_functions/precision_matrix_match.r')
dmulti_ddirch_forecast <- function(mod, seq.depth = 1000, cov_mu, names, cov_sd = NA, n.samp = 1000,
                                   zero_parameter_uncertainty = F,
                                   zero_covariate_uncertainty = F,
                                     zero_process_uncertainty = F,
                                                 make_it_work = F){
  #run some tests.----
  if(is.list(mod) == F){
    stop("Your model object isn't a list. It really needs to be.")
  }

  #grab the model out of list
  j.mod <- mod$jags_model #separate out the jags model.

  #### organize your covariates and covariate standard deviations.----
  preds <- as.character(mod$species_parameter_output$other$predictor)
  #grab covariates that were predictors.
  covs <- cov_mu[,colnames(cov_mu) %in% preds]
  #add intercept.
  covs <- cbind(rep(1,nrow(covs)), covs)
  colnames(covs)[1] <- 'intercept'
  #grab uncertainties in sd, if present.
  if(is.data.frame(cov_sd)){
    cov.sd <- precision_matrix_match(covs, cov_sd)
  }

  #re-order to match predictor order. Conditional otherwise it breaks the intercept only case.
  if(ncol(covs) > 1){
    covs   <-   data.frame(covs[,preds])
    if(is.data.frame(cov_sd)){
      cov.sd <- data.frame(cov.sd[,preds])
    }
  }

  #### Sample parameter and covariate space, make predictions.----
  pred.out <- list()
  cred.out <- list()
  for(j in 1:n.samp){
    #Sample parameters from mcmc output.----
    mcmc <- do.call(rbind,j.mod$mcmc)
    mcmc.sample <- mcmc[sample(nrow(mcmc),1),]
    #if we are fixing parameter uncertainty to zero, do something different.----
    if(zero_parameter_uncertainty == T){
      mcmc.sample <- colMeans(mcmc)
    }
    #grab x.m values, convert to matrix.
    x.m <- mcmc.sample[grep("^x\\.m\\[", names(mcmc.sample))]
    x.m <- matrix(x.m, nrow = ncol(covs), ncol = length(x.m)/ncol(covs))

    #Sample from covariate distributions.----
    #fix covariate uncertainty? Do this by making sd zero.
    if(zero_covariate_uncertainty==T){
      cov_sd <- NA
    }

    #only sample if you supplied uncertainties.
    if(is.data.frame(cov_sd)){
      now.cov <- matrix(NA, ncol = ncol(covs), nrow = nrow(covs))
      for(k in 1:ncol(covs)){now.cov[,k] <- rnorm(nrow(covs),covs[,k], cov.sd[,k])}
      colnames(now.cov) <- preds
      now.cov <- data.frame(now.cov)
      #log transform map values if this is one of your covariates, put covariates back in matrix form.
      #anti-logit relEM, multiply by 100 if this is one of your covariates.
      if('relEM' %in% colnames(now.cov)){now.cov$relEM <- boot::inv.logit(now.cov$relEM) * 100}
      if('map'   %in% colnames(now.cov)){now.cov$map   <- log(now.cov$map)}
      now.cov <- as.matrix(now.cov)
    }

    #If you did not supply covariate uncertainties then now.cov is just covs.
    if(!is.data.frame(cov_sd)){
      now.cov <- as.matrix(covs)
      colnames(now.cov) <- preds
      now.cov <- data.frame(now.cov)
      #log transform map values if this is one of your covariates, put covariates back in matrix form.
      #anti-logit relEM, multiply by 100 if this is one of your covariates.
      if('relEM' %in% colnames(now.cov)){now.cov$relEM <- boot::inv.logit(now.cov$relEM) * 100}
      if('map'   %in% colnames(now.cov)){now.cov$map   <- log(now.cov$map)}
      now.cov <- as.matrix(now.cov)
    }

    #Combine covariates and parameters to make a prediction.----
    pred.x.m <- matrix(NA, ncol=ncol(x.m), nrow = nrow(covs))
    for(k in 1:ncol(x.m)){pred.x.m[,k] <- exp(now.cov %*% x.m[,k])}
    #Sometimes parameter draws generate infinities for groups that didnt fit well. Force it?
    if(make_it_work == T){
      pred.x.m[pred.x.m == Inf] <- max(pred.x.m < Inf)
    }
    #get mean prediction and then draw from multinomial-dirichlet distribution.
    cred.out[[j]] <- pred.x.m / rowSums(pred.x.m)
    #prediction interval passes through Dirichlet and multinomial processes.
    dirichlet_out <- DirichletReg::rdirichlet(nrow(pred.x.m), pred.x.m + 0.1)
    multinom_out <- list()
    for(i in 1:nrow(dirichlet_out)){
      multinom_out[[i]] <- t(stats::rmultinom(1,seq.depth,dirichlet_out[i,]))
    }
    multinom_out <- do.call(rbind, multinom_out)
    pred.out[[j]] <- multinom_out/rowSums(multinom_out)
    #Turn off process uncertainty? If so pred.out is just cred.out.
    if(zero_process_uncertainty == T){
      pred.out[[j]] <- cred.out[[j]]
    }
  }

  #Summarize prediction mean and confidence intervals.----
  pred.mean       <- apply(simplify2array(cred.out), 1:2, mean)
  pred.ci.0.025   <- apply(simplify2array(cred.out), 1:2, quantile, probs = c(0.025))
  pred.ci.0.975   <- apply(simplify2array(cred.out), 1:2, quantile, probs = c(0.975))
  pred.pi.0.025   <- apply(simplify2array(pred.out), 1:2, quantile, probs = c(0.025))
  pred.pi.0.975   <- apply(simplify2array(pred.out), 1:2, quantile, probs = c(0.975))

  #batch into list, column names are group names, row names are siteIDs.
  output <- list(pred.mean, pred.ci.0.025, pred.ci.0.975, pred.pi.0.025, pred.pi.0.975)
  for(k in 1:length(output)){
    colnames(output[[k]]) <- names(mod$species_parameter_output)
    rownames(output[[k]]) <- names
  }
  names(output) <- c('mean','ci_0.025','ci_0.975','pi_0.025','pi_0.975')

  #return output.----
  return(output)
} #end function.----


library(keras)
library(tidyverse)
library(tokenizers)
library(tensorflow)

use_condaenv("r-tensorflow")

df_no_ea_samp = sample_n(df_no_ea,2000)

max_length <- 40

text <- df_no_ea_samp%>%
  pull(text) %>%
  str_c(collapse = " ") %>%
  tokenize_characters(lowercase = TRUE, strip_non_alphanum = FALSE, simplify = TRUE)

print(sprintf("Corpus length: %d", length(text)))
## [1] "Corpus length: 684767"
chars <- text %>%
  unique() %>%
  sort()

print(sprintf("Total characters: %d", length(chars)))
## [1] "Total characters: 74"

dataset <- map(
  seq(1, length(text) - max_length - 1, by = 3),
  ~list(sentence = text[.x:(.x + max_length - 1)],
        next_char = text[.x + max_length])
)

dataset <- transpose(dataset)

vectorize <- function(data, chars, max_length){
  x <- array(0, dim = c(length(data$sentence), max_length, length(chars)))
  y <- array(0, dim = c(length(data$sentence), length(chars)))

  for(i in 1:length(data$sentence)){
    x[i,,] <- sapply(chars, function(x){
      as.integer(x == data$sentence[[i]])
    })
    y[i,] <- as.integer(chars == data$next_char[[i]])
  }

  list(y = y,
       x = x)
}

vectors <- vectorize(dataset, chars, max_length)

create_model <- function(chars, max_length){
  keras_model_sequential() %>%
    layer_lstm(128, input_shape = c(max_length, length(chars))) %>%
    layer_dense(length(chars)) %>%
    layer_activation("softmax") %>%
    compile(
      loss = "categorical_crossentropy",
      optimizer = optimizer_rmsprop(lr = 0.01)
    )
}
fit_model <- function(model, vectors, epochs = 1){
  model %>% fit(
    vectors$x, vectors$y,
    batch_size = 128,
    epochs = epochs
  )
  NULL
}
generate_phrase <- function(model, text, chars, max_length, diversity){

  # this function chooses the next character for the phrase
  choose_next_char <- function(preds, chars, temperature){
    preds <- log(preds) / temperature
    exp_preds <- exp(preds)
    preds <- exp_preds / sum(exp(preds))

    next_index <- rmultinom(1, 1, preds) %>%
      as.integer() %>%
      which.max()
    chars[next_index]
  }

  # this function takes a sequence of characters and turns it into a numeric array for the model
  convert_sentence_to_data <- function(sentence, chars){
    x <- sapply(chars, function(x){
      as.integer(x == sentence)
    })
    array_reshape(x, c(1, dim(x)))
  }

  # the inital sentence is from the text
  start_index <- sample(1:(length(text) - max_length), size = 1)
  sentence <- text[start_index:(start_index + max_length - 1)]
  generated <- ""

  # while we still need characters for the phrase
  for(i in 1:(max_length * 20)){

    sentence_data <- convert_sentence_to_data(sentence, chars)

    # get the predictions for each next character
    preds <- predict(model, sentence_data)

    # choose the character
    next_char <- choose_next_char(preds, chars, diversity)

    # add it to the text and continue
    generated <- str_c(generated, next_char, collapse = "")
    sentence <- c(sentence[-1], next_char)
  }

  generated
}
iterate_model <- function(model, text, chars, max_length,
                          diversity, vectors, iterations){
  for(iteration in 1:iterations){

    message(sprintf("iteration: %02d ---------------\n\n", iteration))

    fit_model(model, vectors)

    for(diversity in c(0.2, 0.5, 1)){

      message(sprintf("diversity: %f ---------------\n\n", diversity))

      current_phrase <- 1:10 %>%
        map_chr(function(x) generate_phrase(model,
                                            text,
                                            chars,
                                            max_length,
                                            diversity))

      message(current_phrase, sep="\n")
      message("\n\n")

    }
  }
  NULL
}
model <- create_model(chars, max_length)

iterate_model(model, text, chars, max_length, diversity, vectors, 40)
## NULL