/*******************************************************************************
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
 * USA
 * 
 * Contact Info:
 * 	Bruce Donald
 * 	Duke University
 * 	Department of Computer Science
 * 	Levine Science Research Center (LSRC)
 * 	Durham
 * 	NC 27708-0129 
 * 	USA
 * 	brd@cs.duke.edu
 * 
 * Copyright (C) 2011 Jeffrey W. Martin and Bruce R. Donald
 * 
 * <signature of Bruce Donald>, April 2011
 * Bruce Donald, Professor of Computer Science
 ******************************************************************************/


package edu.duke.donaldLab.share.mpi;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintStream;
import java.util.HashMap;

import org.apache.log4j.Logger;

import edu.duke.donaldLab.share.io.Logging;
import edu.duke.donaldLab.share.mpi.messages.ClientException;

import mpi.MPI;
import mpi.MPIException;
import mpi.Request;
import mpi.Status;

public class MPIEndpoint
{
	/**************************
	 *   Definitions
	 **************************/
	
	private static final int RankServer = 0;
	private static final boolean DefaultAppend = true;
	private static final boolean ClientsWriteToLog = false;
	// UNDONE: support communication modes
	//private static final CommunicationMode DefaultCommunicationMode = CommunicationMode.Standard;
	
	
	/**************************
	 *   Data Members
	 **************************/
	
	private static final Logger m_log;
	private static int m_rank;
	private static int m_numNodes;
	private static String m_nodeName;
	private static boolean m_isServer;
	private static HashMap<Request,Object[]> m_buffers;
	private static File m_logFile;
	

	/**************************
	 *   Accessors
	 **************************/
	
	public static int getRank( )
	{
		return m_rank;
	}
	
	public static int getNumNodes( )
	{
		return m_numNodes;
	}
	
	public static String getNodeName( )
	{
		return m_nodeName;
	}
	
	public static boolean isServer( )
	{
		return m_isServer;
	}
	
	public static int getNumClients( )
	{
		return m_numNodes - 1;
	}
	
	
	/**************************
	 *   Constructors
	 **************************/
	
	static
	{
		m_log = Logging.getLog( MPIEndpoint.class );
		m_rank = -1;
		m_numNodes = 0;
		m_nodeName = null;
		m_isServer = false;
		m_buffers = new HashMap<Request,Object[]>();
		m_logFile = null;
	}
	
	
	/**************************
	 *   Methods
	 **************************/
	
	public static void initClientServer( String[] args, File logFile, ClientServerCreator creator )
	{
		try
		{
			MPIEndpoint.init( args, logFile );
			
			if( MPIEndpoint.isServer() )
			{
				// UNDONE: init log4j for server
				creator.createServer().run();
			}
			else
			{
				try
				{
					// UNDONE: init log4j for client
					creator.createClient().run();
				}
				catch( Exception ex )
				{
					MPIEndpoint.sendMessageToServer( new ClientException( ex ) );
				}
				// yes, catch errors too so the server can be notified
				catch( Error er )
				{
					MPIEndpoint.sendMessageToServer( new ClientException( er ) );
				}
			}
			
			MPIEndpoint.cleanup();
		}
		catch( Exception ex )
		{
			MPIEndpoint.reportException( ex );
			MPIEndpoint.cleanup();
		}
		
		/* Jeff: 10/25/2009 - NOTE:
			don't put cleanup() in a finally block here.
			For some unknown reason, it changes the way exceptions are thrown/caught. =(
		*/
	}
	
	public static void init( String[] args, File logFile )
	throws MPIException
	{
		// save parameters
		m_logFile = logFile;
		
		// get info from MPI
		MPI.Init( args );
		m_rank = MPI.COMM_WORLD.Rank();
		m_numNodes = MPI.COMM_WORLD.Size();
		m_nodeName = getNodeName( m_rank );
		m_isServer = m_rank == RankServer;
		
		// init the log file if needed
		if( m_isServer )
		{
			log( "Log initialized!", false );
		}
	}
	
	public static void cleanup( )
	{
		try
		{
			MPI.Finalize();
		}
		catch( MPIException ex )
		{
			MPIEndpoint.reportException( ex );
		}
	}
	
	public static void log( String msg )
	{
		logAs( msg, DefaultAppend, m_rank );
	}
	
	public static void log( String msg, boolean append )
	{
		logAs( msg, append, m_rank );
	}
	
	public static void logAs( String msg, int rank )
	{
		logAs( msg, DefaultAppend, rank );
	}
	
	public static void logAs( String msg, boolean append, int rank )
	{
		// prepend the node name
		msg = getNodeName( rank ) + ": " + msg + "\n";
		
		if( isServer() || ClientsWriteToLog )
		{
			try
			{
				/* Jeff: 12/30/2008 - NOTE:
					This isn't exactly the most efficient way of keeping a log file.
					However, having real-time updates to this log file is very important,
					so we explicitly open and close the file with every write to ensure
					that all the underlying channels (possibly even network channels)
					flush their buffers.
				*/
				FileWriter writer = new FileWriter( m_logFile, append );
				writer.write( msg );
				writer.close();
				writer = null;
			}
			catch( IOException ex )
			{
				System.out.println( "Error writing to alert file!" );
				System.out.println( "Alert:\n\t" + msg );
				ex.printStackTrace( System.out );
			}
		}
		else
		{
			/* NOTE:
				Clients can just write to stdout.
				MPI handles sending these messages, but they often
				lag far behind what's really happening.
			*/
			m_log.info( msg );
		}
	}
	
	public static void sendMessage( Message message, int destRank )
	throws MPIException
	{
		message.setSourceRank( m_rank );
		
		// DEBUG
		//alert( "Sending " + message.getType() + " to " + destRank + "..." );
		
		MPI.COMM_WORLD.Send( new Object[] { message }, 0, 1, MPI.OBJECT, destRank, 0 );
		
		// DEBUG
		//alert( "Sent " + message.getType() + " to " + destRank + "!" );
	}
	
	public static Request sendMessageAsync( Message message, int destRank )
	throws MPIException
	{
		message.setSourceRank( m_rank );
		
		// DEBUG
		//alert( "Sending " + message.getType() + " to " + destRank + "..." );
		
		return MPI.COMM_WORLD.Isend( new Object[] { message }, 0, 1, MPI.OBJECT, destRank, 0 );
		
		// DEBUG
		//alert( "Sent " + message.getType() + " to " + destRank + "!" );
	}
	
	public static void sendMessageToServer( Message message )
	throws MPIException
	{
		sendMessage( message, RankServer );
	}
	
	public static void sendMessageToClients( Message message )
	throws MPIException
	{
		// UNDONE: can use broadcast?
		//MPI.COMM_WORLD.Bcast( );
		
		for( int i=1; i<getNumNodes(); i++ )
		{
			sendMessage( message, i );
		}
	}
	
	public static Message waitForMessage( )
	throws MPIException, InterruptedException
	{
		/* Jeff: 10/25/2009 - NOTE:
			Recv appears to be a busy wait. =(
			I'll just write a sleeping wait using the async methods
		*/
		Request request = registerForMessage();
		Message message = getRegisteredMessage( request );
		while( message == null )
		{
			Thread.sleep( 100 );
			message = getRegisteredMessage( request );
		}
		
		return message;
	}
	
	public static Message waitForMessageBusy( )
	throws MPIException
	{
		Object[] buffer = new Object[1];
		MPI.COMM_WORLD.Recv( buffer, 0, 1, MPI.OBJECT, MPI.ANY_SOURCE, MPI.ANY_TAG );
		return (Message)buffer[0];
	}
	
	public static Request registerForMessage( )
	throws MPIException
	{
		Object[] buffer = new Object[1];
		Request request = MPI.COMM_WORLD.Irecv( buffer, 0, 1, MPI.OBJECT, MPI.ANY_SOURCE, MPI.ANY_TAG );
		
		// save the buffer so we can find it later
		m_buffers.put( request, buffer );
		
		return request;
	}
	
	public static Message getRegisteredMessage( Request request )
	throws MPIException
	{
		Message message = null;
		
		// check to see if there's a message waiting
		Status status = request.Test();
		if( status != null )
		{
			// read the message from the buffer
			message = (Message)m_buffers.get( request )[0];
			
			// cleanup the buffer
			m_buffers.remove( request );
		}
		
		return message;
	}
	
	/* is this ever used?
	public static void waitForSignalFromClients( Message message )
	throws MPIException
	{
		log( "Waiting for " + numSlaves + " slaves..." );
		
		// wait for all slaves to send the signal
		int clientsResponded
		while( slavesLoaded < numSlaves )
		{
			// wait for the loaded message
			Message msg = waitForNextMessage();
			assert( msg.getType() == message.getType() );
			slavesLoaded++;
		}
		
		log( "All clients complete." );
	}
	*/
	
	public static void reportException( Exception ex )
	{
		reportExceptionAs( ex, m_rank );
	}
	
	public static void reportExceptionAs( Exception ex, int rank )
	{
		// write the exception info to a buffer
		ByteArrayOutputStream buf = new ByteArrayOutputStream();
		ex.printStackTrace( new PrintStream( buf ) );
		
		// convert the buffer to a string and alert
		logAs( "Exception!\n" + buf.toString(), rank );
	}
	
	
	/**************************
	 *   Functions
	 **************************/
	
	private static String getNodeName( int rank )
	{
		if( rank == RankServer )
		{
			return String.format( "%s-%s", getProcessorName(), "server" );
		}
		else
		{
			return String.format( "%s-%s-%02d", getProcessorName(), "client", rank );
		}
	}
	
	private static String getProcessorName( )
	{
		String name = "Unknown";
		
		try
		{
			name = MPI.Get_processor_name();
		}
		catch( MPIException ex )
		{
			// just print it for posterity
			ex.printStackTrace( System.err );
		}
		
		return name;
	}
}
