/*******************************************************************************
 * 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.kinemage;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;

import edu.duke.donaldLab.share.bond.Bond;
import edu.duke.donaldLab.share.bond.BondGraph;
import edu.duke.donaldLab.share.bond.BondGraphBuilder;
import edu.duke.donaldLab.share.bond.BreadthFirstBondIterator;
import edu.duke.donaldLab.share.geom.Vector3;
import edu.duke.donaldLab.share.mapping.AddressMapper;
import edu.duke.donaldLab.share.nmr.Assignment;
import edu.duke.donaldLab.share.nmr.DistanceRestraint;
import edu.duke.donaldLab.share.protein.Atom;
import edu.duke.donaldLab.share.protein.AtomAddressInternal;
import edu.duke.donaldLab.share.protein.Element;
import edu.duke.donaldLab.share.protein.Protein;
import edu.duke.donaldLab.share.protein.AtomAddressReadable;
import edu.duke.donaldLab.share.protein.Residue;
import edu.duke.donaldLab.share.protein.Subunit;

public class KinemageBuilder
{
	/**************************
	 *   Definitions
	 **************************/
	
	private static final Vector3 DefaultPosition = new Vector3();
	private static final int DefaultColor = 0;
	private static final int DefaultWidth = 1;
	private static final double DefaultLength = 1.0;
	
	
	/**************************
	 *   Static Methods
	 **************************/
	
	public static void appendLine( Kinemage kinemage, Vector3 source, Vector3 target )
	{
		appendLine( kinemage, source, target, "Line", DefaultColor, DefaultWidth );
	}
	
	public static void appendLine( Kinemage kinemage, Vector3 source, Vector3 target, String name, int color, int width )
	{
		Group group = new Group( name );
		group.addOption( "dominant" );
		kinemage.getRoot().addNode( group );
		List list = new List( "vector", "line" );
		list.addOption( "width= " + width );
		group.addNode( list );
		
		// add the points
		Point point = new Point( "source", new Vector3( source ) );
		list.addNode( point );
		point = new Point( "target", new Vector3( target ) );
		point.addOption( ColorScheme.DarkList[color] );
		list.addNode( point );
	}
	
	public static void appendOrientations( Kinemage kinemage, java.util.List<Vector3> orientations )
	{
		appendOrientations( kinemage, orientations, "Orientations", DefaultColor, DefaultWidth, DefaultLength );
	}
	
	public static void appendOrientations( Kinemage kinemage, java.util.List<Vector3> orientations, String name, int color, int width, double length  )
	{
		// add the group
		Group group = new Group( name );
		group.addOption( "dominant" );
		kinemage.getRoot().addNode( group );
		
		// add the source ball
		Vector3 origin = Vector3.getOrigin();
		List list = new List( "ball", "source" );
		group.addNode( list );
		Point point = new Point( "source", origin );
		point.addOption( ColorScheme.DarkList[color] );
		point.addOption( "r= " + Double.toString( length / 40.0 ) );
		list.addNode( point );
		
		// add each orientation
		for( Vector3 orientation : orientations )
		{
			// compute the target position
			Vector3 target = new Vector3( orientation );
			target.normalize();
			target.scale( length );
			
			// add the direction line
			list = new List( "vector", "line" );
			list.addOption( "width= " + width );
			group.addNode( list );
			point = new Point( "source", origin );
			list.addNode( point );
			point = new Point( "target", new Vector3( target ) );
			point.addOption( ColorScheme.DarkList[color] );
			list.addNode( point );
			
			// add the target ball
			list = new List( "dot", "target" );
			group.addNode( list );
			point = new Point( "target", target );
			point.addOption( ColorScheme.DarkList[color] );
			point.addOption( "width= 7" );
			list.addNode( point );
		}
	}
	
	public static void appendVector( Kinemage kinemage, Vector3 orientation )
	{
		appendVector( kinemage, orientation, DefaultPosition, "Vector", DefaultColor, DefaultWidth, DefaultLength );
	}
	
	public static void appendVector( Kinemage kinemage, Vector3 orientation, Vector3 position )
	{
		appendVector( kinemage, orientation, position, "Vector", DefaultColor, DefaultWidth, DefaultLength );
	}
	
	public static void appendVector( Kinemage kinemage, Vector3 orientation, Vector3 position, String name, int color, int width, double length )
	{
		// add the group
		Group group = new Group( name );
		group.addOption( "dominant" );
		kinemage.getRoot().addNode( group );
		
		// compute the target position
		Vector3 target = new Vector3( orientation );
		target.scale( length );
		target.add( position );
		
		// add the direction line
		List list = new List( "vector", "line" );
		list.addOption( "width= " + width );
		group.addNode( list );
		Point point = new Point( "source", new Vector3( position ) );
		list.addNode( point );
		point = new Point( "target", new Vector3( target ) );
		point.addOption( ColorScheme.DarkList[color] );
		list.addNode( point );
		
		// add the source marker
		list = new List( "dot", "source" );
		group.addNode( list );
		point = new Point( "source", position );
		point.addOption( ColorScheme.DarkList[color] );
		list.addOption( "width= 7" );
		list.addNode( point );
	}
	
	public static void appendAxes( Kinemage kinemage )
	{
		appendAxes( kinemage, DefaultWidth, DefaultLength );
	}
	
	public static void appendAxes( Kinemage kinemage, int width, double length )
	{
		Point point = null;
		Group group = new Group( "Axes" );
		group.addOption( "dominant" );
		kinemage.getRoot().addNode( group );
		List list = new List( "vector", "Axes" );
		list.addOption( "width= " + width );
		group.addNode( list );
		
		// add the x-axis
		point = new Point( "leftx", Vector3.getOrigin() );
		list.addNode( point );
		Vector3 x = Vector3.getUnitX();
		x.scale( length );
		point = new Point( "rightx", x );
		point.addOption( "red" );
		list.addNode( point );
		
		// add the y-axis
		point = new Point( "lefty", Vector3.getOrigin() );
		point.addOption( "P" );
		list.addNode( point );
		Vector3 y = Vector3.getUnitY();
		y.scale( length );
		point = new Point( "righty", y );
		point.addOption( "green" );
		list.addNode( point );
		
		// add the z-axis
		point = new Point( "leftz", Vector3.getOrigin() );
		point.addOption( "P" );
		list.addNode( point );
		Vector3 z = Vector3.getUnitZ();
		z.scale( length );
		point = new Point( "rightz", z );
		point.addOption( "blue" );
		list.addNode( point );
	}
	
	public static void appendProtein( Kinemage kinemage, Protein protein )
	throws IOException
	{
		appendProtein( kinemage, protein, "Protein" );
	}
	
	public static void appendProtein( Kinemage kinemage, Protein protein, String name )
	throws IOException
	{
		appendProtein( kinemage, protein, name, null );
	}
	
	public static void appendProtein( Kinemage kinemage, Protein protein, String name, String filterString )
	throws IOException
	{
		// load the bond graphs
		ArrayList<BondGraph> bondGraphs = BondGraphBuilder.getInstance().build( protein );
		
		// list the backbone atoms
		Group proteinGroup = new Group( name );
		kinemage.getRoot().addNode( proteinGroup );
		
		// parse the filter if needed
		ArrayList<AtomAddressInternal> filter = parseFilter( protein, filterString );
		
		// for each subunit
		for( Subunit subunit : protein.getSubunits() )
		{
			BondGraph bondGraph = bondGraphs.get( subunit.getId() );
			
			// get a group for the subunit
			Subgroup subunitGroup = new Subgroup( "Subunit " + subunit.getName() );
			proteinGroup.addNode( subunitGroup );
			
			// add the backbone chain
			List backboneChain = new List( "vector", "Backbone Chain" );
			backboneChain.addOption( "color= " + ColorScheme.DarkList[subunit.getId() % ColorScheme.DarkList.length] );
			backboneChain.addOption( "width= 3" );
			backboneChain.addOption( "master= {Backbone}" );
			subunitGroup.addNode( backboneChain );
			for( Residue residue : subunit.getResidues() )
			{
				// get the backbone atoms in order
				Atom atomN = residue.getAtomByName( "N" );
				Atom atomCa = residue.getAtomByName( "CA" );
				Atom atomC = residue.getAtomByName( "C" );
				
				// skip empty residues
				if( atomN == null || atomCa == null || atomC == null )
				{
					continue;
				}
				
				String nameN = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomN.getId() ) ).toString();
				String nameCa = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomCa.getId() ) ).toString();
				String nameC = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomC.getId() ) ).toString();
				
				backboneChain.addNode( new Point( nameN, atomN.getPosition() ) );
				backboneChain.addNode( new Point( nameCa, atomCa.getPosition() ) );
				backboneChain.addNode( new Point( nameC, atomC.getPosition() ) );
			}
			
			// prep heavy atoms list
			List heavyAtoms = new List( "vector", "Heavy Atoms" );
			heavyAtoms.addOption( "color= " + ColorScheme.DarkList[subunit.getId() % ColorScheme.DarkList.length] );
			heavyAtoms.addOption( "width= 3" );
			heavyAtoms.addOption( "off" );
			heavyAtoms.addOption( "master= {Heavy Atoms}" );
			subunitGroup.addNode( heavyAtoms );
			
			// prep hydrogen atoms list
			List hydrogenAtoms = new List( "vector", "Hydrogens" );
			hydrogenAtoms.addOption( "color= " + ColorScheme.LightList[subunit.getId() % ColorScheme.LightList.length] );
			hydrogenAtoms.addOption( "width= 3" );
			hydrogenAtoms.addOption( "off" );
			hydrogenAtoms.addOption( "master= {Hydrogens}" );
			subunitGroup.addNode( hydrogenAtoms );
			
			// find the first heavy atom in the protein
			AtomAddressInternal firstAddress = null;
			for( AtomAddressInternal address : subunit.getAtomIndex() )
			{
				if( !isHydrogen( protein, address ) )
				{
					firstAddress = address;
					break;
				}
			}
			
			// heavy atoms (do BFS in the bond graph)
			BreadthFirstBondIterator iterBond = new BreadthFirstBondIterator( bondGraph, firstAddress );
			while( iterBond.hasNext() )
			{
				ArrayList<Bond> bonds = iterBond.next();
				for( Bond bond : bonds )
				{
					// does the bond satisfy the residue filter?
					if( !doesBondSatisfyFilter( bond, filter ) )
					{
						continue;
					}
					
					// add hydrogens to a separate list
					if( isHydrogen( protein, bond.getLeftAddress() ) || isHydrogen( protein, bond.getRightAddress() ) )
					{
						addBond( protein, hydrogenAtoms, bond );
					}
					else
					{
						addBond( protein, heavyAtoms, bond );
					}
				}
			}
		}
	}
	
	public static void appendBackbone( Kinemage kinemage, Subunit subunit )
	{
		appendBackbone( kinemage, subunit, "Subunit Backbone", DefaultColor, DefaultWidth );
	}
	
	public static void appendBackbone( Kinemage kinemage, Protein protein )
	{
		appendBackbone( kinemage, protein, "Backbone", DefaultColor, DefaultWidth );
	}
	
	public static void appendBackbone( Kinemage kinemage, Subunit subunit, String name, int color, int width )
	{
		appendBackbone( kinemage, new Protein( new Subunit( subunit ) ), name, color, width );
	}
	
	public static void appendBackbone( Kinemage kinemage, Protein protein, String name, int color, int width )
	{
		// list the backbone atoms
		Group backboneGroup = new Group( name );
		backboneGroup.addOption( "dominant" );
		kinemage.getRoot().addNode( backboneGroup );
		
		// for each subunit
		for( Subunit subunit : protein.getSubunits() )
		{
			// add the backbone chain
			String subunitName = "Subunit " + subunit.getName();
			List subunitList = new List( "vector", subunitName );
			subunitList.addOption( "color= " + ColorScheme.DarkList[color] );
			subunitList.addOption( "width= " + width );
			subunitList.addOption( "master= {" + subunitName + "}" );
			backboneGroup.addNode( subunitList );
			for( Residue residue : subunit.getResidues() )
			{
				// get the backbone atoms in order
				Atom atomN = residue.getAtomByName( "N" );
				Atom atomCa = residue.getAtomByName( "CA" );
				Atom atomC = residue.getAtomByName( "C" );
				
				// skip empty residues
				if( atomN == null || atomCa == null || atomC == null )
				{
					continue;
				}
				
				String nameN = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomN.getId() ) ).toString();
				String nameCa = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomCa.getId() ) ).toString();
				String nameC = AddressMapper.mapAddress( protein, new AtomAddressInternal( subunit.getId(), residue.getId(), atomC.getId() ) ).toString();
				
				subunitList.addNode( new Point( nameN, atomN.getPosition() ) );
				subunitList.addNode( new Point( nameCa, atomCa.getPosition() ) );
				subunitList.addNode( new Point( nameC, atomC.getPosition() ) );
			}
		}
	}
	
	public static void appendDistanceRestraints( Kinemage kinemage, Protein protein, java.util.List<DistanceRestraint<AtomAddressInternal>> restraints )
	{
		// add the distance restraints as vectors
		Group group = new Group( "NOEs" );
		kinemage.getRoot().addNode( group );
		List list = new List( "vector", "NOEs" );
		list.addOption( "width= 3" );
		//list.addOption( "alpha= 0.6" );
		group.addNode( list );
		
		for( DistanceRestraint<AtomAddressInternal> restraint : restraints )
		{
			for( Assignment<AtomAddressInternal> assignment : restraint )
			{
				// is the restraint satisfied?
				Atom leftAtom = protein.getAtom( assignment.getLeft() );
				Atom rightAtom = protein.getAtom( assignment.getRight() );
				double atomDistance = leftAtom.getPosition().getDistance( rightAtom.getPosition() );
				boolean isSatisfied = atomDistance <= restraint.getMaxDistance();
				
				Point leftPoint = new Point(
					AddressMapper.mapAddress( protein, assignment.getLeft() ).toString(),
					protein.getAtom( assignment.getLeft() ).getPosition()
				);
				leftPoint.addOption( "P" );
				list.addNode( leftPoint );
				
				// get the right atom
				Point rightPoint = new Point(
					AddressMapper.mapAddress( protein, assignment.getRight() ).toString(),
					protein.getAtom( assignment.getRight() ).getPosition()
				);
				rightPoint.addOption( isSatisfied ? ColorScheme.SatisfiedNoe : ColorScheme.UnsatisfiedNoe );
				list.addNode( rightPoint );
			}
		}
	}
	
	public static void appendPoints( Kinemage kinemage, Collection<Vector3> points )
	{
		appendPoints( kinemage, points, "Points", DefaultColor, 7 );
	}
	
	public static void appendPoints( Kinemage kinemage, Collection<Vector3> points, String name, int color, int width )
	{
		// add the group
		Group group = new Group( name );
		group.addOption( "dominant" );
		kinemage.getRoot().addNode( group );
		
		for( Vector3 position : points )
		{
			// add the point
			List list = new List( "dot", "source" );
			group.addNode( list );
			Point point = new Point( "dot", position );
			point.addOption( ColorScheme.DarkList[color] );
			list.addOption( "width= " + Integer.toString( width ) );
			list.addNode( point );
		}
	}
	
	
	/**************************
	 *   Static Functions
	 **************************/
	
	private static boolean isHydrogen( Protein protein, AtomAddressInternal address )
	{
		return protein.getAtom( address ).getElement() == Element.Hydrogen;
	}
	
	private static void addBond( Protein protein, List list, Bond bond )
	{
		// get the left atom
		Point leftPoint = new Point(
			AddressMapper.mapAddress( protein, bond.getLeftAddress() ).toString(),
			protein.getAtom( bond.getLeftAddress() ).getPosition()
		);
		leftPoint.addOption( "P" );
		list.addNode( leftPoint );
		
		// get the right atom
		Point rightPoint = new Point(
			AddressMapper.mapAddress( protein, bond.getRightAddress() ).toString(),
			protein.getAtom( bond.getRightAddress() ).getPosition()
		);
		list.addNode( rightPoint );
	}
	
	private static ArrayList<AtomAddressInternal> parseFilter( Protein protein, String filterString )
	{
		// short cut
		if( filterString == null )
		{
			return null;
		}
		
		ArrayList<AtomAddressInternal> masks = new ArrayList<AtomAddressInternal>();
		
		String[] parts = filterString.split( "," );
		for( String part : parts )
		{
			String[] subparts = part.split( ":" );
			AtomAddressReadable address = new AtomAddressReadable(
				subparts[0].charAt( 0 ),
				Integer.parseInt( subparts[1] ),
				"CA"
			);
			masks.addAll( AddressMapper.mapAddress( protein, address ) );
		}
		
		return masks;
	}
	
	private static boolean doesBondSatisfyFilter( Bond bond, ArrayList<AtomAddressInternal> masks )
	{
		// no filter? consider it satisfied
		if( masks == null )
		{
			return true;
		}
		
		// see if the bond satisfies any of the masks
		for( AtomAddressInternal mask : masks )
		{
			if( doesAddressSatisfyMask( bond.getLeftAddress(), mask ) && doesAddressSatisfyMask( bond.getRightAddress(), mask ) )
			{
				return true;
			}
		}
		
		return false;
	}
	
	private static boolean doesAddressSatisfyMask( AtomAddressInternal address, AtomAddressInternal mask )
	{
		return address.getSubunitId() == mask.getSubunitId()
			&& address.getResidueId() == mask.getResidueId();
	}
}

