library(survival);


#####  Program for the pump failure data


#### Prior Pars.

#### Cost or repair vs cost of maintenance, for utlity func
CRCM.ratio <- 4.46;
down.time <- 1; ###hours of down time due to repair or maintenance

### Prior for theta
alpha1 <- 1/2;
beta1 <- 1/2;

### Prior for lambda

mX <- 7*24 ##hours, 10 days
sX <- 7*24 ##+- hours of operation (10 days), a priori

### a priori mean and squared coeficient of variantion for alpha
mu.alpha <- 3/2
C2.alpha <- 1/27

if (((sX/mX)^2 - (1/mu.alpha + C2.alpha)) < 0)
	cat("ERROR: WE REQUIRE", (sX/mX),"= (sX/mX)^2 > sqrt(1/mu.alpha + C2.alpha)=",
		sqrt(1/mu.alpha + C2.alpha), "\n" )
alpha2 <- 1 + ((sX/mX)^2 + 1)/((sX/mX)^2 - (1/mu.alpha + C2.alpha))
beta2  <- mX*(alpha2 - 1)/mu.alpha

#### Now jump to the end, to the MakeAllPlots() function.




rber <- function( n, p) { (runif(n) < p); }



ReadData <- function( sel=0, sort.data=FALSE, print.data=FALSE) {
	
	dat <- read.table("DatosBomba.txt", header=TRUE);
	
	if (sel[1] == 0) { ### automatically chooses all data
		ntt <<- dim(dat)[1];
		sel <- 1:ntt;
	}
	else {
		ntt <<- length(sel);
	}
	
	if (sort.data) {
		sort.per <- order(dat[,1]);
		dat <- dat[ sort.per, ];
	}
	
	tt <<- dat[sel,1];
	mt <<- sum(tt);
	
	eps <<- dat[sel,2]-1;
	##eps[rber( ntt, p=0.5)] <<- 1;

	meps <<- sum(eps);
	
	if (meps == ntt) {
		cat("WARNING: No failure times!!\n", "tt=", tt, "eps=", eps, "\n");
	}
	
	nomeps <<- ntt-meps
	
	log.mteps <<- sum(log(tt[which(eps == 0)]))
	
	if (print.data) {
		cat("i  t_i  eps_i\n");
		for (i in 1:ntt)
			cat( i, tt[i], eps[i], "\n");
		cat("\n");
	}
	
}


PlotThetaDens <- function( al1=alpha1, be1=beta1, xlim=c(0,1),
	xlab="Theta", ylab="Density", main="Posterior for theta (and prior in grey)") {
	
	curve( dbeta( x, meps+al1, ntt-meps+be1),
		xlim=xlim, xlab=xlab, ylab=ylab, main=main);
	
	curve( dbeta( x, al1, be1), xlim=c(0.0001,0.9999), col="grey", add=TRUE);

}

SimThetaPost <- function( m, al1=alpha1, be1=beta1) {
	rbeta( m, meps+al1, ntt-meps+be1);
}




PlotCondPriorW <- function( la, al) {
	curve( dgamma( x, shape=al-1, rate=la), xlim=c(0,3*((al-1)/la)),
		xlab="W", ylab="Density",
		main=paste("Prior for W given lambda=", la, "and alpha=", al) );
}

SimCondPriorW <- function( m=1, la, al) { rgamma( m, shape=al-1, rate=la); }

PlotCondPostLambda <- function( al, al2=alpha2, be2=beta2,
	xlim=c(0,3*sh/rt), xlab="Lambda", ylab="Density",
	main=paste("Cond. posterior of Lambda given alpha=", al), ...) {
		
	sh <- al2+ntt+(al-1)*nomeps;
	rt <- be2+mt;
	curve( dgamma( x, shape=sh, rate=rt),
		xlim=xlim, xlab=xlab, ylab=ylab, main=main, ...);
}

SimCondPostLambda <- function( m=1,  al, al2=alpha2, be2=beta2) {
	sh <- al2+ntt+(al-1)*nomeps;
	rt <- be2+mt;	
	rgamma( m, shape=sh, rate=rt);
}



###	Here Nu = (la,al) not including theta. U = -log post

NuPostU <- function( la, al, al2=alpha2, be2=beta2) {
	(1-(ntt+al2+(al-1)*nomeps))*log(la) + la*(mt+be2) + ## la | al
	-1*al*log.mteps + nomeps*lgamma(al); ## Rest, Uniform prior for al
}

VecNuPostU <- function(x) {
	NuPostU( x[1], x[2], al2=alpha2, be2=beta2);
}

SuppNuPost <- function( la, al) {
	if ((la > 0) && ((1 < al) && (al <= 2))) TRUE else FALSE;
}

VecSuppNuPost <- function(x) { SuppNuPost( x[1], x[2]); }


AlphaPostU <- function(al, al2=alpha2, be2=beta2) {
	(al*nomeps)*log(be2+mt) +
	-1*al*log.mteps +
	-1*lgamma( al2 + al*nomeps + meps) + nomeps*lgamma(al)
}

PlotAlphaPost <- function( al2=alpha2, be2=beta2, add=TRUE, ...) {
	
	C <- AlphaPostU( 1.5, al2=al2, be2=be2);
	dens <- function(al) {
		exp( -1*AlphaPostU( al, al2=al2, be2=be2) + C);
		}
	C <- C - log(integrate( dens, 1, 2)$value);
	
	curve( dens(x), add=add, ...);
}

SimPrioNu <- function( al2=alpha2, be2=beta2) {
	c( rgamma( 1, shape=al2, rate=be2), runif( 1, min=1, max=2));
}

### Simulate from the marginal of alpha with the gamma terms removed
### quite similar to the actual posterior marginal
SimChopOffAlphaPost <- function( m=1, al2=alpha2, be2=beta2) {

	logC <- log.mteps - nomeps*log(be2+mt) + log.mteps
	#= D + log.mteps ... we add this term to have a wider prop. and better mixing
	
	1 + log(runif(m)*(exp(logC)-1) + 1)/logC;
}


SimNuPost <- function( m, sel=0, sort.data=FALSE, al2=alpha2, be2=beta2,
	 al0=SimChopOffAlphaPost( 1, al2=al2, be2=be2), plot.time.series=TRUE) {

	ReadData(sel=sel, sort.data=sort.data);
	
	

	### Do a MCMC to simulate from alpha.  Use SimChopOffAlphaPost
	### as a propolsal
	
	Tr = 1000+10*m; ## ten times the required sample size + burn in

	al <- al0; ### Initial values
	
	U <- -1*lgamma( al2 + al*nomeps + meps) + nomeps*lgamma(al) + al*log.mteps
	
	output <- matrix( 0, nrow=Tr+1, ncol=3);
	
	output[ 1, 2] <- al;
	output[ 1, 3] <- AlphaPostU( al, al2, be2);
	
	for (j in 2:(Tr+1)) {
	### the only remaining part in the M-H are the gamma terms.
	### al is always in the correct support, 1 < al <= 2
	
		propal <- SimChopOffAlphaPost( 1, al2=al2, be2=be2); 
			
		propU <- -1*lgamma(al2 + nomeps*propal + meps) + nomeps*lgamma(propal) + propal*log.mteps
			
		A <- exp(U-propU); ### Acceptance ratio
			
		if (runif(1) < A) { 
			### Accepted
				
			al <- propal;
			U  <- propU;
		}
			
		output[ j, 2] <- al;
		output[ j, 3] <- AlphaPostU( al, al2, be2);
	} 
		
	if (plot.time.series)
		plot( 1:(Tr+1), -1*output[,3], type="l", main="MCMC output for Alpha",
		 		xlab="Iteration", ylab="LogPost");
	
	### We will output only m points sampling every 3 iterations 
	### output[ (1:m)*3, 2];
	### But first we need to sample from la | al_i, data
	
	output[ 1000+(1:m)*10, 1] <-
		SimCondPostLambda( m=m, al=output[ (1:m)*3, 2], al2=alpha2, be2=beta2); 
	
	output[ 1000+(1:m)*10, ];

}


SampDist <- function( x, smpl) { length(which(smpl <= x))/length(smpl); }

SimAll <- function(m=10000, sel=0, sort.data=FALSE,
	layout.mat=matrix( 1:8, nrow=4, ncol=2, byrow=TRUE),
	al1=alpha1, be1=beta1, al2=alpha2, be2=beta2, n=101) {
	
	layout(layout.mat);

	### col 1 has the lambdas and col 2 the alphas col 3 the U's
	output <- SimNuPost( m=m, sel=sel, sort.data=sort.data, al2=al2, be2=be2)
	
	### The predective, given Nu: Ga( al_i, la_i) 
	pred <- rgamma( n=m, shape=output[,2], rate=output[,1]);
	
	### The independent marginal for theta
	simth <- SimThetaPost( m, al1, be1);
	
	### The cond dist of W given Nu: Ga( al_i-1, la_i)
	simW <- rgamma( n=m, shape=output[,2]-1, rate=output[,1]);


	### Plot the posterior (and prior) for theta
	PlotThetaDens( al1, be1);
	

	### Plot the posterior and prior for Lambda
	hist( output[,1], freq=FALSE, breaks=20,
		xlab="Lambda", main="Post. for Lambda  (and prior in grey)");
	curve( dgamma( x, shape=al2, rate=be2), add=TRUE, col="grey");
	

	### Plot the posterior (apporx) and the exact posterior for Alpha
	hist( output[,2], freq=FALSE, breaks=20,
		xlab="Alpha",
		main="Exact (black line) and Monte Carlo post. for Alpha\nPrior in grey");
	PlotAlphaPost( al2=al2, be2=be2);
	### and the prior
	curve( dunif( x, min=1, max=2), col="grey", add=TRUE);


	#### Calculate the cdf os the pred. post for X
	mX <- mean(pred);
	sdX <- sqrt(var(pred));
	len <- length(pred);
	
	cat("Mean for pred X:", mX, "and sd. dev:", sdX, "\n");
	
	pred.dens <- density( pred, from=0);
	XS <- pred.dens$x;
	Dist <- sapply( XS, function(x) { length(which(pred <= x))/len; } )
	
	#### Cut off for Dist < 0.99 to avoid noise
	mx <- max(which(Dist < 0.99))
	XS <- XS[1:mx];
	Dist <- Dist[1:mx]

	#### and print the critical time (tst or tstar) at Eps
	PEps <- (al1+meps)/(al1+be1+ntt); ## P(eps=1 | Data)
	tmp <- abs(Dist-(1-PEps));
	tst <-which(tmp == min(tmp))[1]; ### tstar = XS[tst]
	
	### Calculate the empirical cdf ditribution
	### (never mind the matainence times)
	EmpDist <- sapply( XS, function(x) { length(which(tt <= x))/ntt; } )



	#### Plot the post. pred for X and W	
	plot(density(simW), col="red", xlab="Units of X",
		main="Pred. post. for X (black) and W (red)");
	lines(pred.dens);
	
	### Add the observations
	jitt <- runif( ntt, 0, 0.005);
	points( tt, jitt);
	
	### Mark the maintenance times
	mn <- which(eps == 1);
	points( tt[mn], jitt[mn], pch=20, col="red");
	
	### Mark tstar
	lines( rep( XS[tst], 2), c(0,0.005), col="blue");
	


	### Plot the Kaplam-Meier estimator
	pump.Surv <- Surv( tt, 1-eps);
	pump.survfit <- survfit(pump.Surv);
	plot(pump.survfit, col="magenta", xlim=c( XS[1], XS[mx]),
		xlab="Units of X", ylab="Probability",
main="Survival func. for the pred. post. of X (black),\n Empirical (green) and kaplan-meier (magenta)");
	
	#### add the cdf and 1-cdf of the post pred for X
	lines( XS, 1-Dist);
	### and the emp. cdf
	lines( XS, 1-EmpDist, col="green");
	


	#### Plot the hazard function
	plot( XS, pred.dens$y[1:mx]/(1-Dist), type="l",
		xlab="Units of X", ylab="Risk", main="Hazard Function");	
	### Finally plot the h function
	Plotsimh(list( X=XS, densX=pred.dens$y[1:mx], DistX=Dist),
		main="(t + d) x hazard function");
	
	cat("P(eps=0)=", 1-PEps, "tc=", XS[tst]);
	

	list( predEps=PEps, tc=XS[tst],
		th=simth, la=output[,1], al=output[,2],
		predX=pred, densX=pred.dens$y[1:mx], DistX=Dist, X=XS, mx=mx, predW=simW);
} 




####Plot the h function, the hazard function times t of a Ga( alpha, 1)
Ploth <- function(alpha=seq( 1, 2, by=0.2),
	xlim=c(0,1), ylim=c(0,1.5), xlab="t", ylab="h1(t)") {
	
	h <- function(x, shape) { x*dgamma( x, shape=shape)/(1-pgamma(x,shape=shape)); }
	
	curve( h(x, shape=alpha[1]), xlim=xlim, ylim=ylim, xlab=xlab, ylab=ylab);
	for (al in alpha)
		curve( h(x, shape=al), xlim=xlim, add=TRUE);
}


####Plot the h function, the hazard function times (t+d) of the simulated values
#### use as input the output of SimAll.  ratio is the CR/CM ratio to obtain the optimal t*
Plotsimh <- function( sim, ratio=CRCM.ratio, d=down.time,
	ylim=c(0,1), xlab="hours", ylab="ratio", ...) {

	h <- (sim$X + d)*sim$densX/(1-sim$DistX);
	
	### Assuming h is increasing the proper cutoff and t*
	sel <- which(((ylim[1] < h) & (h < ylim[2])));

	plot( sim$X[sel], h[sel], type="l", axes=FALSE,
		ylim=ylim, xlab=xlab, ylab=ylab, ...);	

	### Add the observations
	jitt <- runif( ntt, 0.0, 0.05);
	
	### Mark the failure times
	mn <- which(eps == 0);
	points( tt[mn], jitt[mn], pch=20);
	### Mark the maintenance times
	mn <- which(eps == 1);
	points( tt[mn], jitt[mn]);


		
	axis( 1, at=axTicks(1)); ## Tick marks for the x axis
	axis( 4, at=axTicks(2)); ## Tick marks for the y axis
	
	rt <- seq( from=1/ylim[2] + 1, to=10, by=0.5); ### CR/CM ratios to be plotted
	axis( 2, at=1/(rt-1), labels=paste(rt)) 

	if (0 < ratio) {
		istar <- max(which(h < 1/(ratio-1)))[1];
		tstar <-  sim$X[istar]; 
		
		cat("t* =", tstar, "\n");

		### line for the intercept
		lines( c( sim$X[sel[1]], sim$X[sel[max(sel)]]), c(1/(ratio-1),1/(ratio-1)), col="grey");

		#### Tick marks for the ratio and t*
		axis( 2, at=1/(ratio-1), labels="", col="red", tcl=-0.5);
		axis( 2, at=1/(ratio-1), labels="", col="red", tcl=1);
		axis( 1, at=tstar, labels="", col="red", tcl=-0.5);
		axis( 1, at=tstar, labels="", col="red", tcl=1);
		
	}	
	
	tstar;
}

### Produce all plots for the paper using global parameters
MakeAllPlots <- function(m=10000, pointsize=16) {
	
	### The Correcting Factor plot
	cat("Creating:", "../text/CorrFactor.png\n"); 
	png("../text/CorrFactor.png", pointsize=pointsize)
	curve( 1/(x-1), xlim=c(1.2,5), xlab="ratio", ylab="");
	axis( 1, at=1.2);
	dev.off();

	cat("Redoing all the sampling ...\n");
	sim <<- SimAll(m);

	### The h function plot	
	cat("\nCreating:", "../text/hfun.png\n"); 
	png("../text/hfun.png", pointsize=pointsize);
	Plotsimh(sim);
	dev.off();

	### Plot the posterior (and prior) for theta
	cat("Creating:", "../text/PostTheta.png\n");
	png("../text/PostTheta.png", pointsize=pointsize); 
	PlotThetaDens(main="", xlab="Probability", xlim=c(0,0.5));
	dev.off();
	
	### Plot the posterior and prior for Lambda
	cat("Creating:", "../text/PostLambda.png\n");
	png("../text/PostLambda.png", pointsize=pointsize); 
	hist( sim$la, freq=FALSE, breaks=20,
		xlab="1/hour", main="");
	curve( dgamma( x, shape=alpha2, rate=beta2), add=TRUE, col="grey");
	dev.off();

	### Plot the posterior (apporx) and the exact posterior for Alpha
	cat("Creating:", "../text/PostAlpha.png\n");
	png("../text/PostAlpha.png", pointsize=pointsize); 
	hist( sim$al, freq=FALSE, breaks=20, xlab="", main="");
	PlotAlphaPost();
	### and the prior
	curve( dunif( x, min=1, max=2), col="grey", add=TRUE);
	dev.off();

	#### Plot the post. pred for X	
	cat("Creating:", "../text/Pred.png\n");
	png("../text/PostPred.png", pointsize=pointsize); 
	hist( sim$predX, xlab="hours", freq=FALSE, main="", breaks=20);
	###lines( sim$X, sim$densX);
	axis( 1, at=sim$tc, labels="", col="blue");
	dev.off();

	### plot the post. pred. survival function	
	### along with the Kaplam-Meier estimator
	cat("Creating:", "../text/KMandPostPredSurv.png\n");
	png("../text/KMandPostPredSurv.png", pointsize=pointsize); 
	pump.Surv <- Surv( tt, 1-eps);
	pump.survfit <- survfit(pump.Surv);
	plot(pump.survfit, col="grey", xlim=c( sim$X[1], sim$X[sim$mx]),
		xlab="hours", ylab="Probability", main="");
	#### add the cdf and 1-cdf of the post pred for X
	lines( sim$X, 1-sim$Dist)
	dev.off()


}