# R program to read in and monthly-average Southern Ocean station and shipboard CO2 data

## could remove all seasonal and climatological aggregation and plotting code since only monthly file used

readsurfobs=function(){

library('ncdf4')

# General options:
endyear=2020
localdatadir='local-data-files'
gvp60dir='obspack_co2_1_GLOBALVIEWplus_v6.0_2020-09-11/data/nc'
gvp70dir='obspack_co2_1_GLOBALVIEWplus_v7.0_2021-08-18/data/nc'

# LMG options:
uwlatbin=2; uwlatoffset=0 # by 2s, 56-66
noaalmgfiltco=0.3 # hour-to-hour jumps less than 0.3 ppm excludes 16%
ncarlmgfiltco=0.3 # hour-to-hour jumps less than 0.3 ppm excludes 20%
noaaagghourly=T # flag to only aggregate once
ncaragghourly=T # flag to only aggregate once

# Aggregation options:
minnmon=2 # lowest number of months allowed for seasonal average
meanwin=c(1998.9,2020.2) # window for calculating means, inclusive

# Plotting options
zoomyr=1999
ylm=c(-0.75,0.75)

# Select which sites to use:
stationlist=read.table('SO_CO2_stationlist.txt',header=T,stringsAsFactors=F) # use = 0 for do not use, 2 for SO, 1 was for reference but no longer used - use refcol below
print(stationlist)

# Read in surface data:

# set up arrays
allsta=data.frame(cbind(rep(seq(1957,endyear),each=12),rep(seq(1,12),times=endyear-1957+1)))
colnames(allsta)=c('year','month')
yrfrac=allsta$year+(allsta$mon-0.5)/12
monseas=rep(1,nrow(allsta)) # DJF
monseas[allsta$month>2&allsta$month<6]=2 # MAM
monseas[allsta$month>5&allsta$month<9]=3 # JJA
monseas[allsta$month>8&allsta$month<12]=4 # SON
seasyear=trunc(yrfrac+1/12) # shift Dec to next year
allseas=data.frame(cbind(aggregate(monseas,by=list(seas=monseas,year=seasyear),mean)$x,aggregate(yrfrac,by=list(seas=monseas,year=seasyear),mean)$x))
colnames(allseas)=c('seas','yrfrac')
alllat=NULL
alllon=NULL
allalt=NULL

# loop over records
for(i in c(1:nrow(stationlist))){

	staco2=NULL
	sta=stationlist$station[i]
	lab=stationlist$lab[i]
	meth=stationlist$method[i]
	use=stationlist$use[i]
	labcode=stationlist$labcode[i]
	print(c(sta,lab,meth,use))

	if(lab=='NOAA'){

		if(meth!='underway'){

			if(sta=='DRP'){ type='shipboard' } else { type='surface' }
			if(meth=='insitu'){ flag='allvalid' } else { flag='representative' }  # SPO = allvalid
			nc=nc_open(paste(localdatadir,'/',gvp60dir,'/co2_',tolower(sta),'_',type,'-',meth,'_1_',flag,'.nc',sep=''))
			gvp=data.frame(cbind(ncvar_get(nc,'time_decimal'),t(ncvar_get(nc,'time_components')),ncvar_get(nc,'value')*1E6)) ; colnames(gvp)=c('date','year','mon','day','hour','min','sec','co2')
			qcflag=ncvar_get(nc,'qcflag'); gvp$co2[substr(qcflag,1,1)!='.']=NA; gvp$co2[substr(qcflag,2,2)!='.']=NA
			gvpmon=aggregate(gvp$co2,by=list(year=gvp$year,month=gvp$mon),mean,na.rm=T) ; gvpmon=gvpmon[order(gvpmon$year+gvpmon$mon/12),]
			staco2=data.frame(cbind(gvpmon$year,gvpmon$month,gvpmon$x)); colnames(staco2)=c('year','month','co2')

		} else { # LMG

			if(noaaagghourly){ # only need to do once

				type='shipboard'
	                        fl=nc_open(paste(localdatadir,'/',gvp70dir,'/co2_gould_',type,'-insitu_1_allvalid.nc',sep='')) # filename uses 'gould'
	                        flgvp=data.frame(cbind(ncvar_get(fl,'time_decimal'),t(ncvar_get(fl,'time_components')),ncvar_get(fl,'value_original_scale')*1E6,ncvar_get(fl,'latitude'),ncvar_get(fl,'longitude'))) ; colnames(flgvp)=c('date','year','mon','day','hour','min','sec','co2','lat','lon')
				# GV+ 7.0 is X2019, so have to use 'value_original_scale'

	                        # aggregate to hourly medians
	                        noaalmghourly=aggregate(flgvp,by=list(year=flgvp$year,mon=flgvp$mon,day=flgvp$day,hour=flgvp$hour),median,na.rm=T)
	                        noaalmgdt=ISOdatetime(noaalmghourly$year,noaalmghourly$mon,noaalmghourly$day,noaalmghourly$hour,rep(0,nrow(noaalmghourly)),rep(0,nrow(noaalmghourly)),tz='UTC')
	                        noaalmghourly=noaalmghourly[order(noaalmgdt),]
	                        noaalmgdt=noaalmgdt[order(noaalmgdt)]

	                        # filter based on 2-hour differences
	                        co2steps1=c(999,diff(noaalmghourly$co2))
	                        co2steps2=c(diff(noaalmghourly$co2),999)
	                        noaalmghourly$co2[abs(co2steps1)>noaalmgfiltco|abs(co2steps2)>noaalmgfiltco]=NA # 0.3 cuts 16%

				# exclude prior to Dec 2005
				noaalmghourly=noaalmghourly[noaalmghourly$year>2005|noaalmghourly$mon>11,]

	                        # exclude box around PSA
	                        psalat=c(-64.82,-64.72) # 1/10 of deg lat or +/-3 nm
	                        psalon=c(-64.10,-64.00)
	                        noaalmghourly$co2[noaalmghourly$lat>psalat[1]&noaalmghourly$lat<psalat[2]&noaalmghourly$lon>psalon[1]&noaalmghourly$lon<psalon[2]]=NA

				noaaagghourly=F
			}

			# select latitude bin
                        noaalmglatbin=trunc((noaalmghourly$lat-uwlatoffset)/uwlatbin)*uwlatbin+uwlatoffset # north edge of bin
                        latlmg=noaalmghourly[noaalmglatbin==stationlist$lat[i]+uwlatbin/2,]
                        monlmg=aggregate(latlmg,by=list(year=latlmg$year,month=latlmg$mon),mean,na.rm=T)
                        lmgyrfrac=monlmg$year+monlmg$month/12-0.5
                        monlmg=monlmg[order(lmgyrfrac),]
                        staco2=data.frame(cbind(monlmg$year,monlmg$mon,monlmg$co2))
                        names(staco2)=c('year','month','co2')

		}

	} else if(lab=='NCAR'){ # LMG, assume meth=='underway'

		if(ncaragghourly){ # only need to do once

			alllmg=read.table(paste(localdatadir,'/ncar-lmg/lmg_all_results.txt',sep=''),header=T,stringsAsFactors=F) # 1-min data
			alllmg=alllmg[!is.na(alllmg$lat),] # ~6000 missing lats

			# aggregate to hourly medians
			alllmg=alllmg[,names(alllmg)!='leg'&names(alllmg)!='crsname'] # can not average text variables
			ncarlmghourly=aggregate(alllmg,by=list(year=alllmg$year,mon=alllmg$mon,day=alllmg$day,hour=alllmg$hour),median,na.rm=T)
			ncarlmgdt=ISOdatetime(ncarlmghourly$year,ncarlmghourly$mon,ncarlmghourly$day,ncarlmghourly$hour,rep(0,nrow(ncarlmghourly)),rep(0,nrow(ncarlmghourly)),tz='UTC')
			ncarlmghourly=ncarlmghourly[order(ncarlmgdt),]
			ncarlmgdt=ncarlmgdt[order(ncarlmgdt)]

			# filter based on 2-hour differences
			co2steps1=c(999,diff(ncarlmghourly$co2))
			co2steps2=c(diff(ncarlmghourly$co2),999)
			ncarlmghourly$co2[abs(co2steps1)>ncarlmgfiltco|abs(co2steps2)>ncarlmgfiltco]=NA # 0.3 cuts 20% of the data

			# exclude box around PSA
			psalat=c(-64.82,-64.72) # 1/10 of deg lat or +/-3 nm
			psalon=c(-64.10,-64.00)
			ncarlmghourly$co2[ncarlmghourly$lat>psalat[1]&ncarlmghourly$lat<psalat[2]&ncarlmghourly$lon>psalon[1]&ncarlmghourly$lon<psalon[2]]=NA

			ncaragghourly=F
		}

		# select latitude bin
		ncarlmglatbin=trunc((ncarlmghourly$lat-uwlatoffset)/uwlatbin)*uwlatbin+uwlatoffset
		latlmg=ncarlmghourly[ncarlmglatbin==stationlist$lat[i]+uwlatbin/2,]
		#latlmg=latlmg[,names(latlmg)!='leg'] # get rid of leg flag
		#latlmg=latlmg[,names(latlmg)!='crsname'] # get rid of cruise name
		monlmg=aggregate(latlmg,by=list(year=latlmg$year,month=latlmg$mon),mean,na.rm=T)
		lmgyrfrac=monlmg$year+monlmg$month/12-0.5
		monlmg=monlmg[order(lmgyrfrac),]
		staco2=data.frame(cbind(monlmg$year,monlmg$month,monlmg$co2))
		names(staco2)=c('year','month','co2')

	} else if(lab=='SIO_O2'){

		infile=paste(localdatadir,'/sio/',tolower(sta),'cav.csv',sep='')
		nhead=41
		names=readLines(infile,n=nhead)[nhead] ; names=unlist(strsplit(names,c(',',', '))) ; names=names[names!='']; names[1]='Date'; names=gsub(' ','',names)
		indata=read.csv(infile,skip=nhead+1,stringsAsFactors=F); colnames(indata)=names

		# aggregate to monthly
		indata$year=as.numeric(substr(indata$Date,1,4))
		indata$mon=as.numeric(substr(indata$Date,6,7))
		indatamon=aggregate(indata[,c('Value','StD','N','T-Excel','T-Decimal','Spline','Val-Spl')],by=list(year=indata$year,month=indata$mon),mean,na.rm=T); indatamon=indatamon[order(indatamon['T-Decimal']),] # allowing months with 1 sample
		staco2=data.frame(cbind(indatamon$year,indatamon$month,indatamon$Value)); colnames(staco2)=c('year','month','co2')

		# adjust SIO O2 Program CO2 from VH344-X2020 to WMO X2007 scale based on linear fit to NCAR primaries:
		sioo22gmd1=-0.002903; sioo22gmd0=1.041
		staco2$co2=staco2$co2*sioo22gmd1+sioo22gmd0+staco2$co2

	} else if(lab=='SIO_CDK'){

		infile=paste(localdatadir,'/sio/monthly_merge_co2_spo.csv',sep=''); nhead=59 # the obs have been slid parallel to a 4-harmonic (+ stiff-spline) fit to 2400 on the 15th of each month
		names=c('year','month','exceldate','date','co2','deseas','fit','fitdeseas','filled','filleddeseas')
		indata=read.csv(infile,skip=nhead,stringsAsFactors=F); colnames(indata)=names
		staco2=data.frame(cbind(indata$year,indata$month,indata$co2)); colnames(staco2)=c('year','month','co2')
		staco2$co2[staco2$co2==-99.99]=NA

		# adjust SIO CO2 Program CO2 from X12 to WMO X2007 scale based on linear fit to NCAR primaries (same as for SIO O2):
		sioo22gmd1=-0.002903; sioo22gmd0=1.041
		staco2$co2=staco2$co2*sioo22gmd1+sioo22gmd0+staco2$co2

## needed?
		staco2=staco2[order(staco2$year+staco2$month/12),]

	} else if(lab=='CSIRO'){

		type='surface'
		if(meth=='flask'){
			fl=nc_open(paste(localdatadir,'/',gvp60dir,'/co2_',tolower(sta),'_',type,'-flask_2_representative.nc',sep=''))
			flgvp=data.frame(cbind(ncvar_get(fl,'time_decimal'),t(ncvar_get(fl,'time_components')),ncvar_get(fl,'value')*1E6)) ; colnames(flgvp)=c('date','year','mon','day','hour','min','sec','co2')
			staco2=aggregate(flgvp$co2,by=list(year=flgvp$year,month=flgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		} else if(meth=='insitu'){
			is=nc_open(paste(localdatadir,'/wdcgg/nc/co2/hourly/co2_',tolower(sta),'_surface-insitu_16_9999-9999_hourly.nc',sep=''))
			isgvp=data.frame(cbind(t(ncvar_get(is,'start_time_components')),ncvar_get(is,'value'))) ; colnames(isgvp)=c('year','mon','day','hour','min','sec','co2')
			isgvp$co2[ncvar_get(is,'QCflag')!=1]=NA
			staco2=aggregate(isgvp$co2,by=list(year=isgvp$year,month=isgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		}
		names(staco2)=c('year','month','co2')

	} else if(lab=='SAWS'){ # CPT, get from ObsPack (could try WDCGG or SAWS too?)

                is=nc_open(paste(localdatadir,'/',gvp60dir,'/co2_cpt_surface-insitu_36_marine.nc',sep='')) # all 'qcflag' = 1
                isgvp=data.frame(cbind(t(ncvar_get(is,'time_components')),ncvar_get(is,'value')*1E6)) ; colnames(isgvp)=c('year','mon','day','hour','min','sec','co2')
                # for custom, now using WDCGG 222Rn to filter (WDCGG CO2 appears to be on X2019 so still using ObsPack)

                # filter for 222Rn < 150
                rn=nc_open(paste(localdatadir,'/wdcgg/nc/222rn/hourly/222rn_cpt_surface-insitu_7_9999-9999_hourly.nc',sep=''))
                rncust=data.frame(cbind(t(ncvar_get(rn,'start_time_components')),ncvar_get(rn,'value'))) ; colnames(rncust)=c('year','mon','day','hour','min','sec','rn')
                isgvp=merge(isgvp,rncust,by=c('year','mon','day','hour','min','sec'),all=F)
                isgvp$co2[isgvp$rn>150]=NA
                isgvp$co2[is.na(isgvp$rn)]=NA
		staco2=aggregate(isgvp$co2,by=list(year=isgvp$year,month=isgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		names(staco2)=c('year','month','co2')

	} else if(lab=='LSCE'){ # AMS, get from ObsPack (could try WDCGG or LSCE too?)

		# now using GV+ 7.0 with flagging ### use GV+ 6.1 instead? if so, set value_original_scale back to value
		is=nc_open(paste(localdatadir,'/',gvp70dir,'/co2_ams_surface-insitu_11_allvalid.nc',sep=''))
		isgvp=data.frame(cbind(t(ncvar_get(is,'time_components')),ncvar_get(is,'value_original_scale')*1E6)) ; colnames(isgvp)=c('year','mon','day','hour','min','sec','co2') # GV+ 7.0 is X2019, so have to use 'value_original_scale'
		isgvp$co2[ncvar_get(is,'obs_flag')==0]=NA
		staco2=aggregate(isgvp$co2,by=list(year=isgvp$year,month=isgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		names(staco2)=c('year','month','co2')

	} else if(lab=='TU'){ # SYO, get from ObsPack (could try WDCGG or TU too?)

		is=nc_open(paste(localdatadir,'/',gvp60dir,'/co2_',tolower(sta),'_surface-insitu_8_allvalid.nc',sep='')) # all 'obs_flag' = 1
	        isgvp=data.frame(cbind(t(ncvar_get(is,'time_components')),ncvar_get(is,'value')*1E6)) ; colnames(isgvp)=c('year','mon','day','hour','min','sec','co2')
		staco2=aggregate(isgvp$co2,by=list(year=isgvp$year,month=isgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		names(staco2)=c('year','month','co2')

	} else if(lab=='NIWA'){ # BHD, get from ObsPack (get from NIWA directly?)

		indata=read.table(paste(localdatadir,'/niwa/co2_bhd_surface-insitu_57_1978-2019_hourly.txt',sep=''),header=T)
		indata=indata[indata$QCflag==1,] # within steady southerly period
		isgvp=data.frame(cbind(indata$year,indata$month,indata$day,indata$hour,indata$minute,indata$second,indata$value)); colnames(isgvp)=c('year','mon','day','hour','min','sec','co2')
		staco2=aggregate(isgvp$co2,by=list(year=isgvp$year,month=isgvp$mon),mean,na.rm=T) ; staco2=staco2[order(staco2$year+staco2$mon/12),]
		names(staco2)=c('year','month','co2')

	}

	# aggregate records by season
	if(!is.null(staco2)){
		oldnames=colnames(allsta)
		allsta=merge(allsta,staco2,by=c('year','month'),all.x=T)
		colnames(allsta)=c(oldnames,paste(sta,'_',lab,'_',meth,sep=''))
		# aggregate by season
		seasco2=aggregate(allsta[,ncol(allsta)],by=list(seas=monseas,year=seasyear),mean,na.rm=T)$x # with na.rm=T so returns value even if only 1 month present (alldiffs filtered for < 2 months below)
		oldnames=colnames(allseas)
		allseas=cbind(allseas,seasco2) # since aggregating after merge, all rows present
		colnames(allseas)=c(oldnames,paste(sta,'_',lab,'_',meth,sep=''))
		alllat=c(alllat,stationlist$lat[i])
		alllon=c(alllon,stationlist$lon[i])
		allalt=c(allalt,stationlist$masl[i])
	}

} # loop on surface record

# print out time periods
for(i in c(3:ncol(allsta))){
	ind=c(1:nrow(allsta))[!is.na(allsta[,i])]
	print(paste(colnames(allsta)[i],',',month.abb[allsta$month[ind[1]]],allsta$year[ind[1]],'-',month.abb[allsta$month[tail(ind,1)]],allsta$year[tail(ind,1)]))
}

# write out results
sel=allsta$year>=trunc(meanwin[1])&allsta$year<=trunc(meanwin[2])
write(paste(names(allsta),collapse=' '),'SO_CO2_monthly.txt')
write(t(allsta[sel,]),'SO_CO2_monthly.txt',append=T,ncol=ncol(allsta))
write('record lat lon alt','SO_CO2_locations.txt')
write(rbind(names(allsta)[3:ncol(allsta)],alllat,alllon,allalt),'SO_CO2_locations.txt',append=T,ncol=4)
allloc=read.table('SO_CO2_locations.txt',header=T,stringsAsFactors=F)

#### code from here to end not really needed as only SO_CO2_monthly.txt used

# Calc diffs for all stations using NOAA in situ SPO record as the reference
alldiffsmon=allsta # monthly resolution
refcol=which(names(allsta)=='SPO_NOAA_insitu')
for(i in c(3:ncol(allsta))){
	alldiffsmon[,i]=allsta[,i]-allsta[,refcol]
}
alldiffsmonclim=aggregate(alldiffsmon[alldiffsmon$year>=meanwin[1]&alldiffsmon$year<=meanwin[2],3:ncol(alldiffsmon)],by=list(month=alldiffsmon$month[alldiffsmon$year>=meanwin[1]&alldiffsmon$year<=meanwin[2]]),mean,na.rm=T) # allows missing months
alldiffsmonclimcomp=apply(alldiffsmonclim[,which(stationlist$use==2)+1],1,mean,na.rm=T)

alldiffs=data.frame(cbind(aggregate(monseas,by=list(seas=monseas,year=seasyear),mean)$x,aggregate(yrfrac,by=list(seas=monseas,year=seasyear),mean)$x,aggregate(alldiffsmon[,3:ncol(alldiffsmon)],by=list(seas=monseas,year=seasyear),mean,na.rm=T)[,3:ncol(alldiffsmon)]))
allnmon=data.frame(cbind(aggregate(monseas,by=list(seas=monseas,year=seasyear),mean)$x,aggregate(yrfrac,by=list(seas=monseas,year=seasyear),mean)$x,aggregate(!is.na(alldiffsmon[,3:ncol(alldiffsmon)]),by=list(seas=monseas,year=seasyear),sum)[,3:ncol(alldiffsmon)]))
alldiffs[,3:ncol(alldiffs)][allnmon[,3:ncol(alldiffs)]<minnmon]=NA ## allow seasons with one missing month, but not two
colnames(alldiffs)[1:2]=c('seas','yrfrac')

# Calc long term mean and sd of diffs, from seasonal differences
meandiff=apply(alldiffs[alldiffs$yrfrac>=meanwin[1]&alldiffs$yrfrac<=meanwin[2],],2,mean,na.rm=T)
sddiff=apply(alldiffs[alldiffs$yrfrac>=meanwin[1]&alldiffs$yrfrac<=meanwin[2],],2,sd,na.rm=T)
sumdiff=apply(!is.na(alldiffs[alldiffs$yrfrac>=meanwin[1]&alldiffs$yrfrac<=meanwin[2],]),2,sum,na.rm=T)
meandiffseas=NULL
sddiffseas=NULL
for(seas in c(1:4)){
	meandiffseas=rbind(meandiffseas,apply(alldiffs[alldiffs$yrfrac>=meanwin[1]&alldiffs$yrfrac<=meanwin[2]&alldiffs$seas==seas,],2,mean,na.rm=T))
	sddiffseas=rbind(sddiffseas,apply(alldiffs[alldiffs$yrfrac>=meanwin[1]&alldiffs$yrfrac<=meanwin[2]&alldiffs$seas==seas,],2,sd,na.rm=T))
}
meandiffseas=data.frame(meandiffseas)
sddiffseas=data.frame(sddiffseas)
sddiffseas$seas=meandiffseas$seas

# Make plots:
library('RColorBrewer')
cols=rep(brewer.pal(12,'Paired'),5) # can handle up to 60 stations
pchs=rep(21,length(cols)) # NOAA
pchs[which(stationlist$lab=='SIO_CDK')]=22
pchs[which(stationlist$lab=='CSIRO')]=23
pchs[which(stationlist$lab=='SIO_O2')]=24
pchs[which(stationlist$lab=='NCAR')]=25
pchs[which(stationlist$lab=='TU')]=0
pchs[which(stationlist$lab=='LSCE')]=1
pchs[which(stationlist$lab=='SAWS')]=2
pchs[which(stationlist$lab=='NIWA')]=5
nmulti=sum(grepl('Multi',names(alldiffs)))
numsta=ncol(allsta)-2-nmulti
bgs=cols # these are fill colors for pch 21-25
cols[which(stationlist$method=='underway')]=rgb(0,0,0) # these are edge colors for pch 21-25

seasname=c('DJF','MAM','JJA','SON')
for(seas in c(1,3)){

# lat grad of diffs

png(paste('so_station_co2diff_gradient_',meanwin[1],'-',meanwin[2],'_',seasname[seas],'.png',sep=''),height=1200,width=1200,pointsize=30)
par(mar=c(5,5,4,2)+0.1)

plot(as.numeric(stationlist$lat),meandiffseas[meandiffseas$seas==seas,3:ncol(meandiffseas)],type='n',xlim=c(-90,max(as.numeric(stationlist$lat))),ylim=ylm,main=substitute(paste('Southern Ocean Seasonal ',CO[2],' Gradient - ',v),list(v=seasname[seas])),ylab=expression(paste(Delta,CO[2],' (ppm)')),xlab='Latitude (degrees N)',cex.main=1.2,cex.axis=1.2,cex.lab=1.2)
abline(h=0)
mtext(paste(stationlist$station[refcol-2],' ',stationlist$lab[refcol-2],' ',stationlist$method[refcol-2],' Subtracted, ',meanwin[1],'-',meanwin[2],' Average and SD',sep=''),3,0.3)
segwd=0.2
stasel=c(1:numsta)[c(1:numsta)+2!=refcol&!is.na(meandiffseas[meandiffseas$seas==seas,3:(length(meandiff)-nmulti)])]
y=c(0,as.numeric(meandiffseas[meandiffseas$seas==seas,3:length(meandiff)][stasel]))
x=c(-90,stationlist$lat[stasel])
w=c(10000,as.numeric(1/sddiffseas[sddiffseas$seas==seas,3:length(meandiff)][stasel]^2)) # 10000 for SPO equiv of SD of 0.01
new=data.frame(x=seq(-90,ceiling(max(x)),1))
lines(new$x,predict.lm(lm(y ~ poly(x,3),weights=w),new),lwd=2)
for(i in stasel){
	segments(stationlist$lat[i],meandiffseas[meandiffseas$seas==seas,i+2]-sddiffseas[sddiffseas$seas==seas,i+2],stationlist$lat[i],meandiffseas[meandiffseas$seas==seas,i+2]+sddiffseas[sddiffseas$seas==seas,i+2],col=cols[i])
	segments(stationlist$lat[i]-segwd,meandiffseas[meandiffseas$seas==seas,i+2]-sddiffseas[sddiffseas$seas==seas,i+2],stationlist$lat[i]+segwd,meandiffseas[meandiffseas$seas==seas,i+2]-sddiffseas[sddiffseas$seas==seas,i+2],col=cols[i])
	segments(stationlist$lat[i]-segwd,meandiffseas[meandiffseas$seas==seas,i+2]+sddiffseas[sddiffseas$seas==seas,i+2],stationlist$lat[i]+segwd,meandiffseas[meandiffseas$seas==seas,i+2]+sddiffseas[sddiffseas$seas==seas,i+2],col=cols[i])
        points(stationlist$lat[i],meandiffseas[meandiffseas$seas==seas,i+2],pch=pchs[i],bg=bgs[i],col=cols[i],cex=1.5,lwd=2)
	if(any(!is.na(meandiffseas[meandiffseas$seas==seas,i+2]))) text(stationlist$lat[i],ylm[1],stationlist$sta[i],col=cols[i],srt=90,offset=0,adj=c(0,0.5))
}
labsel=stationlist$lab[stasel]; pchsel=pchs[stasel]; pchsel=pchsel[!duplicated(labsel)]; labsel=labsel[!duplicated(labsel)]
legend('topleft',labsel,pch=pchsel,pt.bg='black',cex=1.0,col='gray30',ncol=2)
if(any(stationlist$method=='underway')) legend('topright','Underway data points outlined in black',cex=0.75,bty='n')

dev.off()

} # loop on season

# Plot composite seasonal cycle

compflag=paste(tolower(substr(stationlist$station[stationlist$use==2],1,1)),collapse='')
refflag='s'

png(paste('so_station_composite_co2diff_seascycle_',compflag,'-',refflag,'_',meanwin[1],'-',meanwin[2],'.png',sep=''),height=1200,width=1800,pointsize=30)
par(mar=c(5,5,4,2)+0.1)

plot(seq(0.5,11.5),alldiffsmonclim[c(7:12,1:6),3],type='n',xlim=c(0,12),ylim=ylm,main=expression(paste('Southern Ocean Climatological Monthly Mean ',CO[2])),ylab=expression(paste(Delta,CO[2],' (ppm)')),xlab='Month',cex.main=1.2,cex.lab=1.2,axes=F)
box()
axis(2,cex.axis=1.2)
axis(1,at=c(0:12),labels=F,cex.axis=1.5)
axis(1,seq(0.5,11.5),labels=c('J','A','S','O','N','D','J','F','M','A','M','J'),cex.axis=1.3,tick=F)
mtext(paste(stationlist$station[refcol-2],' ',stationlist$lab[refcol-2],' ',stationlist$method[refcol-2],' Subtracted, ',meanwin[1],'-',meanwin[2],sep=''),3,0.3)
abline(h=0)
for(i in c(1:numsta)[stationlist$use==2]){
        points(seq(0.5,11.5),alldiffsmonclim[c(7:12,1:6),i+1],type='b',pch=pchs[i],bg=bgs[i],col=cols[i],cex=1.5,lwd=2)
}
lines(seq(0.5,11.5),alldiffsmonclimcomp[c(7:12,1:6)],col='grey10',lwd=10)

legend('topleft',names(alldiffsmonclim)[c(1:numsta)[stationlist$use==2]+1],col=cols[c(1:numsta)[stationlist$use==2]],pch=pchs[c(1:numsta)[stationlist$use==2]],pt.bg=bgs[c(1:numsta)[stationlist$use==2]],cex=0.75,pt.cex=0.75,lwd=2,ncol=3) # pt.lwd?

dev.off()

} # end of readsurfobs()


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

## Author: Zach Rivers, @riverspharmd, zrivers@fredhutch.org
## Contributors:

## Purpose:

## Depends on:
## 07PredictorExploration.r
## 10COSTTidyJoin.r
## tidyverse
## gtsummary
## flextable

## Input:
## IntData/dat_wide.csv
## IntData/cost.csv

## Output:

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

# Packages______________________________________________________________________
library(tidyverse)
library(gtsummary)
library(flextable)

# Functions_____________________________________________________________________
source("Functions/importing.r")
source("Functions/regressions.r")
source("Functions/tidying.r")
# Data In_______________________________________________________________________
pred_wide <- read_csv(
    "IntData/dat_wide.csv"
)

## COST

cost <- read_csv("IntData/cost.csv",
    col_types = "ccccn"
)
# Data Prep_____________________________________________________________________
## Clean COST-------------------------------------------------------------------
cost_outcomes <- build_outcome(dat = cost)

## Prep Predictors-------------------------------------------------------------
### Factor
pred_wide <- label_factors(dat = pred_wide)

### Build predictor datset
pred_wide <- build_predictors(dat = pred_wide)

### Label predictor dataset
var_label(pred_wide) <- lr_labels

## Join with outcomes-----------------------------------------------------------
pred_outcomes <- join_pred_out(
    pred_dat = pred_wide,
    outcome_dat = cost_outcomes
)
# Run Models____________________________________________________________________

### Build Model Pieces
pat_var <- c(
    "age_p",
    "gender_p",
    "ethnicity_p",
    "race_p",
    "education_p",
    "comorbid_sum_p"
)
clin_var <- c(
    "location",
    "stage",
    "dx"
)
care_var <- c(
    "age_c",
    "gender_c",
    "ethnicity_c",
    "race_c",
    "education_c",
    "comorbid_sum_c"
)

dyad_var <- c(
    "marital",
    "income",
    "children"
)
### Build Model Inputs
pat_input <- list(
    p1 = c(pat_var),
    p2 = c(pat_var, clin_var),
    p3 = c(pat_var, care_var),
    p4 = c(pat_var, dyad_var),
    p5 = c(pat_var, clin_var, care_var),
    p6 = c(pat_var, clin_var, dyad_var, "income * location"),
    p7 = c(pat_var, care_var, dyad_var),
    p8 = c(pat_var, clin_var, care_var, dyad_var, "income * location")
)

care_input <- list(
    c1 = c(care_var),
    c2 = c(care_var, clin_var),
    c3 = c(care_var, pat_var),
    c4 = c(care_var, dyad_var),
    c5 = c(care_var, clin_var, pat_var),
    c6 = c(care_var, clin_var, dyad_var, "income * location"),
    c7 = c(care_var, pat_var, dyad_var),
    c8 = c(care_var, clin_var, pat_var, dyad_var, "income * location")
)

dyad_input <- c(pat_input[1:7], care_input[c(1, 2, 4, 6, 8)])
mod_input <- list(pat_input, care_input, dyad_input)
### Build Model Outcomes
outcome_list <- list("p_fintox", "c_fintox", "con")

### Build Table Spanners
p_spanner <- c(
    "Bivariate",
    "Patient",
    "Patient and Clinical",
    "Patient and Caregiver",
    "Patient and Dyad",
    "Patient, Clinical, and Caregiver",
    "Patient, Clinical, and Dyad",
    "Patient, Caregiver, and Dyad",
    "All Predictors"
)
c_spanner <- c(
    "Bivariate",
    "Caregiver",
    "Caregiver and Clinical",
    "Caregiver and Patient",
    "Caregiver and Dyad",
    "Caregiver, Clinical, and Patient",
    "Caregiver, Clinical, and Dyad",
    "Caregiver, Patient, and Dyad",
    "All Predictors"
)

dyad_spanner <- c(p_spanner[1:8], c_spanner[c(2, 3, 4, 7, 9)])
span_list <- list(p_spanner, c_spanner, dyad_spanner)

### Model output list
joined_model_list <- list(NA, NA, NA)
### Run models
# https://aosmith.rbind.io/2019/06/24/function-for-model-fitting/
for (i in 1:3) {
    log_reg_list <- map(.x = mod_input[[i]], ~ mv_log_reg(
        input = .x,
        resp = outcome_list[[i]],
        dat = pred_outcomes
    ))

    ### Tidy model fits
    log_reg_list_tidy <- map(
        .x = log_reg_list,
        ~ tidy_mv_log_reg(mod_fit = .x)
    )

    ### Join Models

    joined_model_list[[i]] <- tbl_merge(log_reg_list_tidy,
        tab_spanner = span_list[[i]]
    )
    print(paste("Completed Loop", i, "of Three"))
}
joined_model_list
# Data out______________________________________________________________________


#####################################################################################################
#################### Aim: estimate summary measures
#####################################################################################################
rm(list=ls())
dev.off()

library(eurostat)
library(MASS)
library(sjmisc)
library(tidyr)
library(dplyr)
library(fBasics)

data.dir <- "C:\\Users\\Magdalena\\demography\\happyindex\\data\\SILC"
lt.data.dir <- "C:\\Users\\Magdalena\\demography\\happyindex\\data\\EUROSTAT"
GBD.data.dir <- "C:\\Users\\Magdalena\\demography\\happyindex\\data\\GBD"

##prevalence data
setwd(data.dir)
chronicprev <- read.table(file="chronicprev.csv", sep=",", header=TRUE)
chronicperc <- tapply(X=chronicprev$Freq, INDEX=list(chronic=chronicprev$chronic,
                                                     age=chronicprev$agegr, country=chronicprev$country, sex=chronicprev$sex), FUN=sum)

GALIprev <- read.table(file="GALIprev.csv", sep=",", header=TRUE)
GALIperc <- tapply(X=GALIprev$Freq, INDEX=list(GALI=GALIprev$GALI,
                                                     age=GALIprev$agegr, country=GALIprev$country, sex=GALIprev$sex), FUN=sum)

SRprev <- read.table(file="SRprev.csv", sep=",", header=TRUE)
SRperc <- tapply(X=SRprev$Freq, INDEX=list(SR=SRprev$SR,
                                               age=SRprev$agegr, country=SRprev$country, sex=SRprev$sex), FUN=sum)

multiprev <- read.table(file="multiprev.csv", sep=",",  header=TRUE)
multiperc <- tapply(X=multiprev$Freq, INDEX=list(multi=multiprev$multi,
                                           age=multiprev$agegr, country=multiprev$country, sex=multiprev$sex), FUN=sum)

#cname <- c( "Austria", "Belgium", "Bulgaria", "Switzerland", "Cyprus", "Czechia", "Denmark", "Estonia", "Greece",
#             "Spain", "Finland", "France", "Croatia","Hungary","Ireland","Italy", "Lithuania", "Luxembourg", "Latvia", "Netherlands","Norway", "Poland",
#             "Portugal", "Romania", "Serbia",  "Sweden", "Slovenia", "Slovakia", "United Kingdom")

#lt data
setwd(lt.data.dir)

Lx <- read.table(file="Lx.csv", sep=",", header=TRUE)
l15 <- read.table(file="l15.csv", sep=",", header = TRUE)

Lx <- tapply(Lx$values, INDEX=list(age=Lx$agegr, country=Lx$cno, sex=Lx$sex), FUN=sum)
lx15 <-  tapply(l15$values, INDEX=list(country=l15$cno, sex=l15$sex), FUN=sum)

##########################################################################################
######## Lx redistributed by health states
##########################################################################################
Lxchronic <- array(0, dim=dim(chronicperc), dimnames=dimnames(chronicperc))
for (i in 1:2){
  Lxchronic[i,,,] <- Lx[,,]*chronicperc[i,,,]
}
#
LxGALI <- array(0, dim=dim(GALIperc), dimnames=dimnames(GALIperc))
for (i in 1:3){
  LxGALI[i,,,]  <- GALIperc[i,,,]*Lx[,,]
}
#
LxSR <- array(0, dim=dim(SRperc), dimnames=dimnames(SRperc))
for (i in 1:5){
  LxSR[i,,,]  <- SRperc[i,,,]*Lx[,,]
}
#
Lxmulti <- array(0, dim=dim(multiperc), dimnames=dimnames(multiperc))
for (i in 1:30){
  Lxmulti[i,,,]  <- multiperc[i,,,]*Lx[,,]
}

##############################################################################
###Lx reditributed with well-being weights
#weights do not differentiated between age groups
#############################################################################
setwd(data.dir)
chronicallw <- read.table(file="chronicweights.csv", sep=",", header=TRUE)
selfweights <- read.table(file="selfweights.csv", sep=",", header=TRUE)
GALIallw <- read.table(file="GALIweights.csv", sep=",", header=TRUE)
multiweights <- read.table(file="multiweights.csv", sep=",", header=TRUE)

###chronic
#prepare weights
chronicw2 <- tapply(X=chronicallw$chronic, INDEX=list(country=chronicallw$country,sex=chronicallw$sex), FUN=sum)
Lxchronicw <- Lxchronic
for (i in 1:10){
  Lxchronicw[2,i,,] <- Lxchronic[2,i,,]*chronicw2
}

####GALI
#prepare weights
GALIallw2 <- GALIallw %>%
  pivot_longer(cols=GALI1:GALI2,
               names_to="GALI", values_to="weights")

GALI2w <- tapply(X=GALIallw2$weights,
                 INDEX=list(GALI=GALIallw2$GALI, country=GALIallw2$country,sex=GALIallw2$sex), FUN=sum)

LxGALIw <- LxGALI
for (j in 1:2){
  for (i in 1:10){
    LxGALIw[j+1,i,,]  <- LxGALI[j+1,i,,]*GALI2w[j,,]
  }
}

####self-rated health
#prepare weights
selfweights2 <- selfweights %>%
  pivot_longer(cols = SR2:SR5,
               names_to = "sr", values_to = "weights")

selfw2 <- tapply(X=selfweights2$weights,
                 INDEX=list(SR=selfweights2$sr, country=selfweights2$country,sex=selfweights2$sex), FUN=sum)

LxSRw <- LxSR
for (j in 1:4){
  for (i in 1:10){
    LxSRw[j+1,i,,]  <- LxSR[j+1,i,,]*selfw2[j,,]
  }
}

####multiple
#prepare weights
multiw2 <- tapply(X=multiweights$weight,
                 INDEX=list(multi=multiweights$level,country=multiweights$country,sex=multiweights$sex), FUN=sum)

Lxmultiw <- Lxmulti
for (j in 1:29){
  for (i in 1:10){
    Lxmultiw[j+1,i,,]  <- Lxmultiw[j+1,i,,]*multiw2[j,,]
  }
}

##########################################################################################
########## WAHE indicator
##########################################################################################
WAHEchr15 <- array(0, dim=dim(Lxchronicw[1,1,,]), dimnames=(dimnames(Lxchronicw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    WAHEchr15[i,j] <- sum(Lxchronicw[,,i,j])/lx15[i,j]
  }
}

WAHEGALI15 <- array(0, dim=dim(LxGALIw[1,1,,]), dimnames=(dimnames(LxGALIw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    WAHEGALI15[i,j] <- sum(LxGALIw[,,i,j])/lx15[i,j]
  }
}

WAHESR15 <- array(0, dim=dim(LxSRw[1,1,,]), dimnames=(dimnames(LxSRw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    WAHESR15[i,j] <- sum(LxSRw[,,i,j])/lx15[i,j]
  }
}

WAHEmulti15 <- array(0, dim=dim(Lxmultiw[1,1,,]), dimnames=(dimnames(Lxmultiw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    WAHEmulti15[i,j] <- sum(Lxmultiw[,,i,j], na.rm=TRUE)/lx15[i,j]
  }
}

###########################################################################################
###################HE indicator
###########################################################################################
HEchr15 <- array(0, dim=dim(Lxchronic[1,1,,]), dimnames=(dimnames(Lxchronic[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    HEchr15[i,j] <- sum(Lxchronic[1,,i,j])/lx15[i,j]
  }
}

HEGALI15 <- array(0, dim=dim(LxGALIw[1,1,,]), dimnames=(dimnames(LxGALIw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    HEGALI15[i,j] <- sum(LxGALIw[1,,i,j])/lx15[i,j]
  }
}

HESR15 <- array(0, dim=dim(LxSRw[1,1,,]), dimnames=(dimnames(LxSRw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    HESR15[i,j] <- sum(LxSRw[1,,i,j])/lx15[i,j]
  }
}

HEmulti15 <- array(0, dim=dim(Lxmultiw[1,1,,]), dimnames=(dimnames(Lxmultiw[1,1,,])))
for (i in 1:29){
  for (j in 1:2){
    HEmulti15[i,j] <- sum(Lxmultiw[1,,i,j],na.rm=TRUE)/lx15[i,j]
  }
}

########################################################################################################
################# e15
########################################################################################################
e15 <- array(0, dim=dim(lx15), dimnames=dimnames(lx15))

for (j in 1:2){
  for (i in 1:29){
    e15[i,j]  <- sum(Lx[,i,j])/lx15[i,j]
  }
}

########################################################################################################
### DALE
########################################################################################################
setwd(GBD.data.dir)
gball <- read.table(file="gbd2.csv", sep=",", header=TRUE)

###order by country as in other indicators
cname <- c( "Austria", "Belgium", "Bulgaria", "Switzerland", "Cyprus", "Czechia", "Denmark", "Estonia", "Greece",
             "Spain", "Finland", "France", "Croatia","Hungary","Ireland","Italy", "Lithuania", "Luxembourg", "Latvia", "Netherlands","Norway", "Poland",
             "Portugal", "Romania", "Serbia",  "Sweden", "Slovenia", "Slovakia", "United Kingdom")
cno <- 1:29
gball$cno <- 0
for (i in 1:length(cno)){
  gball$cno[rownames(gball)==cname[i]] <- i
}

#order by cno
DALE <- gball[order(gball$cno),]

##########################################################################################################
#########################  write out results, also appenix data table
##########################################################################################################

males15 <- cbind(cname,e15[,1], DALE[,2], HEchr15[,1], HEGALI15[,1], HESR15[,1], HEmulti15[,1],
                 WAHEchr15[,1], WAHEGALI15[,1], WAHESR15[,1],  WAHEmulti15[,1])
fem15 <- cbind(cname,e15[,2], DALE[,1], HEchr15[,2], HEGALI15[,2], HESR15[,2], HEmulti15[,2],
                 WAHEchr15[,2], WAHEGALI15[,2], WAHESR15[,2],  WAHEmulti15[,2])

colind <- c("country","LE", "DALE", "HE.chr", "HE.GALI",  "HE.sr", "HE.multi", "WAHE.chr", "WAHE.GALI",  "WAHE.sr", "WAHE.multi")
colnames(males15) <- colind
colnames(fem15) <- colind

setwd(data.dir)
write.table(males15,file="malesmph15.csv", sep=",")
write.table(fem15, file="femsmph15.csv", sep=",")


drawClusterHeatMap	<- function( seuratObj , heatMapHeight, heatMapWidth){

# This snippet draws a heatmap of gene expression data (found in seuratObj@data) and annotates clusters
# according seuratObj@ident with colors taken from setClusterColors.r
# The order of clusters is set manually, so that eHMP and ltHMP are at the bottom of the legend and in the left of the graph
#
# written by Vsevolod Makeev 2017 - 2021

require( ComplexHeatmap)
require( circlize)
require( viridis)
require( proxy)
require( dendsort)

source( "R/setClusterColors.r")

heatMapGenes	<- rownames( seuratObj@data)
dataMatrix 	<- seuratObj@data

nCol		<- 1024

clusterColors		<- setClusterColors( seuratObj)
names(clusterColors)	<- levels( seuratObj@ident)

if( seuratObj@project.name == "WT"){

	orderedCellTypes	<- ordered( seuratObj@ident, levels = c( "eHMP", "ltHMP", "I", "M", "X", "7", "4"))
	orderedCellTypes	<- orderedCellTypes[ order( orderedCellTypes)]

}else if( seuratObj@project.name == "sox10_mutants"){
	orderedCellTypes	<- seuratObj@ident
	orderedCellTypes	<- orderedCellTypes[ order( orderedCellTypes)]
}else{ cat( "Unknown project type /n")}

row_dend 	<- dendsort(hclust( dist(  dataMatrix, method = "cosine", pairwise = TRUE)))
colPanelFun	 = colorRamp2( quantile( dataMatrix, seq(0, 1, by = 1/(nCol - 1))), viridis( nCol))

annotBar		<- HeatmapAnnotation(
				Cluster = orderedCellTypes,
				col = list( Cluster = clusterColors),
				height = unit(30, "points"),
				which = "column",
				annotation_name_gp = gpar( fontsize = 16, fontface = "bold"),
				simple_anno_size = unit( 0.7, "cm"),
				annotation_legend_param = list( ncol = 2, nrow = length( levels(seuratObj@ident)))
					)
hMap		<- Heatmap( 	dataMatrix[ , names(orderedCellTypes)],
			name 	= "gene expression",
			width 	= unit( heatMapWidth,  "inches"),
			height 	= unit( heatMapHeight, "inches"),
			heatmap_legend_param = list(
				title = "log10(Exp)",
				legend_height = unit( 2, "inches"),
				grid_width = unit( 0.5, "inches"),
				title_gp = gpar( fontsize = 12, fontface = "bold")
						),
			col	=  colPanelFun,
			bottom_annotation = annotBar,
			column_order = names(orderedCellTypes),
			column_split = orderedCellTypes,
#			cluster_column_slices = FALSE,
			cluster_rows = row_dend,
			cluster_columns = FALSE,
			row_dend_reorder = FALSE,
			column_dend_reorder = FALSE,
			show_column_names = FALSE,
			row_names_gp = gpar(fontsize = 12),
#			column_names_gp = gpar(fontsize = 16),
			clustering_distance_rows = function(x) dist( x, method = "cosine", pairwise = TRUE),
#			clustering_distance_columns = function(x) dist( x, method = "cosine", pairwise = TRUE),
			use_raster = TRUE, raster_device = "png"
			)

return( hMap)

}


################################################
## 2-1_Compute-metrics_emptyDrops.r
##  - Compute emptyDrops per sample
##  - Save an rds with emptyDrops result per sample
################################################

suppressWarnings(library(Seurat))
suppressWarnings(library(DropletUtils))

message("reloading old matrices...")
scdata_list <- readRDS("/output/pre-doublet-scdata_list.rds")

message("Loading configuration...")
config <- RJSONIO::fromJSON("/input/meta.json")

# Check which samples have been selected. Otherwiser we are going to use all of them.
if (length(config$samples)>0){
    samples <- config$samples
}else{
    samples <- names(scdata_list)
}

scdata_list <- scdata_list[samples]

# compute_emptydrops function
#' @description Save the result of emptyDrops per sample.
#' @param scdata Raw sparse matrix with the counts for one sample.
#' @param sample_name Name of the sample that we are preparing.
#'
#' @return
compute_emptydrops <- function(scdata, sample_name) {

  message("Sample --> ", sample_name, "...")

  # [HARDCODE]
  threshold_emptydrops <- 100
  # threshold_emptydrops <- 2000 # use this to simulate unfiltered data
  # automatically checking if data contains enough (50) barcodes that
  # are confidently considered to be empty to serve as training set
  # not necessary as of now: sce <- as.SingleCellExperiment(scdata)
  nr_barcodes_ambient <- sum(colSums(scdata) < threshold_emptydrops)
  range(colSums(scdata) )
  if (nr_barcodes_ambient < 50 ) {
    flag_filtered <- TRUE
    message("Found not enough [", nr_barcodes_ambient, "] ambient barcodes,
          so data is considered to be pre-filtered")
    message("Smallest nr of UMI for any barcode: ", min(colSums(scdata)))
    message("EmptyDrops filter will not be applied")
  } else {
    flag_filtered <- FALSE
    # if non-filtered, run emptyDrops and save results in file
    message("Found enough [", nr_barcodes_ambient, "] ambient barcodes,
          so data is considered to be not filtered")
    emptydrops_out <- emptyDrops(scdata, lower = threshold_emptydrops)

    file_path <- paste("/output/pre-emptydrops-", sample_name, ".rds", sep = "")
    saveRDS(emptydrops_out, file = file_path, compress = FALSE)
  }
}

message("calculating probability of barcodes being background noise...")
for (sample_name in names(scdata_list)) {
  compute_emptydrops(scdata_list[[sample_name]], sample_name)
}

message("Step 2-1 completed.")


#--------------------------#
# LA variation - modelling #
#--------------------------#

# load packages ----
library(dplyr)
library(tidyr)
library(tibble)
library(readxl)
library(readr)
library(purrr)
library(ggplot2)
library(naniar)
library(broom)
library(lme4)
library(ggrepel)
library(reshape2)
library(janitor)
library(tableone)
library(glmnet)
library(doParallel)
library(MuMIn)
library(lmtest)
library(GGally)
library(car)
library(stringr)

#--------------#
# get data ----
#--------------#
source("N:/Documents/Data & Programming/R Scripts/la variation/la_variation_analysis_data.r", echo = FALSE)

#------------------------------------------#
# look at correlation between variables ----
#------------------------------------------#

# using only data from 2010/11 (i.e. mid-point)
dat_forcor <- model_dat %>%
  filter(year == 2010) %>%
  select(ends_with("_bt"), lone_parent_pct, violent_crime, prop10_imd_pct)

# create correlations (method = spearman's)
cormat <- round(cor(dat_forcor, use = "complete.obs", method = 'spearman'),2) %>%
  melt()

# plot these in coloured tiles
ggplot(data = cormat, aes(x = Var1, y = Var2, fill = value)) +
  geom_tile() +
  theme(axis.text.x = element_text(angle = 45, hjust = 1)) +
  scale_fill_gradientn(colours=rainbow(2))
ggsave("N:/Documents/Data & Programming/outputs/la variation/figures/correlation_analysis_variables.png")

# try again but with distribution of each var in plot (using GGpairs):
dat_forcor <- filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0) %>%
       filter(year == 2010) %>%
       dplyr::select(teenage_births_pct= teenbirth_2010, propmumsLSOA_in_10pct_mostdepr= imd_10pc_2010,
              low_birweit_pct = lowbw_2010,
              congenital_anoms_pct= ca_feud_2010,
              lone_parent_pct, ARA_pct= any_ara_pct_bt, violent_crime,
              propLSOA_in_10pct_mostdepr= prop10_imd_pct)

lowerFn <- function(data, mapping, method = "lm", ...) {
  p <- ggplot(data = data, mapping = mapping) +
    geom_point(colour = "black") +
    geom_smooth(method = method, color = "red", ...)+
    theme_classic()
  p
}

ggpairs(dat_forcor, lower = list(continuous = wrap(lowerFn, method = "lm")))
ggsave("N:/Documents/Data & Programming/outputs/la variation/figures/correlation_analysis_variables_scatter.png", units = "mm",
       width = 510, height = 400)

#----------------------------------------------------------------------------------------------#
# fitting lmer models ----
#----------------------------------------------------------------------------------------------#
set.seed(1234)

# model diagnostic functions:
source("N:/Documents/Data & Programming/R Scripts/lmer_diagnostics.R", echo = FALSE)

# pseudo R-squared confidence interval bootstrapping function:
source("N:/Documents/Data & Programming/R Scripts/la variation/la_variation_lmer_bs.R", echo = FALSE)

# function to create dataframe of lmer object results:
res_2_df <- function(model, model_R2, model_name){
  lmer_null_tidy <- broom::tidy(model, conf.int = TRUE) %>% # get model estimates with confidence intervals
    mutate(R2m_actual = r.squaredGLMM(model)[1,1],# get actual marginal R-squared
           R2c_actual = r.squaredGLMM(model)[1,2],# get actual conditional R-squared
           AIC = AIC(model),
           BIC = BIC(model),
           model = model_name,
           R2m_median = distinct(filter(model_R2, type_R2 == "R2m"),median)$median, # from 10,000 bootstrapped samples
           R2m_ll = distinct(filter(model_R2, type_R2 == "R2m"),ll)$ll,
           R2m_ul = distinct(filter(model_R2, type_R2 == "R2m"),ul)$ul,
           R2c_median = distinct(filter(model_R2, type_R2 == "R2c"),median)$median,
           R2c_ll = distinct(filter(model_R2, type_R2 == "R2c"),ll)$ll,
           R2c_ul = distinct(filter(model_R2, type_R2 == "R2c"),ul)$ul)
}

#-----------------#
# Null model   ----
#-----------------#

lmer_null <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                  formula = as.formula("n_child_10000 ~ year_std + (year_std|la_code)"), REML = TRUE)

# get median and 95% CI for pseudo R-squared:
lmer_null_bs <- R2_bootstrap(model = lmer_null)

# create dataframe of model results:
lmer_null_tidy <- res_2_df(model = lmer_null, model_R2 = lmer_null_bs, model_name = "null")

# without random slope:
lmer_null_noslope <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                          formula = as.formula("n_child_10000 ~ year_std + (1|la_code)"), REML = TRUE)

G <- -2*logLik(lmer_null_noslope) + 2*logLik(lmer_null)
pchisq(as.numeric(G), df = 2, lower.tail = FALSE) # p < 0.001

#-----------------------------------#
# Unadjusted association model   ----
#-----------------------------------#
lmer_aras <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                  formula = as.formula("n_child_10000 ~ year_std + any_ara_pct_wt + any_ara_pct_bt + (year_std|la_code)"),
                  REML = TRUE)

# get median and 95% CI for pseudo R-squared:
lmer_aras_bs <- R2_bootstrap(model = lmer_aras)

# create dataframe of model results:
lmer_aras_tidy <- res_2_df(model = lmer_aras, model_R2 = lmer_aras_bs, model_name = "aras")

# without random slope:
lmer_aras_noslope <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                          formula = as.formula("n_child_10000 ~ year_std + any_ara_pct_wt + any_ara_pct_bt + (1|la_code)"), REML = TRUE)

G <- -2*logLik(lmer_aras_noslope) + 2*logLik(lmer_aras)
pchisq(as.numeric(G), df = 2, lower.tail = FALSE) # p < 0.001

#------------------------#
# Full models         ----
#------------------------#

# full model with random intercept for LA and random slope for year:
lmer_full <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                 formula = as.formula(paste("n_child_10000 ~ year_std",
                                            "any_ara_pct_wt", "any_ara_pct_bt",
                                            "teenbirth_2010",
                                            "imd_10pc_2010",
                                            "lowbw_2010",
                                            "ca_feud_2010",
                                            "total_2010",
                                            "lone_parent_pct",
                                            "violent_crime",
                                            "(year_std|la_code)",
                                            sep = "+")), REML = T)

# get bootstrap estimate of pseudo R-squared measures (and confidence intervals):
lmer_full_bs <- R2_bootstrap(model = lmer_full)

# create dataframe of model results:
lmer_full_tidy <- res_2_df(model = lmer_full, model_R2 = lmer_full_bs, model_name = "full")

# test equality of with/between effects
linearHypothesis(lmer_full, "any_ara_pct_wt = any_ara_pct_bt") # insufficient evidence to reject H0

# without random slope:
lmer_full_noslope <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                          formula = as.formula(paste("n_child_10000 ~ year_std",
                                                     "any_ara_pct_wt", "any_ara_pct_bt",
                                                     "teenbirth_2010",
                                                     "imd_10pc_2010",
                                                     "lowbw_2010",
                                                     "ca_feud_2010",
                                                     "total_2010",
                                                     "lone_parent_pct",
                                                     "violent_crime",
                                                     "(1|la_code)",
                                                     sep = "+")), REML = TRUE)

# compare random intercept model with random coefficient model:
G <- -2*logLik(lmer_full_noslope) + 2*logLik(lmer_full)
pchisq(as.numeric(G), df = 2, lower.tail = FALSE)# p < 0.001

# full model without disagregation of within/between LA effects of maternal ARA:
lmer_full_nodisag <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                         formula = as.formula(paste("n_child_10000 ~ any_ara_pct + year_std",
                                                    "teenbirth_2010",
                                                    "imd_10pc_2010",
                                                    "lowbw_2010",
                                                    "ca_feud_2010",
                                                    "total_2010",
                                                    "lone_parent_pct",
                                                    "violent_crime",
                                                    #"prop10_imd_pct",
                                                    "(year_std|la_code)",
                                                    sep = "+")), REML = T)

lmer_full_nodisag_bs <- R2_bootstrap(model = lmer_full_nodisag)

lmer_full_nodisag_tidy <- res_2_df(model = lmer_full_nodisag, model_R2 = lmer_full_nodisag_bs,
                                   model_name = "full_nodisag")

# higher order year_std terms:
lmer_full_nodisag_2 <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                          formula = as.formula(paste("n_child_10000 ~ any_ara_pct + poly(year_std, 2, raw = TRUE)",
                                                     "teenbirth_2010",
                                                     "imd_10pc_2010",
                                                     "lowbw_2010",
                                                     "ca_feud_2010",
                                                     "total_2010",
                                                     "lone_parent_pct",
                                                     "violent_crime",
                                                     #"prop10_imd_pct",
                                                     "(poly(year_std, 2, raw = TRUE)|la_code)",
                                                     sep = "+")), REML = T)

lmer_full_nodisag_2_bs <- R2_bootstrap(model = lmer_full_nodisag_2)

lmer_full_nodisag_2_tidy <- res_2_df(model = lmer_full_nodisag_2, model_R2 = lmer_full_nodisag_2_bs,
                                   model_name = "full_nodisag_poly2")

lmer_full_nodisag_3 <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                            formula = as.formula(paste("n_child_10000 ~ any_ara_pct + poly(year_std, 3, raw = TRUE)",
                                                       "teenbirth_2010",
                                                       "imd_10pc_2010",
                                                       "lowbw_2010",
                                                       "ca_feud_2010",
                                                       "total_2010",
                                                       "lone_parent_pct",
                                                       "violent_crime",
                                                       #"prop10_imd_pct",
                                                       "(poly(year_std, 3, raw = TRUE)|la_code)",
                                                       sep = "+")), REML = T)

lmer_full_nodisag_3_bs <- R2_bootstrap(model = lmer_full_nodisag_3)

lmer_full_nodisag_3_tidy <- res_2_df(model = lmer_full_nodisag_3, model_R2 = lmer_full_nodisag_3_bs,
                                     model_name = "full_nodisag_poly3")

#-------------------------#
# save analysis models ----
#-------------------------#
model_list <- list(lmer_null,
                   lmer_null_noslope,
                   lmer_aras,
                   lmer_aras_noslope,
                   lmer_full,
                   lmer_full_noslope,
                   lmer_full_nodisag,
                   lmer_full_nodisag_2,
                   lmer_full_nodisag_3,
                   lmer_null_bs,
                   lmer_aras_bs,
                   lmer_full_bs,
                   lmer_full_nodisag_bs,
                   lmer_full_nodisag_2_bs,
                   lmer_full_nodisag_3_bs)
save(list = "model_list", file = "N:/Documents/Data & Programming/outputs/la variation/modelling/model_list.rda")

# save raw results:
write_csv(lmer_null_tidy, "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_null_results.csv")
write_csv(lmer_aras_tidy, "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_aras_results.csv")
write_csv(lmer_full_tidy, "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_wslope_results.csv")
write_csv(lmer_full_nodisag_tidy, "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_nodisag_results.csv")

#----------------------#
# model diagnostics ----
#----------------------#

# examine level-1 residuals using histograms and quantile-quantile plots----

# null model:
png("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_null_lvl1_dist.png",
    width = 300, height = 150, res = 300, units = "mm")
par(mfrow = c(1,2))
lmer1_hist(model = lmer_null)
lmer1_qnorm(model = lmer_null)
dev.off()

# aras only model:
png("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_aras_lvl1_dist.png",
    width = 300, height = 150, res = 300, units = "mm")
par(mfrow = c(1,2))
lmer1_hist(model = lmer_aras)
lmer1_qnorm(model = lmer_aras)
dev.off()

# full model:
png("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_lvl1_dist.png",
    width = 300, height = 150, res = 300, units = "mm")
par(mfrow = c(1,2))
lmer1_hist(model = lmer_full)
lmer1_qnorm(model = lmer_full)
dev.off()

# full model - ara disaggregated:
png("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_nodisag_lvl1_dist.png",
    width = 300, height = 150, res = 300, units = "mm")
par(mfrow = c(1,2))
lmer1_hist(model = lmer_full_nodisag)
lmer1_qnorm(model = lmer_full_nodisag)
dev.off()

#--------------------------------------#
# interaction between time and ARA ----
#--------------------------------------#

# extra packages:
library(interplot)

# interaction model:
lmer_full_int <- lmer(data = filter(model_dat, any_bad_linkrate == 0 & any_bad_bwrate ==0),
                      formula = as.formula(paste("n_child_10000 ~ year_std*any_ara_pct",
                                                 "teenbirth_2010",
                                                 "imd_10pc_2010",
                                                 "lowbw_2010",
                                                 "ca_feud_2010",
                                                 "total_2010",
                                                 "lone_parent_pct",
                                                 "violent_crime",
                                                 "(1|la_code)",
                                                 sep = "+")), REML = T)

# get bootstrap estimate of pseudo R-squared measures (and confidence intervals):
lmer_full_int_bs <- R2_bootstrap(model = lmer_full_int)

# collate results:
lmer_full_int_tidy <- res_2_df(model = lmer_full_int, model_R2 = lmer_full_int_bs, model_name = "full_int") #%>%
  #mutate_at(vars(contains("R2"), contains("conf"), "estimate"),list(~sprintf("%.2f", round(.,2)))) %>%
  #mutate(ci = paste0(conf.low," to ", conf.high))

# save results:
write_csv(lmer_full_int_tidy,
          "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_int_results.csv")

# model diagnostics:
lmer1_hist(model = lmer_full_int)
lmer1_qnorm(model = lmer_full_int)

# predictive values:
range(predict(lmer_full_int))

# plot interaction effect:
int_data <- interplot(lmer_full_int, var2= "year_std", var1 = "any_ara_pct",plot = FALSE) %>%
  mutate_at(vars("coef", "ub", "lb"),list(~sprintf("%.2f", round(.,2)))) %>%
  mutate(ci = paste0(ub," to ", lb))

write_csv(int_data, "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_full_int_results_byyr.csv")

interplot(lmer_full_int, var2= "year_std", var1 = "any_ara_pct",
                    point = TRUE) +
  scale_x_continuous("",
                     breaks = seq(0,7,1), labels = paste(seq(2006,2013,1),
                                                         substr(seq(2007,2014),3,4), sep = "/")) +
  scale_y_continuous("Estimated coefficient\n",
                     breaks = seq(-2,4,2), limits = c(-2.2,5.2)) +
  geom_hline(aes(yintercept = 0), linetype = "dashed",colour = "grey60") +
  #geom_text(data = int_data, aes(x = year_std, y = as.numeric(lb), label = paste0(coef, "\n(", ci, ")")),
  #          vjust = -1,
  #          size = 5) +
  theme_classic() +
  geom_errorbar(data = int_data, aes(x = year_std, ymin = as.numeric(ub), ymax = as.numeric(lb)),
                inherit.aes = FALSE, width = 0.2, colour = "#5379a9ff", size = 1.25) +
  geom_point(size = 5, data = int_data, aes(x = year_std, y = as.numeric(coef)),
             shape = 21, fill = "white", colour = "#5379a9ff") +
  #scale_fill_brewer(palette = "YlGnBu") +
  labs(caption = "") +
  guides(fill = "none") +
  #geom_text()
  theme(axis.text = element_text(size = 16),
        axis.title = element_text(size = 16),
        strip.text = element_text(size = 18),
        legend.title=element_text(size=18),
        legend.text=element_text(size=18),
        title = element_text(size = 18))

#require("svglite")
ggsave("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_interaction_plot_poster.png", units = "mm",
       height = 150, width = 300, dpi = 300)

#-----------------------------------------------#
# plot fixed-effect coeffs and their 95% CIs ----
#-----------------------------------------------#
# create variable labels:
lbls <- data.frame(label = c("any_ara_pct_bt" = "Between-cluster: % of live births with maternal history of ARA",
                             "any_ara_pct_wt" = "Within-cluster: % of live births with maternal history of ARA",
                             "any_ara_pct" = "% of live births with maternal history of ARA",
                             "lone_parent_pct" = "% of dependent child households with lone parent",
                             "teenbirth_2010" = "% of live births where mother < 20 years old",
                             "imd_10pc_2010" = "% of live births where maternal LSOA history within 10% most deprived LSOAs in England",
                             "prop10_imd_pct" = "Proportion of LSOAs in LA in 10% most deprived in England",
                             "imd_quint" = "LA IMD quintile - 2010",
                             "total_2010" = "LA population size (000s)",
                             "violent_crime" = "Rate of violent crime (per 100 LA residents)",
                             "lowbw_2010" = "% of live births with LBW",
                             "ca_feud_2010" = "% of live births where child has a CCC",
                             "year_std" = "Time (financial year)",
                             "year_std:any_ara_pct" = "Interaction: Time & % of live births with maternal history of ARA")) %>%
  rownames_to_column(var = "term")

# create dataframe to plot: --------------------------------------------------------------------------------#
lmer_plot <- bind_rows(lmer_null_tidy,
                       lmer_aras_tidy,
                       lmer_full_tidy,
                       lmer_full_nodisag_tidy,
                       lmer_full_int_tidy) %>%
  filter(group == "fixed" & !grepl("intercept", tolower(term))) %>%
  rowwise() %>%
  mutate(model_num = case_when(grepl("null", model) ~ "Model 1",
                               grepl("aras", model) ~ "Model 2",
                               model == "full" ~ "Model 3",
                               grepl("disag", model) ~ "Model 4",
                               grepl("full_int", model) ~ "Model 5"),
         model_ltr = case_when(grepl("null", model) ~ "a",
                               grepl("aras", model) ~ "b",
                               model == "full" ~ "c",
                               grepl("disag", model) ~ "d",
                               grepl("full_int", model) ~ "e"),
         # facet_lbl = paste(model_num,
         #                   paste("AIC = ",round(AIC)),
         #                   paste0("R2m = ", sprintf("%.2f", round(R2m_median,2)),
         #                         " (95% CI: ", sprintf("%.2f", round(R2m_ll,2))," to ", sprintf("%.2f", round(R2m_ul,2)),")"),
         #                   paste0("R2c = ", sprintf("%.2f", round(R2c_median,2)),
         #                         " (95% CI: ", sprintf("%.2f", round(R2c_ll,2))," to ", sprintf("%.2f", round(R2c_ul,2)),")"),
         #                   sep = "\n"),
         facet_lbl_noaic = ifelse(model_num %in% c("Model 3", "Model 4", "Model 5"),
                                  paste0(model_num, "*"), model_num), #as.character(as.expression(bquote(.(model_num)^.(model_ltr)))),
         facet_lbl = ifelse(model_num %in% c("Model 3", "Model 4", "Model 5"),
                            paste0(model_num, "*", "\nAIC = ", round(AIC)),
                            paste0(model_num, "\nAIC = ", round(AIC))),
         aic_lbl = paste0("NULL[AIC:~", AIC,"]")
         ) %>%
  arrange(model_num, term) %>%
  droplevels()

# create rows for each coefficient in each model
dummy <- tidyr::expand(lmer_plot, nesting(term), facet_lbl)

# merge on dummy rows and create order for model terms:
lmer_plot <- left_join(dummy, lmer_plot, by = c("facet_lbl", "term")) %>%
  left_join(lbls, by = "term") %>%
  mutate(order = case_when(term == "year_std" ~ 12,
                           grepl("ARA", label ) & grepl("Within", label) ~ 11,
                           grepl("ARA", label ) & grepl("Between", label) ~ 10,
                           term == "any_ara_pct"  ~ 9,
                           term == "year_std:any_ara_pct" ~ 8,
                           grepl("< 20", label) ~ 7,
                           grepl("LSOA", label) ~ 6,
                           grepl("CCC", label) ~ 5,
                           grepl("LBW", label) ~ 4,
                           grepl("size", label) ~ 3,
                           grepl("lone", label) ~ 2,
                           TRUE ~ 1),
         tvc = ifelse(order > 7, 0, 1), # group terms by whether they are time-varying or not
         int = ifelse(model_num == "Model 5" & term != "year_std:any_ara_pct", 1, 0),
         ltr = case_when(model_num == "Model 5" & term == "year_std" ~ "a",
                         model_num == "Model 5" & term == "any_ara_pct" ~ "b",
                         TRUE ~ NA_character_),
         label = reorder(as.factor(as.character(label)), order),
         text = ifelse(is.na(estimate) | model_num != "Model 4", # create text label of estimate and CI for model 4
                       "", paste0("(", sprintf("%.2f", round(estimate,2)),", ",
                                  sprintf("%.2f", round(conf.low,2)), " to ",
                                          sprintf("%.2f", round(conf.high, 2)), ")")),
         text2 = ifelse(grepl("ara|year",term) & !is.na(estimate), # create text label of estimate and CI for model 4
                       paste0("'", sprintf("%.2f", round(estimate,2))," (",
                                  sprintf("%.2f", round(conf.low,2)), " to ",
                                  sprintf("%.2f", round(conf.high, 2)), ")'"), ""),
         text3 = ifelse(!is.na(ltr), paste(text2, "^", ltr), text2))

# save model plot data:
write.csv(lmer_plot, "N:/Documents/Data & Programming/outputs/la variation/modelling/final_results.csv")
#-------------------------------------------------------------------------------------------------------------#

# caption for plot ----
# lmer_caption <- lmer_plot %>%
#   mutate(cap  = paste(paste0(model_ltr, ": ", "AIC = ",round(AIC)),
#                       paste0("R2m = ", sprintf("%.2f", round(R2m_median,2)),
#                              " (95% CI: ", sprintf("%.2f", round(R2m_ll,2))," to ", sprintf("%.2f", round(R2m_ul,2)),")"),
#                       paste0("R2c = ", sprintf("%.2f", round(R2c_median,2)),
#                              " (95% CI: ", sprintf("%.2f", round(R2c_ll,2))," to ", sprintf("%.2f", round(R2c_ul,2)),")"),
#                       sep = "; ")) %>%
#   filter(!is.na(model_ltr)) %>%
#   distinct(cap) %>%
#   arrange(cap) %>%
#   unlist()

# create plot object: -------------------------------------------------------------------
ggplot(data = lmer_plot,
       aes(y = reorder(label, order),
           x = estimate,
           xmin = conf.low,
           xmax = conf.high,
           color = as.factor(tvc))) +
  facet_wrap(~facet_lbl,
             scales = "free_x",
             ncol = 4,
             labeller = "label_parsed") +
  geom_vline(xintercept = 0,
             linetype = "dashed",
             size = 0.6,
             colour = "grey60",
             alpha = 0.5) +
  geom_errorbarh(height = 0.2,
                 size = 1) +
  geom_point(shape = 19,
             fill = "white",
             size = 2.5) +
  geom_text(aes(label = text),
            position = position_dodge(width = 1),
            vjust = -1,
            show.legend = FALSE,
            color = "black",
            size = 5.5) +
  scale_x_continuous(limits = c(-9, 23),
                     breaks = c(-5,0,5,10,15,20)) +
  scale_y_discrete(labels = function(x) str_wrap(x, width = 25)) +
  ylab("") +
  xlab("\nDifference in the outcome per unit change in predictor") +
  labs(caption = paste0("\n\nAIC = Akaike information criterion; R2m = marginal pseudo R-squared; R2c = conditional pseudo R-squared.\n\n",
                       paste(lmer_caption, collapse = "\n\n"))) +
  theme_classic() +
  theme(axis.text = element_text(size = 16),
        strip.text = element_text(size = 18),
        legend.text = element_text(size = 16),
        axis.title = element_text(size = 16),
        plot.caption = element_text(hjust = 0, size = 16),
        panel.border = element_rect(color="black", fill=NA)) +
  scale_colour_manual("",
                      labels = c("time-varying\n(2006/07 to 2013/14)\n ",
                                 "non-time-varying\n(2010/11 data only)\n "),
                      values = c("#1F78B4", "#D55E00"))

ggsave("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_coef_plots.png",
       units = "mm", height = 350, width = 550, dpi = "print")

#--------------------------------------#
# model diagnostic plot ----
#--------------------------------------#
model_check <-
  ggplot(data = filter(distinct(lmer_plot, model_num, .keep_all = TRUE), !is.na(model_num)),
       aes(group = 1)) +
  # facet_wrap(~facet_lbl_noaic,
  #            scales = "free_x",
  #            ncol = 5) +
  # geom_point(aes(x = model_num,
  #                y = AIC/10000,
  #                shape = "AIC",
  #                colour = "AIC"),
  #            size = 3) +
  geom_line(aes(x = model_num,
                 y = R2m_median,
                 colour = "% variation in the outcome explained by the fixed-effects",
                group = 1)) +
  geom_line(aes(x = model_num,
                 y = R2c_median,
                 colour = "% variation in the outcome explained by the whole model",
                group = 1)) +
  # geom_line(aes(x = model_num,
  #                y = AIC/10000,
  #               colour = "AIC")) +
  geom_errorbar(aes(x = model_num,
                    ymin = R2m_ll,
                    ymax = R2m_ul),
                    #colour = "% variation in the outcome explained by the fixed-effects"),
                width = 0.02) +
  geom_errorbar(aes(x = model_num,
                    ymin = R2c_ll,
                    ymax = R2c_ul),
                    #colour = "% variation in the outcome explained by the whole model"),
                width = 0.02) +
  geom_point(aes(x = model_num,
                 y = R2m_median,
                 fill = "% variation in the outcome explained by the fixed-effects",
                 shape = "% variation in the outcome explained by the fixed-effects"),
             size = 3) +
  geom_point(aes(x = model_num,
                 y = R2c_median,
                 fill = "% variation in the outcome explained by the whole model",
                 shape = "% variation in the outcome explained by the whole model"),
             size = 3) +
  theme_classic() +
  geom_text(aes(x = model_num, y = R2m_ul, label = paste0(round(R2m_median*100),"% (",
                                                          round(R2m_ll*100),
                                                          "-",
                                                          round(R2m_ul*100),"%)")),
                #colour = "% variation in the outcome explained by the fixed-effects"),
            vjust = -1,
            show.legend = FALSE,
            size = 4) +
  geom_text(aes(x = model_num, y = R2c_ul, label = paste0(round(R2c_median*100),"% (",
                                                          round(R2c_ll*100),
                                                          "-",
                                                          round(R2c_ul*100),"%)")),
                #colour = "% variation in the outcome explained by the whole model")
            vjust = -1,
            show.legend = FALSE,
            size = 4) +
  # geom_text(aes(x = model_num, y = AIC/10000, label = round(AIC),
  #               colour = "AIC"), vjust = -0.9,
  #           show.legend = FALSE,
  #           size = 6) +
  scale_shape_manual("", values = c(21, 22)) +
  #scale_color_manual("", values = c("#D95F02", "#1B9E77", "#7570B3")) +
  scale_color_manual("", values = c("#bbcc33", "dodgerblue4")) +
  scale_fill_manual("", values = c("#bbcc33", "dodgerblue4")) +
  scale_y_continuous("",labels = scales::percent,
                     limits =c(0,1.10), breaks = seq(0,1,0.2)) +
                     #sec.axis = sec_axis(~.*10000, name = "AIC\n")) +
  scale_x_discrete("") +
  #labs(caption = "Note: 95% confidence intervals given in brackets.", hjust = 0) +
  theme(axis.text = element_text(size = 12),
        plot.title = element_text(size = 18),
        strip.text = element_text(size = 14),
        axis.title = element_text(size = 14),
        #axis.text = element_text(size = 20),
        #plot.title = element_text(size = 26),
        #strip.text = element_text(size = 20),
        #axis.title = element_text(size = 20),
        legend.text = element_text(size = 14),
        legend.spacing.x = unit(1.0, 'cm'),
        axis.text.x = element_blank(),
        axis.ticks.x = element_blank(),
        panel.border = element_rect(color="black", fill=NA),
        legend.position = "bottom")

model_check
ggsave("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_prop_explained.png",
       units = "mm", height = 210, width = 250, dpi = "print")

#--------------------------------------#
# only tvc plot ----
#--------------------------------------#

tvc_plot <- ggplot(data = filter(lmer_plot, tvc == 0 & !is.na(model_num) ),
                 aes(y = reorder(label, order),
                     x = estimate,
                     xmin = conf.low,
                     xmax = conf.high#,
                     #colour = as.factor(int)
                 )) +
  facet_wrap(~ facet_lbl,
             ncol = 5,
             scales = "free_x") +
  geom_vline(xintercept = 0,
             linetype = "dashed",
             size = 0.6,
             colour = "grey60",
             alpha = 0.3) +
  geom_errorbarh(height = 0.1) +
                 #size = 1,
                 #colour = "#44bb99") +
  geom_point(size = 3,
             shape = 21,
             fill = "#44bb99") +
  geom_text(aes(label = text3),
            #position = position_dodge(width = 1),
            vjust = -1.2,
            show.legend = FALSE,
            color = "black",
            size = 4,
            parse = TRUE) +
  scale_x_continuous(limits = c(-5, 18),
                     breaks = c(-5,0,5,10,15)) +
  scale_y_discrete(labels = function(x) str_wrap(x, width = 25)) +
  ylab("") +
  xlab("\nDifference in the outcome per unit change in predictor\n") +
  #labs(caption = paste0("\n\nAIC = Akaike information criterion; R2m = marginal pseudo R-squared; R2c = conditional pseudo R-squared.\n\n",
  #                      paste(lmer_caption, collapse = "\n\n"))) +
  theme_classic() +
  theme(axis.text = element_text(size = 12),
        plot.title = element_text(size = 18),
        strip.text = element_text(size = 14),
        axis.title = element_text(size = 14),
         panel.border = element_rect(color="black", fill=NA)) #+
# scale_colour_manual("",
#                     labels = c("interaction model\n ",
#                                "non-time-varying\n(2010/11 data only)\n "),
#                     values = c("#1F78B4", "#D55E00"))

tvc_plot
ggsave( "N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_coef_plots_tvconly.png",
       units = "mm", height = 200, width = 550, dpi = "print")

# plot tvc only and model diagnostics together:
svg(filename = "N:/Documents/Data & Programming/outputs/la variation/modelling/figure2.svg",
    width = 18, height = 11.5)
egg::ggarrange(tvc_plot + labs(title = "A"), model_check + labs(title = "B"), ncol=1)
dev.off()

#--------------------------------------#
# Final model only plot ----
#--------------------------------------#
ggplot(data = filter(lmer_plot, grepl("disag", model)),
       aes(y = reorder(label, order),
           x = estimate,
           xmin = conf.low,
           xmax = conf.high,
           color = as.factor(tvc))) +
  geom_vline(xintercept = 0,
             linetype = "dashed",
             size = 0.6,
             colour = "grey60",
             alpha = 0.5) +
  geom_errorbarh(height = 0.2,
                 size = 1) +
  geom_point(shape = 19,
             fill = "white",
             size = 2.5) +
  geom_text(aes(label = ifelse(term == "any_ara_pct", text, "")),
            #position = position_dodge(width = 1),
            vjust = -1,
            show.legend = FALSE,
            color = "black",
            size = 5.5) +
  scale_x_continuous(limits = c(-9, 23),
                     breaks = c(-5,0,5,10,15,20)) +
  scale_y_discrete(labels = function(x) str_wrap(x, width = 25)) +
  ylab("") +
  xlab("\nDifference in the outcome per unit change in predictor") +
  #labs(caption = paste0("\n\nAIC = Akaike information criterion; R2m = marginal pseudo R-squared; R2c = conditional pseudo R-squared.\n\n",
  #                      paste(lmer_caption, collapse = "\n\n"))) +
  theme_classic() +
  theme(axis.text = element_text(size = 16),
        strip.text = element_text(size = 18),
        legend.text = element_text(size = 16),
        axis.title = element_text(size = 16),
        plot.caption = element_text(hjust = 0, size = 16),
        panel.border = element_rect(color="black", fill=NA)) +
  scale_colour_manual("",
                      labels = c("time-varying\n(2006/07 to 2013/14)\n ",
                                 "non-time-varying\n(2010/11 data only)\n "),
                      values = c("#1F78B4", "#D55E00"))

ggsave("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_coef_plot_final.png",
       units = "mm", height = 300, width = 350, dpi = "print")

#---------------------------#
# plot for presentations ----
#---------------------------#
library(grid)
library(gridExtra)
library(gtable)
library(ggpubr)

ggarrange(ggplotGrob(tvc_plot + labs(title = "A")),
          ggplotGrob(model_check + labs(title = "B")),
          ncol =1, nrow = 2,
          align = "v",
          heights = c(1,0.9))
ggsave("N:/Documents/Data & Programming/outputs/la variation/modelling/lmer_plot.png",
       units = "mm", height = 410, width = 600, dpi = "print")


#forecast only AR models  only
rm(list=ls())

#function returns only 1 value
pre<-function(post,p){
#post-starting value ,p-pvalue
#for all p coeffienct calculation
sham_sunder<-NULL #for adding p coeff
for( rum in seq_along(p)) {
sham_sunder<-sum(sham_sunder,p[rum]*t[post-rum] )  #left term of equation
}
z=sham_sunder+t[post]
return(z)
}  #function

v<-read.table("K:\\Project_Data_Sagar\\Halwa\\BADMASHI\\pot_tamasha\\chocolate_potential_x.txt",skip=1)
w<-v[1:20,2:4] #taking only 3 columns & 20 rows (excludes X's row & AA column)

ohm<-NULL #for appending p-zero order file names
dhino_dhino<-NULL #for appending file_name & r_square
woh_hai<-list.files("C:\\kindle_patrika\\PDB\\pdb_MoserFear\\")  #all stationary files

for ( bari_bari in woh_hai){
vaishakhi<-substr(bari_bari,1,nchar(bari_bari)-16) #contains file name
pyar<-read.table(paste("C:\\kindle_patrika\\PDB\\pdb_MoserFear\\",bari_bari,sep=""),skip=1 )  #opening potential file

shukrana<-as.character(pyar$V2)  #original AA seq
fikrana<-pyar$V3  #original AA States
muskurana<-pyar$V4  #original potential values

t<-pyar$V4 #assigning ori pot to another variable(4 convenience)
ishare<-ar(t)  #calculates AR model
hari_haran<- ishare$ar  #No of ar coefficient

#if AR p-order is zero checking
if(!length(hari_haran)==0){

#forecasting
sajna<-NULL #for appending forecasted value
shraddha<-length(hari_haran) #stoes length of p
for( phulwa in (shraddha+1):length(t)){
sajna<-append(sajna,pre(phulwa,hari_haran))   #1---> starting vlaue ,2---> p-value
}
sajna<-append(t[seq_along(hari_haran)],sajna,) #appending remaining element(has final predicted values)

x<-t   #original value
y<-sajna    #predicted value
f<-lm(y~x)
gen_r_square<-summary(f)$r.squared

#wikipedia method of R-square
r_wiki_square<-1- ( sum((x-y)^2)/ sum((x-mean(x))^2)  )

#for AA & state predictions
motivated<-NULL # for appending AA seq
robotic_vkj<-NULL #for appending states of AA

for( papa in sajna ){
muskurahat=papa     #taking predicted potential value for comparing with matrix
z<-abs(w-muskurahat)==min(abs(w-muskurahat))   #comparing with matrix
great_thumb<-as.character(v[c(which(z[,1]),which(z[,2]),which(z[,3])),]$V1)    #taking single AA character
motivated<-append(motivated,great_thumb)  #appending to predicted AA seq
jannatein<-which(c(length(which(z[,1])),length(which(z[,2])),length(which(z[,3])) )==1)   #gives state of AA
robotic_vkj<-append(robotic_vkj,jannatein)    #appending predicted AA states
}

#Assembling data & writing
sab_kuch<-append(rep(NA,length(t)-1),r_wiki_square,after=0)   #last column r_square with NA
behind<-cbind(t,sajna,sab_kuch,shukrana,motivated,fikrana,robotic_vkj,deparse.level=0)  #written matter for each file (ori_pot,pred_pot,r_sq,ori_AA,pred_AA,ori_states,pred_states)

write.table(behind,paste("C:\\Users\\msc2\\Desktop\\Halwa\\BADMASHI\\forecasting\\AR\\forecast_AR_stationary\\",vaishakhi,".txt",sep=""),col.names=c("ori_pot","pred_pot","R_Square","ori_AA","pred_AA","ori_states","pred_states"),row.names=F)
dhino_dhino<-rbind(dhino_dhino,cbind(vaishakhi,gen_r_square,r_wiki_square,deparse.level=0),deparse.level=0) #file_name,General_R_sq,R_wiki_Sq
}  #p order checking

else{  ohm<-append(ohm,vaishakhi)}

} #whole file iterating

#writing p-zero order files
if(!is.null(ohm)){write.table(ohm,"C:\\Users\\msc2\\Desktop\\Halwa\\BADMASHI\\forecasting\\AR\\AR_stat_p_zero_models.txt",col.names="P-value zero files",row.names=F)}

#writing info file
write.table(dhino_dhino,"C:\\Users\\msc2\\Desktop\\Halwa\\BADMASHI\\forecasting\\AR\\AR_forecast_stat_R_sq_info.txt",col.names=c("PDB_id","Gen_R_square","Wiki_R_square"),row.names=F)


#Usage:
#Rscript cell_cluster_multi.r cluster_number
library("monocle")
library("reshape")
library('ggplot2')
args <- commandArgs(trailingOnly=TRUE)
if(length(args)==0)
{
	print("Usage: Rscript cell_cluster.r cluster_number cluster_labels")
}
cluster_number <- args[1]
HSMM <- readRDS('post_quality_control_data.rds')
HSMM <- clusterCells(HSMM, num_clusters=as.numeric(cluster_number))#, tol = 1e-6,max_components = 10, verbose=TRUE,param.gamma=300)
saveRDS(HSMM,'post_quality_control_data_multi.rds')
#HSMM <- readRDS('post_quality_control_data_multi.rds')
print(plot_cell_trajectory(HSMM, color="Cluster"))
ggsave("cell_cluster_plot.png")
ggsave("cell_cluster_plot.tiff", width = 7, height = 7,dpi=600)
cell_name <- pData(HSMM)$name
cluster_num <- pData(HSMM)$Cluster
write.table(cell_name,'cluster_name_multi.txt',sep="\t",quote=FALSE,row.names=FALSE,col.names=FALSE)
write.table(cluster_num,'cluster_clusternum_multi.txt',sep="\t",quote=FALSE,row.names=FALSE,col.names=FALSE)
dev.off()


timer_naive = function(hours, minutes, seconds)
{
  while (hours > 0 || minutes > 0 || seconds > 0)
  {
    if (seconds == 0)
    {
      if (minutes > 0)
      {
        minutes = minutes - 1
        seconds = 60
      }
      else if (minutes == 0 && hours > 0)
      {
        hours = hours - 1
        minutes = 59
        seconds = 60
      }
    }

    seconds = seconds - 1
    cf = list(hours=hours, minutes=minutes, seconds=seconds)
    print_clockface(cf, "time remaining: ")
    Sys.sleep(1)
  }

  cat("\n")
  invisible(TRUE)
}


library(watershed)
library(raster)
library(sf)
library(ggplot2)
library(flume)

# dem1 = raster("data/elevation/eu_dem_v11_E50N10.TIF")
# dem2 = raster("data/elevation/eu_dem_v11_E50N20.TIF")
# dem = merge(dem1, dem2, file = )
# ext = extent(5073000, 5350000, 1874000, 2035000)
# dem = crop(dem, ext, file="data/elevation/vjosa_dem.tif")
thresh = 5e7
dem = raster("data/elevation/vjosa_dem.tif")
vj = delineate(dem, threshold = thresh, reach_len=7000)
Tp = pixel_topology(vj)

## for plotting
# dem_cr = crop(dem, vj)
# dem_df = as.data.frame(rasterToPoints(dem_cr))

# remove small headwater reaches
vj_v = vectorise_stream(vj[["stream"]], Tp=Tp)
# Tr = reach_topology(vj, Tp=Tp)
# hws = which(colSums(Tr != 0) == 0)
# lens = st_length(vj_v)
# hw_drop = hws[which(lens[hws] < units::set_units(3500, "m"))]
# vj[["stream"]][values(vj[["stream"]]) %in% hw_drop] = NA
# Tp = pixel_topology(vj)

# # renumber the reaches
# rids = unique(values(vj$stream))
# rids = sort(rids) ## also gets rid of NA
# vj$stream = match(vj$stream, rids)
# # make into vector to do further revisions in QGIS
# vj_v = vectorise_stream(vj[["stream"]], Tp=Tp)
# vj_v$len = st_length(vj_v)
# st_write(vj_v, "data/vjosa_stream.shp", append = FALSE)

# ggplot() + geom_raster(data=dem_df, aes(x = x, y = y, fill = vjosa_dem)) +
# 	geom_sf(data = vj_v, aes(colour=catchment)) + scico::scale_fill_scico(palette = "turku")

discharge = read.csv("data/discharge.txt")
discharge = discharge[discharge$watershed == "Vjosa" & discharge$season == "Fall",]
discharge = st_as_sf(discharge, coords = c('x', 'y'), crs=3035)
# ggplot() + geom_raster(data=dem_df, aes(x = x, y = y, fill = vjosa_dem)) +
	# geom_sf(data = vj_v, aes(colour=catchment)) + scico::scale_fill_scico(palette = "turku") +
	# geom_sf(data = discharge)

discharge = st_snap(discharge, vj_v, tolerance=100)
discharge$ca = catchment(vj, type = "points", y = st_coordinates(discharge))
discharge = discharge[discharge$ca > thresh,]

vj_v$catchment = catchment(vj, type = 'reach', Tp = Tp)
vj_v = cbind(vj_v, hydraulic_geometry(discharge$ca, discharge$discharge, vj_v$catchment))
Tp_r = reach_topology(vj, Tp)

# save intermediate shit in case of a crash
vj_v$CA = units::set_units(vj_v$CA, "km^2")
saveRDS(vj_v, "res/vjosa_stream.rds")
saveRDS(discharge, "res/vjosa_discharge.rds")
writeRaster(vj, "res/vjosa_watershed.grd", overwrite = TRUE)
saveRDS(Tp, "res/vjosa_pix_topo.rds")
saveRDS(Tp_r, "res/vjosa_each_topo.rds")
layout = st_coordinates(lwgeom::st_startpoint(vj_v))
vj_v$Q = units::drop_units(vj_v$Q)
vj_v$z = units::drop_units(vj_v$z)
vj_v$w = units::drop_units(vj_v$w)

# one last attempt to edit the topology manually to get rid of tiny reaches
# Tp_r_bak = Tp_r
# vj_v_bak = vj_v

vj_v = vj_v_bak
Tp_r = Tp_r_bak
layout = st_coordinates(lwgeom::st_startpoint(vj_v))

rdel = list(
	c(2:19,308),
	21:33,
	c(35:66,306:307),
	c(35:66,306:307),
	c(81:89,305),
	91:92,
	c(94:101,303:304),
	c(94:101,303:304),
	c(104:127, 302),
	c(129:139, 299:301),
	141:153,
	141:153,
	c(157:174, 298, 155),
	c(157:174, 298, 155),
	178:187,
	178:187,
	191,
	191,
	273:277,
	273:277,
	287,
	c(193:202, 271),
	c(193:202, 271),
	206,
	206,
	c(210:212, 214)
)

rconn = list(
	c(20, 1),
	c(34,20),
	c(67,34),
	c(80,34),
	c(90,80),
	c(93,90),
	c(102,93),
	c(103,93),
	c(128,103),
	c(140,128),
	c(154,140),
	c(156,140),
	c(176,156),
	c(177,156),
	c(189,177),
	c(190,177),
	c(272,190),
	c(192,190),
	c(279,272),
	c(282,272),
	c(286,283),
	c(203,192),
	c(270,192),
	c(208,207),
	c(209,207),
	c(213,209))

lens = rowSums(Tp_r)
for(i in 1:length(rdel)) {
	dd = rdel[[i]]
	cc = rconn[[i]]
	Tp_r[cc[1], cc[2]] = sum(lens[dd]) + lens[cc[1]]
}

rdel = unique(unlist(rdel))
Tp_r = Tp_r[-rdel, -rdel]
vj_v = vj_v[-rdel,]
layout = st_coordinates(lwgeom::st_startpoint(vj_v))

vjosa = river_network(Tp_r, discharge = vj_v$Q, area = vj_v$z*vj_v$w, length = 7000,
	layout = layout)
plot(vjosa, edge.arrow.size=0, vertex.size=2, vertex.label=NA)

## notes
# there are many small reaches in the main stem. These need to be combined into simpler reaches


# for manipulating  rodent movement data in MARK
## analyze survival and dispseral probabilities for species
# Use multistrata models that give S(survival), p(capture probability) and Psi(transition probability)
# Will compare data among species and guilds in post-hoc analyses.
# Input files (.txt) generated using rodent_wrapper.r and movement_fxns.r

# Do psi and S significantly differ among species/guilds?

runRMARK = function(wd){

#---------------------------------------------------------------------------------
#          bring in the data and source files
#---------------------------------------------------------------------------------

#grab all the .txt files to loop over for analysis
files = list.files(wd, pattern = "mark.txt", full.name=T, recursive=T)

rm(list=ls()[!ls() %in% c("files")])   # clears the memory of everything except the file list, since R MARK is very memory hungry

for (f in 1:length(files)){

  require(RMark)

  # bring in the inp files and convert to RMark format
  ms_data = import.chdata(files[f], field.types=c("n","f"))

  #convert to factor
  ms_data$species = as.factor(ms_data$species)
  spname = ms_data$species[1]

  print(spname)
  cat("Imported data.", file="outfile.txt", sep="\n")

#---------------------------------------------------------------------------------
#          process multistrata data, includes capture at home, and dipsersal transitions
#---------------------------------------------------------------------------------
  # Build up the model.
  # begin.time == first period number
  ms_process = process.data(ms_data, model = "Multistrata", begin.time = 130, group = c("species"))

  #ddl = design data list
  ms_ddl = make.design.data(ms_process)

#---------------------------------------------------------------------------------
#          make dummy variables and covariates
#---------------------------------------------------------------------------------
  # Add dummy variables for operating on specific states(strata) or transitions
  # Species start in 1, movement between 1 and 2 (or vice versa) indicates a relatively long distance movement
  # Switching states indicates making a transition

  # SURVIVAL probability given that the individual is in A
  ms_ddl$S$inA = 0
  ms_ddl$S$inA[ms_ddl$S$stratum == "1"] = 1

  # SURVIVAL probability given that the individual is in B
  ms_ddl$S$inB = 0
  ms_ddl$S$inB[ms_ddl$S$stratum == "2"] = 1

  # RECAPTURE probability given that the individual is in A
  ms_ddl$p$strataA = 0
  ms_ddl$p$strataA[ms_ddl$p$stratum == "1"] = 1

  # RECAPTURE probability given that the individual is in B
  ms_ddl$p$strataB = 0
  ms_ddl$p$strataB[ms_ddl$p$stratum == "2"] = 1

  # TRANSITION probability given that the individual switches states
  # fix 1-->2 = 2-->1, since we are now interested in "switching"
  #TODO: Are the dummy variables below needed?
  # Transition probability given that the individual  A ---> B
  ms_ddl$Psi$toA = 0
  ms_ddl$Psi$toA[ms_ddl$Psi$stratum == "2" & ms_ddl$Psi$tostratum == "1"] = 1

  # Transition probability given that the individual  B ---> A
  ms_ddl$Psi$toB = 0
  ms_ddl$Psi$toB[ms_ddl$Psi$stratum == "1" & ms_ddl$Psi$tostratum == "2"] = 1

#--------------------------------------------------------------------------------
#           Build up the models
#---------------------------------------------------------------------------------
#          Define model structures for S (survival probability)
#---------------------------------------------------------------------------------
  Snull = list(formula = ~1)

#---------------------------------------------------------------------------------
#          Define model structures for p (capture probability)
#---------------------------------------------------------------------------------
# fix recapture probabilities for unsampled or omitted months
#    skipped_periods = c(237, 241, 267, 277, 278, 283, 284, 300, 311, 313, 314, 318, 321, 323, 337, 339, 344, 351): p = 0

  # select periods that were omitted from the study - untrapped
  p237 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 237,]))
  p241 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 241,]))
  p267 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 267,]))
  p277 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 277,]))
  p278 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 278,]))
  p283 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 283,]))
  p284 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 284,]))
  p300 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 300,]))
  p311 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 311,]))
  p313 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 313,]))
  p314 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 314,]))
  p318 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 318,]))
  p321 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 321,]))
  p323 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 323,]))
  p337 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 337,]))
  p339 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 339,]))
  p344 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 344,]))
  p351 = as.numeric(row.names(ms_ddl$p[ms_ddl$p$time == 351,]))

  # set those periods to p = 0, because they *can't* be anything else
  p237val = rep(0, length(p237))
  p241val = rep(0, length(p241))
  p267val = rep(0, length(p267))
  p277val = rep(0, length(p277))
  p278val = rep(0, length(p278))
  p283val = rep(0, length(p283))
  p284val = rep(0, length(p284))
  p300val = rep(0, length(p300))
  p311val = rep(0, length(p311))
  p313val = rep(0, length(p313))
  p314val = rep(0, length(p314))
  p318val = rep(0, length(p318))
  p321val = rep(0, length(p321))
  p323val = rep(0, length(p323))
  p337val = rep(0, length(p337))
  p339val = rep(0, length(p339))
  p344val = rep(0, length(p344))
  p351val = rep(0, length(p351))

  cat("Fixed omitted periods to zero.", sep="\n", file="outfile.txt", append=TRUE)

  # look for effects on recapture probability, given that some p are fixed to 0 (listed below)
  # link = "logit" is the default. "cloglog" may be esp. useful when there are fewer recaptures

  #Null Model
  pnull = list(formula = ~1, fixed = list(index = c(p237, p241, p267, p277, p278, p283, p284, p300, p311, p313, p314,
                                                     p318, p321, p323, p337, p339, p344, p351),
                                           value = c(p237val, p241val, p267val, p277val, p278val, p283val, p284val, p300val, p311val,
                                                     p313val, p314val, p318val, p321val, p323val, p337val, p339val,
                                                     p344val, p351val), link = "cloglog"))

  cat("Model for period effect on recapture probability.", sep="\n", file="outfile.txt", append=TRUE)

#---------------------------------------------------------------------------------
#          Define model structures for Psi (transition probability)
#---------------------------------------------------------------------------------
  Psinull = list(formula = ~1, link = "logit")   # KTS: Psinull is correct if we want psi 1 >> 2 == psi 2 >> 1

  cat("Defined model structure for Psi", sep="\n", file="outfile.txt", append=TRUE)

#---------------------------------------------------------------------------------
#          Run Models and collect results
#---------------------------------------------------------------------------------
  #send results to new folder - change working directory
  #wd = "~/portal-rodent-dispersal/mark_output/"
  setwd("C://Users//sarah//Documents//GitHub//portal-rodent-dispersal//mark_output//")

  cat("running the multistrata models", sep="\n", file="outfile.txt", append=TRUE)

  # SIMANNEAL should be best for multistrata models, but may take longer to run
  movemodel = mark(ms_process, ms_ddl, model.parameters = list(S=Snull,  p=pnull, Psi=Psinull),
                   options="SIMANNEAL", external=FALSE)

  cat("Null model is finished", sep="\n", file="outfile.txt", append=TRUE)

#-----------------------------------------------------
#             summarize results
#-----------------------------------------------------
  #ms_results <- collect.models(type = "Multistrata")
  #print(ms_results)
  cat("Summarized results.", sep="\n", file="outfile.txt", append=TRUE)

  print (movemodel$results$beta[1:3,])
  print (movemodel$results$real[1:3,])

#---------------------------------------------------------------------------------
#          Write result data to csv files
#---------------------------------------------------------------------------------
  write.csv(paste(movemodel$results$beta, paste(wd, "movemodel_beta_", spname, ".csv", sep=""))
  write.csv(movemodel$results$real, paste(wd, "movemodel_real_", spname, ".csv", sep=""))

  cat("End Code. Look for your csv files.", sep="\n", file="outfile.txt", append=TRUE)
  print( paste(spname, " is done.", sep = ""))

  rm(list=ls()[!ls() %in% c("f", "files")])   # clears the memory of everything except the file list and iterator
}

#END OF FUNCTION
}


###Install the relevant packages
setwd("/netscr/trhenry/Duke Project/RunningFolder")
#Load in Data
#Word data can be any number of words, but the final word data matrix has first column date (in integer form), second column domain (in integer form), then the word counts.
wordData <- read.csv("../wordDataTrimmed.csv")[,-1]

#Due to memory use issues, this matrix is specially constructed. It is provided for the whole year of linkage information.
linkData <- read.csv("processedLinkData.csv")[,-1]
processedLinkData <- linkData
#change them to matrices
linkData <- as.matrix(linkData)
wordData <- as.matrix(wordData)

#This loads in all functions and compiles all relevant c++ code
source("FullSampler2.R")
gc()

##These are your control vectors, pretty self explanatory. If you change the number of topics, you need to change the domainInterest Prior
control = list(nDomains = 467, nDays = 366)
nTopics <- 21
priorList = list(alpha =.1, beta=.1, chi = 1, PostRegressionMean =0, PostRegressionSD = 1000, DomainRateMean = 4, DomainRateSD = 1000, domainInterest = rep(1, 21), activationPrior = .2,
  netRegMeans = c(0,0,0,0,0), netRegSD = c(1000,1000,1000,1000,1000), nBlocksPrior = 25, nMembPrior = 1, blockIntConcentration = 50)
iterList = list(dayone.iter = 10, totalIter = 1000, netRegIter = 10, blockAssignIter = 10)
miscList = list(lag.window = 62, PostRegressionProp = .25, DomainRateProp = .5, netLag = 7, netProposals = c(1,.25,.25,.25,.25))

#Parameter order is: Word data, link data, nTopics, nBlocks, controlVector, IterationList, prior, MiscList, and the directory name for the runs.
topicNetworkSampler(wordData, linkData, 21, 35, control, iterList, priorList, miscList, "Run21")

#This function will produce the information as compressed Rdata files in special folders inside the run directory. You can load these files with load(name of the rData file)


#' Make Venn diagram of shared taxa (ASVs, OTUs) across sample groups
#'
#' Make Venn diagram of shared taxa (ASVs, OTUs) across sample groups from a phyloseq object. Overlap can be weighted by relative abundance
#' @param ps A phyloseq object
#' @param group The grouping factor. Should match variable in sample_data(ps)
#' @param fraction The fraction (0 to 1) of samples in a group in which the taxa should be present to be included in the count.
#' @param weight If TRUE, the overlaps are weighted by abundance
#' @param type "percent" or "counts"
#' @param relative Should abundances be made relative
#' @param plot If TRUE return a plot, if FALSE return a list with shared and unique taxa
#' @param ... Additional arguments
#' @keywords venn diagram
#' @return An venn plot
#' @import phyloseq
#' @import eulerr
#' @importFrom stats aggregate as.formula
#' @export

if (!requireNamespace("eulerr", quietly= TRUE)){install.packages("eulerr")}

ps_venn <- function(ps, group, fraction = 0, weight = FALSE, type = "percent", relative = TRUE, plot = TRUE, ...){

  if(relative){
    ps <- transform_sample_counts(ps, function(x) x/sum(x))
  }

  if(taxa_are_rows(ps)){
    ps_melted <- reshape2::melt(otu_table(ps))
  } else {
    ps_melted <- reshape2::melt(t(otu_table(ps)))
  }
  ps_melted <- merge(ps_melted, sample_data(ps), by.x = "Var2", by.y = "row.names")

  ps_agg <- aggregate(as.formula(paste("value ~ Var1 +",group)), data = ps_melted, function(x) (sum(x > 0)/length(x) >= fraction) * mean(x))

  ps_mat <- reshape2::dcast(as.formula(paste("Var1 ~ ",group)), data = ps_agg, value.var = "value")

  rownames(ps_mat) <- ps_mat[, 1]
  ps_mat <- ps_mat[, -1]
  ps_mat_bin <- (ps_mat>0)*1

  if(plot){
    if(weight){
      df <- eulerr::venn(ps_mat_bin, weights = rowMeans(ps_mat))
    } else {
      df <- eulerr::venn(ps_mat_bin)
    }
    plot(df, quantities = list(type=type), ...)
  } else {
    # Find taxa in all combinations
    combis <- lapply(2:ncol(ps_mat), function(k) lapply(lapply(1:(ncol(combn(1:ncol(ps_mat_bin), m = k))),
                                                               function(y) ps_mat_bin[, combn(1:ncol(ps_mat_bin), m = k)[, y]]),
                                                        function(x) rownames(x[rowSums(x) >= k, , drop=FALSE])))

    # Find taxa in singles
    singles <- apply(ps_mat_bin, 2, function(x) names(x[x > 0]))

    # Keep only those NOT in the same combination space
    singles <- lapply(seq_along(singles), function(x) setdiff(singles[[x]], do.call(c, singles[-x])))
    combis <- lapply(combis, function(cc) lapply(seq_along(cc), function(x) setdiff(cc[[x]], do.call(c, cc[-x]))))

    # Names
    names(singles) <- colnames(ps_mat_bin)
    for(i in 2:ncol(ps_mat)){
      names(combis[[i-1]]) <- apply(combn(colnames(ps_mat_bin), m = i), 2, function(x) paste(x, collapse = "__"))
    }

    # Recursively go through combination space from complex to simple to keep only those in unique combinations
    combis <- rev(combis)
    combis_new <- list()
    for(i in seq_along(combis)){
      if(i == 1) {
        combis_new[[i]] <- combis[[i]]
      } else {
        combis_new[[i]] <- lapply(combis[[i]], function(x) setdiff(x, unlist(combis_new)))
      }
    }
    combis_new <- c(singles, unlist(combis_new, recursive = FALSE))
    return(combis_new[sapply(combis_new, function(x) length(x)>0)])
  }
}