path: root/src/main/java/org/openslx/graph/Graph.java

package org.openslx.graph;

import java.awt.BasicStroke;
import java.awt.Color;
import java.awt.Font;
import java.awt.Graphics2D;
import java.awt.RenderingHints;
import java.awt.Stroke;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.math3.util.FastMath;

import com.google.gson.annotations.Expose;

/**
 * The Graph stores the Nodes and Edges, and InferenceHeurisics to allow
 * the structure of the graph to be modified.
 */
public class Graph implements java.io.Serializable
{

	private static final long serialVersionUID = 5488288903546706793L;
	private String _label;
	private String _caption = "servers are grey, clients are blue, sugar is sweet, and so are you";
	private Map<Node, Node> _nodes = new HashMap<>();
	private Map<Edge, Edge> _edges = new HashMap<>();
	
	@Expose
	private Collection<Node> nodes = _nodes.values();
	@Expose
	private Collection<Edge> edges = _edges.values();

	private double minX = Double.POSITIVE_INFINITY;
	private double maxX = Double.NEGATIVE_INFINITY;
	private double minY = Double.POSITIVE_INFINITY;
	private double maxY = Double.NEGATIVE_INFINITY;
	private double maxWeight = 0;

	private int _frameCount = 0;

	private static final Color COLOR_BG = new Color( 250, 250, 255 );
	private static final Color COLOR_TITLE = new Color( 200, 200, 255 );
	private static final Color COLOR_LABEL = new Color( 0, 0, 20 );
	private static final Color COLOR_LABEL_SHADOW = new Color( 220, 220, 220 );
	private static final Color COLOR_TITLE2 = new Color( 190, 190, 210 );
	private static final Color COLOR_EDGE = new Color( 100, 100, 255 );

	private static final int IMG_X = 1280;
	private static final int IMG_Y = 1024;

	private final Graphics2D graphicsContext;
	private final PngRenderer renderer;

	public Graph( String label )
	{
		_label = label;
		BufferedImage targetImage = new BufferedImage( IMG_X, IMG_Y, BufferedImage.TYPE_INT_ARGB );
		graphicsContext = targetImage.createGraphics();
		graphicsContext.setRenderingHint( RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON );
		renderer = new PngRenderer( targetImage );
	}

	// Add a Node to the Graph.
	public void addNode( Node node )
	{
		// Only add the Node to the HashMap if it's not already in there.
		if ( _nodes.containsKey( node ) )
			return;
		_nodes.put( node, node );
	}

	// Add an Edge to the Graph. Increment the weighting if it already exists.
	public void setEdge( Node source, Node target, double weight )
	{
		// Do not add self-edges or weights that are not positive.
		if ( source.equals( target ) || weight <= 0 ) {
			return;
		}

		addNode( source );
		addNode( target );

		// Add the Edge to the HashMap, or find the existing entry.
		Edge edge = _edges.get( new Edge( source, target ) );
		if ( edge == null ) {
			source = _nodes.get( source );
			target = _nodes.get( target );
			edge = new Edge( source, target );
			_edges.put( edge, edge );
		}
		// Set the edge weight.
		edge.setWeight( weight );
	}

	// Remove a Node from the Graph, along with all of its emanating Edges.
	public boolean removeNode( Node node )
	{
		if ( !_nodes.containsKey( node ) )
			return false;
		// Remove the Node from the HashMap.
		_nodes.remove( node );

		// Remove all Edges that connect to the removed Node.
		Iterator<Edge> edgeIt = _edges.keySet().iterator();
		while ( edgeIt.hasNext() ) {
			Edge edge = edgeIt.next();
			if ( edge.getSource().equals( node ) || edge.getTarget().equals( node ) ) {
				edgeIt.remove();
			}
		}
		return true;
	}

	// Return true if the Graph contains the Node.
	// (This does not necessarily imply that the Node is visible).
	public boolean contains( Node node )
	{
		return _nodes.containsKey( node );
	}

	// Return true if the Graph contains the Edge.
	public boolean contains( Edge edge )
	{
		return _edges.containsKey( edge );
	}

	// Return the Graph's Node that has the same nick as the supplied Node.
	public Node get( Node node )
	{
		return _nodes.get( node );
	}

	// Return the Graph's Edge that matched the supplied Edge.
	public Edge get( Edge edge )
	{
		return _edges.get( edge );
	}

	public String toString()
	{
		return "Graph: " + _nodes.size() + " nodes and " + _edges.size() + " edges.";
	}

	public String toString2()
	{
		return "Nodes:\n" + _nodes + "\nEdges:\n" + _edges;
	}

	// Apply the temporal decay to the Graph.
	public void decay()
	{
		for ( Iterator<Edge> it = _edges.values().iterator(); it.hasNext(); ) {
			Edge edge = it.next();
			edge.resetWeight();
			if ( !edge.isAlive() ) {
				it.remove();
			}
		}

		List<Node> toRemove = null;
		for ( Node node : _nodes.values() ) {
			node.age();
			if ( !node.isAlive() ) {
				if ( toRemove == null ) {
					toRemove = new ArrayList<>();
				}
				toRemove.add( node );
			}
		}
		if ( toRemove != null ) {
			for ( Node node : toRemove ) {
				removeNode( node );
			}
		}
	}

	// Returns the set of all Nodes that have emanating Edges.
	// This therefore returns all Nodes that will be visible in the drawing.
	public Set<Node> getConnectedNodes()
	{
		Set<Node> connectedNodes = new HashSet<>();
		for ( Edge edge : _edges.values() ) {
			connectedNodes.add( edge.getSource() );
			connectedNodes.add( edge.getTarget() );
		}
		return connectedNodes;
	}

	public Collection<Edge> getEdges()
	{
		return _edges.values();
	}

	// Limit movement values to stop nodes flying into oblivion.
	private static final double MAX_MOVE = 0.75;
	private static final double K = 2;
	private static final double MOVEMENT_SLOWDOWN = 0.001;

	// Applies the spring embedder.
	private void doLayout( Set<Node> visibleNodes, int iterations )
	{

		// For performance, copy each set into an array.
		Node[] nodes = visibleNodes.toArray( new Node[ visibleNodes.size() ] );
		Edge[] edges = _edges.keySet().toArray( new Edge[ _edges.size() ] );

		if ( nodes.length == 0 )
			return;

		// For each iteration...
		for ( int it = 0; it < iterations; it++ ) {

			// Calculate forces acting on nodes due to node-node repulsions...
			for ( int a = 0; a < nodes.length; a++ ) {
				final Node nodeA = nodes[a];
				for ( int b = a + 1; b < nodes.length; b++ ) {
					final Node nodeB = nodes[b];

					double deltaX = nodeB.getX() - nodeA.getX();
					double deltaY = nodeB.getY() - nodeA.getY();
					final double maxDistance = ( nodeA.getDesiredDistance() * nodeB.getDesiredDistance() ) * 0.9;

					double distanceSquared = deltaX * deltaX + deltaY * deltaY;

					if ( distanceSquared > maxDistance * maxDistance )
						continue;

					if ( distanceSquared < 0.01 ) {
						deltaX = Math.random() / 10 + 0.1;
						deltaY = Math.random() / 10 + 0.1;
						distanceSquared = deltaX * deltaX + deltaY * deltaY;
					}

					final double distance = FastMath.sqrt( distanceSquared );
					final double repulsiveForce = ( K * maxDistance / distance ) * 2;

					final double sumDesired = nodeA.getDesiredDistance() + nodeB.getDesiredDistance();
					// A gets pushed away stronger if B has higher desired distance than A
					final double rfA = repulsiveForce * ( nodeB.getDesiredDistance() / sumDesired );
					// Vice versa for B
					final double rfB = repulsiveForce * ( nodeA.getDesiredDistance() / sumDesired );

					nodeB.addForceX( ( rfB * deltaX / distance ) );
					nodeB.addForceY( ( rfB * deltaY / distance ) );
					nodeA.addForceX( - ( rfA * deltaX / distance ) );
					nodeA.addForceY( - ( rfA * deltaY / distance ) );
				}

			}

			// Calculate forces acting on nodes due to edge attractions.
			for ( int e = 0; e < edges.length; e++ ) {
				final Edge edge = edges[e];
				final Node nodeA = edge.getSource();
				final Node nodeB = edge.getTarget();

				final double minDistance = ( nodeA.getDesiredDistance() * nodeB.getDesiredDistance() ) * 1.1;

				final double deltaX = nodeB.getX() - nodeA.getX();
				final double deltaY = nodeB.getY() - nodeA.getY();

				double distanceSquared = deltaX * deltaX + deltaY * deltaY;

				if ( distanceSquared < minDistance * minDistance )
					continue;

				double distance = FastMath.sqrt( distanceSquared );

				if ( distance > minDistance * 2 ) {
					distance = minDistance * 2;
					distanceSquared = distance * distance;
				}

				final double attractiveForce = ( distanceSquared / K ) * 2;
				final double sumDesired = nodeA.getDesiredDistance() + nodeB.getDesiredDistance();
				// A gets pulled towards B stronger if B has higher desired distance than A
				final double afA = attractiveForce * ( nodeB.getDesiredDistance() / sumDesired );
				// Vice versa for B
				final double afB = attractiveForce * ( nodeA.getDesiredDistance() / sumDesired );

				nodeB.addForceX( -afB * deltaX / distance );
				nodeB.addForceY( -afB * deltaY / distance );
				nodeA.addForceX( afA * deltaX / distance );
				nodeA.addForceY( afA * deltaY / distance );

			}

			// Now move each node to its new location...
			for ( int a = 0; a < nodes.length; a++ ) {
				Node node = nodes[a];

				// General weak attraction towards (0|0)
				if ( node.getDesiredDistance() > 5 ) {
					node.addForceX( -node.getX() / 500 );
					node.addForceY( -node.getY() / 500 );
				}

				double xMovement = MOVEMENT_SLOWDOWN * node.getFX();
				double yMovement = MOVEMENT_SLOWDOWN * node.getFY();

				if ( xMovement > MAX_MOVE ) {
					xMovement = MAX_MOVE;
				}
				else if ( xMovement < -MAX_MOVE ) {
					xMovement = -MAX_MOVE;
				}
				if ( yMovement > MAX_MOVE ) {
					yMovement = MAX_MOVE;
				}
				else if ( yMovement < -MAX_MOVE ) {
					yMovement = -MAX_MOVE;
				}

				node.setX( node.getX() + xMovement );
				node.setY( node.getY() + yMovement );

				// Reset the forces
				node.resetForce();
			}

			// Swap two random nodes if it results in less stress on them
			Node a = nodes[(int) ( Math.random() * nodes.length )];
			Node b = nodes[(int) ( Math.random() * nodes.length )];
			if ( !a.equals( b ) ) {
				double initialLen = getEdgeLenSum( edges, a, b );
				a.swapPos( b );
				double swappedLen = getEdgeLenSum( edges, a, b );
				if ( swappedLen + 0.5 > initialLen ) {
					a.swapPos( b ); // Isn't any better, swap back
				}
			}

		}

	}

	private double getEdgeLenSum( Edge[] edges, Node a, Node b )
	{
		double result = 0;
		for ( Edge edge : edges ) {
			if ( edge.hasNode( a ) || edge.hasNode( b ) ) {
				result += edge.getLengthSquared();
			}
		}
		return result;
	}

	// Work out the drawing boundaries...
	public void calcBounds( Set<Node> nodes )
	{

		minX = Double.POSITIVE_INFINITY;
		maxX = Double.NEGATIVE_INFINITY;
		minY = Double.POSITIVE_INFINITY;
		maxY = Double.NEGATIVE_INFINITY;
		maxWeight = 0;

		for ( Node node : nodes ) {

			if ( node.getX() > maxX ) {
				maxX = node.getX();
			}
			if ( node.getX() < minX ) {
				minX = node.getX();
			}
			if ( node.getY() > maxY ) {
				maxY = node.getY();
			}
			if ( node.getY() < minY ) {
				minY = node.getY();
			}
		}

		// Increase size if too small.
		final double minSize = 10;
		if ( maxX - minX < minSize ) {
			double midX = ( maxX + minX ) / 2;
			minX = midX - ( minSize / 2 );
			maxX = midX + ( minSize / 2 );
		}
		if ( maxY - minY < minSize ) {
			double midY = ( maxY + minY ) / 2;
			minY = midY - ( minSize / 2 );
			maxY = midY + ( minSize / 2 );
		}

		// Work out the maximum weight.
		for ( Edge edge : _edges.values() ) {
			if ( edge.getWeight() > maxWeight ) {
				maxWeight = edge.getWeight();
			}
		}

		// Jibble the boundaries to maintain the aspect ratio.
		double xyRatio = ( ( maxX - minX ) / ( maxY - minY ) ) / ( IMG_X / IMG_Y );
		if ( xyRatio > 1 ) {
			// diagram is wider than it is high.
			double dy = maxY - minY;
			dy = dy * xyRatio - dy;
			minY = minY - dy / 2;
			maxY = maxY + dy / 2;
		}
		else if ( xyRatio < 1 ) {
			// Diagram is higher than it is wide.
			double dx = maxX - minX;
			dx = dx / xyRatio - dx;
			minX = minX - dx / 2;
			maxX = maxX + dx / 2;
		}

	}

	private static final Font FONT_64 = new Font( "SansSerif", Font.BOLD, 64 );
	private static final Font FONT_18 = new Font( "SansSerif", Font.BOLD, 18 );
	private static final Font FONT_12 = new Font( "SansSerif", Font.PLAIN, 12 );
	private static final Font FONT_10 = new Font( "SansSerif", Font.PLAIN, 10 );
	private static final Stroke STROKE_2 = new BasicStroke( 2.0f );

	public void drawImage( Set<Node> nodes, double edgeThreshold )
	{

		// Now actually draw the thing...

		graphicsContext.setColor( COLOR_BG );
		graphicsContext.fillRect( 0, 0, IMG_X, IMG_Y );

		graphicsContext.setColor( COLOR_TITLE );
		graphicsContext.setFont( FONT_64 );
		graphicsContext.drawString( _label, 20, 80 );

		graphicsContext.setColor( COLOR_TITLE2 );
		graphicsContext.setFont( FONT_18 );
		graphicsContext.drawString( _caption, 0, IMG_Y - 5 - 50 );
		graphicsContext.setFont( FONT_12 );
		graphicsContext.drawString( "This frame was drawn at " + new Date(), 0, IMG_Y - 5 );

		// Draw all edges...
		for ( Edge edge : _edges.values() ) {
			if ( edge.getWeight() < edgeThreshold ) {
				continue;
			}

			double weight = edge.getWeight();

			Node nodeA = edge.getSource();
			Node nodeB = edge.getTarget();
			int x1 = toImageX( nodeA );
			int y1 = toImageY( nodeA );
			int x2 = toImageX( nodeB );
			int y2 = toImageY( nodeB );
			graphicsContext.setStroke( new BasicStroke( (float) ( FastMath.log( weight + 1 ) * 0.1 ) + 1 ) );
			final int alpha = 102 + (int) ( 153 * weight / maxWeight );
			graphicsContext.setColor( new Color( ( edge.getColor().getRGB() & 0xffffff ) | ( alpha << 24 ), true ) );
			graphicsContext.drawLine( x1, y1, x2, y2 );
		}

		// Draw all nodes...
		graphicsContext.setStroke( STROKE_2 );
		graphicsContext.setFont( FONT_10 );
		for ( Node node : nodes ) {
			final int x = toImageX( node );
			final int y = toImageY( node );
			final int nodeRadius = node.getRadius();
			graphicsContext.setColor( node.getColor() );
			graphicsContext.fillOval( x - nodeRadius, y - nodeRadius, nodeRadius * 2, nodeRadius * 2 );
			graphicsContext.setColor( COLOR_EDGE );
			graphicsContext.drawOval( x - nodeRadius, y - nodeRadius, nodeRadius * 2, nodeRadius * 2 );
			graphicsContext.setColor( COLOR_LABEL_SHADOW );
			graphicsContext.drawString( node.toString(), x - ( node.toString().length() * 3 ) + 1, y - nodeRadius - 2 + 1 );
			graphicsContext.setColor( COLOR_LABEL );
			graphicsContext.drawString( node.toString(), x - ( node.toString().length() * 3 ), y - nodeRadius - 2 );
		}
	}

	private int toImageX( Node node )
	{
		return (int) ( ( IMG_X - 50 ) * ( node.getX() - minX ) / ( maxX - minX ) ) + 25;
	}

	private int toImageY( Node node )
	{
		return (int) ( ( IMG_Y - 55 ) * ( node.getY() - minY ) / ( maxY - minY ) ) + 35;
	}

	public int getFrameCount()
	{
		return _frameCount;
	}

	public String getLabel()
	{
		return _label;
	}

	public void setCaption( String caption )
	{
		_caption = caption;
	}

	public byte[] makeNextImage()
	{
		_frameCount++;

		Set<Node> nodes = getConnectedNodes();

		doLayout( nodes, 200 );
		calcBounds( nodes );

		ByteArrayOutputStream baos = new ByteArrayOutputStream( 100000 );

		try {
			drawImage( nodes, 0.1 );
			//ImageIO.write( targetImage, "png", baos );
			renderer.render( baos );

			writeGraph();
			return baos.toByteArray();
		} catch ( Exception e ) {
			System.out.println( "PieSpy has gone wibbly: " + e );
			e.printStackTrace();
		}
		return null;
	}

	// Serialize this Graph and write it to a File.
	public void writeGraph()
	{
		/*
		try {
			String strippedChannel = _label.toLowerCase().substring( 1 );
			File dir = new File( config.outputDirectory, strippedChannel );
			File file = new File( dir, strippedChannel + "-restore.dat" );
			ObjectOutputStream oos = new ObjectOutputStream( new FileOutputStream( file ) );
			oos.writeObject( SocialNetworkBot.VERSION );
			oos.writeObject( this );
			oos.flush();
			oos.close();
		} catch ( Exception e ) {
			// Do nothing?
		}
		*/
	}

	public Node getNode( String id )
	{
		return _nodes.get( new Node( id ) );
	}

}