# This library provides couple of useful functions for image processing to test image processing in R
# 
# Written by Nir Ben-Dvora  -  nir@bendvora.com
#
library(dtt);
library(plotrix);

image.array <- function (a, asp=TRUE, rem_mar=TRUE) {
# Modify the default function of image to accept an array
	# TBD: check type and dimentions
	# strech image to [0-1] 
	if (min(a)<0 | max(a)>1) {
		st <- (a-min(a))/(max(a)-min(a));
	}  else {
		st <- a;
	}
	oldmar = par("mar"); oldpty=par("pty");
	# Set the proportions of the image, usuall this is desired. asp=FALSE will change that
	if (asp) par(pty="s");
	if (rem_mar) par(mar = c(0,0,3,0)); #rep(0, 4));
	plot.new();
	x <- par("usr")[1L:2];
	y <- par("usr")[3:4];
	rasterImage(st, min(x), min(y), max(x), max(y), interpolate = FALSE);
	par(mar=oldmar, pty=oldpty);
}

plotImg <- function(img, asp=TRUE, rem_mar=TRUE) {
	# use the modified image array to plot gray scale one matrix image
	if (class(img)=="matrix") {
		# create an array with the same matrix
		newimg <- array(0,dim=c(dim(img),3))
		newimg[,,1]<-newimg[,,2]<-newimg[,,3]<-img;
	} else {
		# assuming input is array
		newimg = img;
	}
	image(newimg,asp,rem_mar);
}

# Image equalization
eqImg <- function(img) {
	if (class(img)=="matrix") {
		eqimg <- ecdf(img)(img);
		dim(eqimg) <- dim(img);
	} else {
		eqimg <- img
		eqimg[,,1] <- ecdf(img[,,1])(img[,,1]);
		eqimg[,,2] <- ecdf(img[,,2])(img[,,2]);
		eqimg[,,3] <- ecdf(img[,,3])(img[,,3]);
	}

	return(eqimg);
}

plotImgHist <- function(img) {
	if (class(img)=="matrix") {
		multhist(list(img),beside=TRUE,space=c(0,0),
			main="Image Histogram",xlab="Level",ylab="Count",
			col=c("#C0C0C0"));

	} else {
		multhist(list(img[,,1],img[,,2],img[,,3]),beside=TRUE,space=c(0,0),
			main="R,G,B Histogram",xlab="Level",ylab="Count",
			col=c("#FF9999","#99FF99","#9999FF"));
	}
}

imageProc <- function(mat, block, FUN) {
	# TBD: Add the ability to add other parameters (such as mask, boundaries, ...)
	# This function go block by block in mat and apply function FUN on the block
	# FUN can retrun the same or different block size
	# If block = 1 than FUN is applied pixel by pixel, could be used for convolution or median for example
	# The result is a new combined matrix with the same order of the blocks FUN is applied to
	# FUN must get the following parameters FUN(mat, row, column, block)
	# TBD: Non square blocks are not handled, so at this point the image should be a round multiplications of the blocks
	# In order truncate the image do: mat<-mat[1:((dim(mat)%/%block)*block)[1],1:((dim(mat)%/%block)*block)[2]]

	# Fist apply FUN on each block
	m2=mapply(function(r,c) FUN(mat,r,c,block), row(mat)[seq(1,nrow(mat),block),seq(1,ncol(mat),block)], col(mat)[seq(1,nrow(mat),block),seq(1,ncol(mat),block)],SIMPLIFY=FALSE)
	
	# m2 is a set of values after applying FUN on every block
	# Now combine back the block in the order of applying it, start with row bind and then column bind
	# TBD: If return values is not number, it could be optimized by only setting the dimention - no need for rbind and cbind
	rr=length(seq(1,nrow(mat),block));
	m3=mapply(function(a) {do.call(rbind, m2[seq(a,a+rr-1)])}, seq(1,length(m2),rr), SIMPLIFY=FALSE)
	m4=do.call(cbind, m3)

	return(m4)
}

# List of functions that could be applied on an image
# ---------------------------------------------------
# Return the same untouched block - used for testing
sameBlock <- function(mat,r,c,block) {return(mat[seq(r,r+block-1),seq(c,c+block-1)])}
# Return the sum of the block - this could also be achived with all one mask matrix
sumBlock <- function(mat,r,c,block) {return(sum(mat[seq(r,r+block-1),seq(c,c+block-1)]))}
# Apply Discrete Cosine Transform to a block
dctBlock <- function(mat,r,c,block) {return(mvdct(mat[seq(r,r+block-1),seq(c,c+block-1)])/8)}
# Apply Inverse Discrete Cosine Transform to a block
idctBlock <- function(mat,r,c,block) {return(8*mvdct(mat[seq(r,r+block-1),seq(c,c+block-1)],inverted=TRUE))}
# Apply Discrete Cosine Transform to a block with quantization matrix imageProcQM
dctBlockq <- function(mat,r,c,block) {return(round((mvdct(mat[seq(r,r+block-1),seq(c,c+block-1)])/8) / imageProcQM)*imageProcQM)}
imageProcQM.default=matrix(c(15,11,11,12,15,19,25,32,11,13,10,10,12,15,19,24,
            11,10,14,14,16,18,22,27,12,10,14,18,21,24,28,33,
            15,12,16,21,26,31,36,42,19,15,18,24,31,38,45,53,
            25,19,22,28,36,45,55,65,32,24,27,33,42,53,65,77), ncol=8)
imageProcQM <- imageProcQM.default

# Different Mask Matrices
# ------------------------
blur=matrix(c(1,  2,  1, 2,  4,  2, 1,  2,  1),ncol=3)/16;
blur2=matrix(rep(1/25, 25),ncol=5);
blur3=matrix(rep(1/49, 49),ncol=7);
emboss=matrix(c( 2,  0,  0,  0, -1,  0  ,0,  0, -1),ncol=3);
emboss2=matrix(c(1,  1, -1, 1,  3, -1, 1, -1, -1),ncol=3)/3;
emboss3=matrix(c(-2, -1, 0, -1, 1, 1, 0, 1, 2), ncol=3);
sharpen1=matrix(c(-1, -1, -1, -1, 9, -1,-1, -1, -1),ncol=3);
sharpen2=matrix(c(0, -2,  0 ,-2, 11, -2, 0, -2,  0),ncol=3)/3;
sharpen3=matrix(c(0, -1,  0 ,-1, 5, -1, 0, -1,  0),ncol=3);
laplacian=matrix(c(1, 1, 1, 1, -8, 1, 1, 1, 1),ncol=3);
edge1=matrix(c(1,  1,  1, 1, -7,  1, 1,  1,  1),ncol=3);
edge2=matrix(c(-5,  0,  0, 0,  0,  0, 0,  0,  5),ncol=3);
edge3=matrix(c(0, 1, 0, 1, -4, 1, 0, 1, 0),ncol=3);
edge_enhance=matrix(c(-1,  0,  0, 0, 1,  0, 0,  0,  0),ncol=3);
crispen=matrix(c(0,  -1,  0, -1,  5,  -1, 0,  -1,  0),ncol=3);
crispen2=matrix(c(1,  -2,  1, -2,  5,  -2, 1,  -2,  1),ncol=3);

# Use blur as the default mask, this could be changed with other masks
imageProcMask = blur

# Apply convolution matrix on a pixel, used with pixel:blcok=1
convProc <- function(mat,r,c,block) {
	nbr<-(dim(imageProcMask)[1]-1)/2;
	if (any(c(r-nbr,dim(mat)[1]-r-nbr+1,c-nbr,dim(mat)[2]-c-nbr+1)<=0)) {
	# if out of mat boundary - return the original pixel
		return(mat[r,c]);
	} else {
	# if in the boundaries, apply mask
		return(sum(mat[seq(r-nbr,r+nbr),seq(c-nbr,c+nbr)]*imageProcMask));
	}
}
# Apply Median on a pixel neighbourhood, used with pixel:blcok=1
medianProc <- function(mat,r,c,block) {
	# Currently median of one pixel sorrounding
	nbr<-1;
	if (any(c(r-nbr,dim(mat)[1]-r-nbr+1,c-nbr,dim(mat)[2]-c-nbr+1)<=0)) {
	# if out of mat boundary - return the original pixel
		return(mat[r,c]);
	} else {
	# if in the boundaries, apply median
		return(median(mat[seq(r-nbr,r+nbr),seq(c-nbr,c+nbr)]));
	}
}
# Count number of trailing zeros in a block
tZerosBlock <- function(mat,r,c,block) {
	z=0;
	ind=c(block,block);
	a = mat[seq(r,r+block-1),seq(c,c+block-1)];
	while(ind[1]!=0) {
		if(a[ind[1],ind[2]]==0) {
			z=z+1;
			ind=nextDiagWalk(ind,block,block,-1);
		} else {
			ind=c(0,0);
		}
	}
	return(z);
}
# Generate graph from image. Used with block=1
# The graph nodes are listed as from, two, capacity in graphProcG
graphProc <- function(mat,r,c,block) {
	hist_var <- 10;  #0.3;
	if (r==1 & c==1) {   # initial state
		graphProcG <<- as.undirected(graph.empty()) + vertices(seq(1,(dim(mat)[1])*(dim(mat)[2])));
	}	
	node <- (c-1)*dim(mat)[1]+r;
	if (c>1) { # add left node
		nn <- (c-2)*dim(mat)[1]+(r);
		graphProcG <<- add.edges(graphProcG, c(node, nn));
		graphProcG$capacity <<- c(graphProcG$capacity, exp(-((mat[r,c]-mat[r,c-1])^2)/(2*hist_var^2)));
	}
	if (r>1) { # add top node
		nn <- (c-1)*dim(mat)[1]+(r-1);
		graphProcG <<- add.edges(graphProcG, c(node, nn));
		graphProcG$capacity <<- c(graphProcG$capacity, exp(-((mat[r,c]-mat[r-1,c])^2)/(2*hist_var^2)));
	}
	#if (r==1) print(c);
	return(node);
}


# This function calculates the bi modal middle point based on Otsu algorithm
# The function is translated from wikipedia
# input:  histogram of values from 1-256
# output: middle point between classes
otsu <- function(histogram) {
	total <- sum(histogram); # number of pixels
	sum <- sum(seq(1,256)*histogram);
	sumB <- 0;
	wB <- 0;
	wF <- 0;
	max <- 0;
	between <- 0;
	threshold1 <- 0;
	threshold2 <- 0;
	for (i in 1:256) {
		wB <- wB + histogram[i];
		if (wB == 0) {next};
		wF <- total - wB;
		if (wF == 0) {break};
		sumB <- sumB + (i * histogram[i]);
		mB <- sumB / wB;
		mF <- (sum - sumB) / wF;
		between <- wB * wF * (mB - mF)^2;
		if ( between >= max ) {
			threshold1 <- i;
			if ( between > max ) {
				threshold2 <- i;
            		}
            		max <- between;            
        	}
    	}
	return (( threshold1 + threshold2 ) / 2);
}


#Huffman encoding (no decoding yet):
#------------------------------------
# This function calculates the binary hufmman code for a given probability vector
# Input is a probability vector
# Return is the huffman code for the given vector
# This function wasn't tested for performance and scale
#
# Written by Nir Ben-Dvora
#
bhuffman <- function (pvecin) {
	# Remove all zeros from vec - no code is needed for them
	pvec=pvecin[which(pvecin>0)];

	# Make couple of checks
	if (class(pvec) != "numeric") stop("Input must be numeric");
	if (length(pvec) < 2) stop("Numbers of positive elements must be 2 or more");

	# Recursion stop condition if the number of rows exactly 2
	if (length(pvec) == 2) return(c("0","1"));

	# Compress low probabilities values and call again. Use last min
	min1loc <- length(pvec)-which.min(rev(pvec))+1;
	min1 <- pvec[min1loc];
	pvec <- pvec[-min1loc];
	# min2loc <- which.min(pvec);
	min2loc <- length(pvec)-which.min(rev(pvec))+1;
	min2 <- pvec[min2loc];
	pvec <- pvec[-min2loc];
	pvec <- append(pvec,min1+min2);

	code <- bhuffman(pvec);

	# Get the last entry code, remove it and add split it back to previous entries
	min1code <- paste0(tail(code,n=1),"0");
	min2code <- paste0(tail(code,n=1),"1");
	# Remove last entry and add this code to the min1, min2 locations
	code <- code[-length(code)];
	# Add back the new codes
	code <- append(code,min2code,min2loc-1);
	code <- append(code,min1code,min1loc-1);

	return(code);

}

nextDiagWalk <- function(ind,maxi=8,maxj=8,dir=1) {
	# walk a matrix in a diagonal fashion
	# dir=1 goes upward, if dir=-1 go backward
	i=ind[1];j=ind[2];

	curdir = dir*(-1+(2*((i+j)%%2)));
	ni=i+curdir; nj=j-curdir;

	if (dir==1) {
		if (i==maxi & j==maxj) return(c(0,0));
		if (ni>maxi) {ni=maxi; nj=nj+2;}
		if (nj>maxj) {nj=maxj; ni=ni+2;}
		if (ni<1) ni=1;
		if (nj<1) nj=1;
	} else {
		if (i==1 & j==1) return(c(0,0));
		if (ni<1) {ni=1; nj=nj-2;}
		if (nj<1) {nj=1; ni=ni-2;}
		if (ni>maxi) ni=maxi;
		if (nj>maxj) nj=maxj;
	}		

	return(c(ni,nj));
}