package osiris.cuda;

import org.apache.log4j.Logger;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;

/**
 *
 * Cubin is a way of compiling NVidia CUDA source files, with a suffix of .cu,
 * in to cubin files that can be loaded a CUDA executable code on to a CUDA
 * processor. This class will perform any compilation necessary on the fly. If
 * the CUBIN file exists and is newer than the CU file, then it is used with no
 * recompilation necessary.<P>
 * This class makes no sense with out the use of JCUDA
 * <P>
 * @author Adrian Challinor - adrian dot challinor at osiris dot co dot uk
 *
 */
public class Cubin {
	/**
	 * Logger for this class
	 */
	private static final Logger logger = Logger.getLogger(Cubin.class);
	private String cuFileName;
	private String cubinFileName;
	private String nvcc;
	private String arc;
	private String model;
	private boolean prepared;

	/**
	 * Class Constructor for Cubin.
	 *
	 * @param cuFileName
	 *            - String name of .cu file name.
	 */
	public Cubin(String cuFileName) {
		this.cuFileName = cuFileName;
		nvcc = "/usr/local/cuda/bin/nvcc";
		arc = "sm_11";
		model = System.getProperty("sun.arch.data.model");
		prepared = false;
	}

	/**
	 * The extension of the given file name is replaced with "cubin". If the
	 * file with the resulting name does not exist, it is compiled from the
	 * given file using NVCC. If the CU file is newer than the CUBIN file
	 * then it is also compiled.
	 *
	 * @throws IOException
	 *             If an I/O error occurs
	 */
	public void prepareCubinFile() throws IOException {
		Boolean recompile = false;
		prepared = true;

		// Check that the CU file exists - it has to !
		File cuFile = new File(cuFileName);
		if (!cuFile.exists()) {
			logger.error("File does not exist: " + cuFileName);
			throw new FileNotFoundException(cuFileName);
		}

		// Create a CUBIN file
		int endIndex = cuFileName.lastIndexOf('.');
		if (endIndex == -1) {
			endIndex = cuFileName.length() - 1;
		}
		cubinFileName = cuFileName.substring(0, endIndex + 1) + "cubin";

		// Check if the cubin file exists and is younger than the cu file.
		File cubinFile = new File(cubinFileName);
		if (cubinFile.exists()) {

			if (cuFile.lastModified() > cubinFile.lastModified()) {
				if (logger.isInfoEnabled()) {
					logger.info("Cubin file exists but is younger"); //$NON-NLS-1$
				}
				recompile = true;
			} else {
				if (logger.isInfoEnabled()) {
					logger
							.info("Using existing  - cubinFileName=" + cubinFileName); //$NON-NLS-1$
				}
			}
		} else
			recompile = true;

		if (recompile) {
			// Recompile the cu file to a cubin file.

			if (logger.isInfoEnabled()) {
				logger.info("Compiling  - cuFileName=" + cuFileName); //$NON-NLS-1$
			}

			String command = nvcc + " -m " + model + " -arch " + arc
					+ " -cubin " + cuFile.getPath() + " -o " + cubinFileName;

			if (logger.isDebugEnabled()) {
				logger.debug("Executing\n" + command); //$NON-NLS-1$
			}
			Process process = Runtime.getRuntime().exec(command);

			String errorMessage = new String(toByteArray(process
					.getErrorStream()));
			String outputMessage = new String(toByteArray(process
					.getInputStream()));
			int exitValue = 0;
			try {
				exitValue = process.waitFor();
			} catch (InterruptedException e) {
				Thread.currentThread().interrupt();
				throw new IOException(
						"Interrupted while waiting for nvcc output", e);
			}

			if (logger.isDebugEnabled()) {
				logger.debug("nvcc process exitValue " + exitValue); //$NON-NLS-1$
			}
			if (exitValue != 0) {
				if (logger.isDebugEnabled()) {
					logger.debug("errorMessage:\n" + errorMessage); //$NON-NLS-1$
				}
				if (logger.isDebugEnabled()) {
					logger.debug("outputMessage:\n" + outputMessage); //$NON-NLS-1$
				}
				throw new IOException("Could not create .cubin file: "
						+ errorMessage);
			}
		}
	}

	/**
	 * Returns the current architecture model. The system architecture is used
	 * in the compilation stage and defaults the the current machine architecture.
	 *
	 * @return - String, the architecture model name.
	 */
	public String getModel() {
		return model;
	}

	/**
	 * Set the architecture mode. If you need to set this to a different
	 * architecture, do so before preparing the CU file.
	 * @param model
	 */
	public void setModel(String model) {
		this.model = model;
	}

	/**
	 * Set the NVCC compiler location.
	 *
	 * @return
	 */
	public String getNvcc() {
		return nvcc;
	}

	/**
	 * Set the location of the NVCC compiler. This will default to the Linux
	 * default file location, /usr/local/cuda/bin/nvcc
	 *
	 * @return
	 */
	public void setNvcc(String nvcc) {
		this.nvcc = nvcc;
	}

	/**
	 * Fully reads the given InputStream and returns it as a byte array.
	 *
	 * @param inputStream
	 *            The input stream to read
	 * @return The byte array containing the data from the input stream
	 * @throws IOException
	 *             If an I/O error occurs
	 */
	private byte[] toByteArray(InputStream inputStream) throws IOException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		byte buffer[] = new byte[8192];
		while (true) {
			int read = inputStream.read(buffer);
			if (read == -1) {
				break;
			}
			baos.write(buffer, 0, read);
		}
		return baos.toByteArray();
	}

	/**
	 * Get the name of the CU input file, as specified in the
	 * @return
	 */
	public String getCuFileName() {
		return cuFileName;
	}

	/**
	 * Get the resulting CUBIN file. If the file has not yet been prepared,
	 * this will happen automatically.
	 *
	 * @return
	 * @throws IOException
	 */
	public String getCubinFileName() throws IOException {
		if (!prepared)
			prepareCubinFile();
		return cubinFileName;
	}

	/**
	 * returns the current CUDA processor architecture. The default is SM-11
	 * @return - String
	 */
	public String getArc() {
		return arc;
	}

	/**
	 * Sets the current CUDA processor architecture. The default is SM-11
	 * @param arc - String
	 */
	public void setArc(String arc) {
		this.arc = arc;
	}

}
