runSim <- function(nTrials = 25, nDogs = 30,
                   e0 = 100, s0 = 10,
                   ks = 10, ke = 1){
	simDogs = matrix(nrow=nDogs, ncol=nTrials)
	for(i in 1:nDogs){
		e = e0; s = s0
		for(j in 1:nTrials){
			tot = e + s
			res = sample(0:1, 1, prob=c(s/tot, e/tot))
			if(res == 0){ s = s + ks}
			if(res == 1){ e = max(e - ke, 0)}
			simDogs[i, j] = res
		}
	}
	return(simDogs)
}


calcStats <- function(dat){
	stats  = matrix(nrow = nrow(dat), ncol = 5)
	colnames(stats) = c("TrialsBeforeFirstAvoidance",
	                    "TrialsBeforeSecondAvoidance",
	                    "TotalShocks",
	                    "TrialofLastShock",
	                    "Alternations")
	for(i in 1:nrow(dat)){
		x = dat[i, ]
		stats[i, ] = c(which(x == 0)[1],
		               which(x == 0)[2],
		               sum(x),
		               tail(which(x == 1), 1),
		               sum(diff(x) != 0))
	}
	summary = t(apply(stats, 2, function(x) c(mean(x), sd(x))))
	colnames(summary) = c("Mean", "SD")
	return(list(stats = stats, summary = summary))
}


simDogs = runSim(nTrials = 25, nDogs = 30,
                  e0 = 100, s0 = 10,
                  ks = 25, ke = 15)

#realStats = calcStats(realDogs)
simStats  = calcStats(simDogs)

#print(realStats$summary, 3)
print(simStats$summary,  3)