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

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

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.ChemicalShift;
import edu.duke.donaldLab.share.nmr.DistanceRestraint;
import edu.duke.donaldLab.share.protein.AminoAcid;
import edu.duke.donaldLab.share.protein.Atom;
import edu.duke.donaldLab.share.protein.AtomAddressInternal;
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 PseudoatomBuilder
{
	/**************************
	 *   Definitions
	 **************************/
	
	private static final String PseudoatomsPath = "/resources/pseudo.atoms";
	private static final String PseudoatomNameChars = "mpq";
	
	
	/**************************
	 *   Data Members
	 **************************/
	
	private static PseudoatomBuilder m_instance;
	private Pseudoatoms m_pseudoatoms;
	
	
	/**************************
	 *   Constructors
	 **************************/
	
	private PseudoatomBuilder( )
	throws IOException
	{
		// read in the defs
		InputStream in = getClass().getResourceAsStream( PseudoatomsPath );
		m_pseudoatoms = PseudoatomReader.read( in );
	}
	
	
	/**************************
	 *   Static Methods
	 **************************/
	
	public static PseudoatomBuilder getInstance( )
	throws IOException
	{
		if( m_instance == null )
		{
			m_instance = new PseudoatomBuilder();
		}
		
		return m_instance;
	}
	
	public static boolean distanceRestraintsHavePseudoatoms( ArrayList<DistanceRestraint<AtomAddressReadable>> restraints )
	{
		// NOTE: this function is somewhat heuristic.
		// It just checks for atoms names whose first character is in a list
		for( DistanceRestraint<AtomAddressReadable> restraint : restraints )
		{
			for( Assignment<AtomAddressReadable> assignment : restraint )
			{
				Character leftChar = Character.toLowerCase( assignment.getLeft().getAtomName().charAt( 0 ) );
				Character rightChar = Character.toLowerCase( assignment.getRight().getAtomName().charAt( 0 ) );
				
				boolean hasPseudoatoms =
					PseudoatomNameChars.indexOf( leftChar ) >= 0
					|| PseudoatomNameChars.indexOf( rightChar ) >= 0;
				if( hasPseudoatoms == true )
				{
					return true;
				}
			}
		}
		
		return false;
	}
	
	public static Pseudoatoms getPseudoatoms( )
	{
		try
		{
			return getInstance().m_pseudoatoms;
		}
		catch( IOException ex )
		{
			System.err.println( "WARNING: could not load Pseudoatoms!" );
			ex.printStackTrace( System.err );
		}
		
		return null;
	}
	
	
	/**************************
	 *   Methods
	 **************************/
	
	public int build( Protein protein )
	{
		int numBuilt = 0;
		
		// for each subunit...
		for( Subunit subunit : protein.getSubunits() )
		{
			numBuilt += build( subunit );
		}
		
		return numBuilt;
	}
	
	public int build( Subunit subunit )
	{
		int numBuilt = 0;
		
		// for each residue
		for( Residue residue : subunit.getResidues() )
		{
			AminoAcid aminoAcid = residue.getAminoAcid();
			
			// get the pseudoatoms if any
			Set<String> pseudoatomNames = m_pseudoatoms.getPseudoatomNames( aminoAcid );
			if( pseudoatomNames == null )
			{
				continue;
			}
			
			// for each pseudoatom...
			for( String pseudoatomName : pseudoatomNames )
			{
				ArrayList<String> atomNames = m_pseudoatoms.getAtoms( aminoAcid, pseudoatomName );
				addPseudoatom( pseudoatomName, atomNames, residue );
				
				numBuilt++;
			}
		}
		
		subunit.updateAtomIndices();
		
		return numBuilt;
	}
	
	public int buildDistanceRestraints( Protein protein, List<DistanceRestraint<AtomAddressReadable>> restraints )
	{
		int numBuilt = 0;
		
		for( DistanceRestraint<AtomAddressReadable> restraint : restraints )
		{
			double maxLeftCorrection = 0.0;
			double maxRightCorrection = 0.0;
			
			for( Assignment<AtomAddressReadable> assignment : restraint )
			{
				double leftCorrection = interpretAddress( protein, assignment.getLeft() );
				double rightCorrection = interpretAddress( protein, assignment.getRight() );
				
				// update the max distances
				maxLeftCorrection = Math.max( maxLeftCorrection, leftCorrection );
				maxRightCorrection = Math.max( maxRightCorrection, rightCorrection );
			}
			
			// apply the largest corrections
			restraint.setMaxDistance( restraint.getMaxDistance() + maxLeftCorrection + maxRightCorrection );
			
			if( maxLeftCorrection > 0.0 || maxRightCorrection > 0.0 )
			{
				numBuilt++;
			}
		}
		
		return numBuilt;
	}
	
	public int buildShifts( Protein protein, List<ChemicalShift> shifts )
	{
		int numBuilt = 0;
		for( ChemicalShift shift : shifts )
		{
			for( AtomAddressReadable address : shift.getAddresses() )
			{
				/* NOTE:
					Sometimes they don't add in the wildcards for pseudoatoms
					so, we'll have to test and see if this is a pseudoatom or not
				*/
				
				// if this address doesn't exist in the protein
				if( AddressMapper.mapAddress( protein, address ).isEmpty() )
				{
					// and the wildcard version exists in the protein
					AtomAddressReadable tweakedAddress = new AtomAddressReadable( address );
					tweakedAddress.setAtomName( tweakedAddress.getAtomName() + "#" );
					List<AtomAddressInternal> mappedAddresses = AddressMapper.mapAddress( protein, tweakedAddress );
					if( !mappedAddresses.isEmpty() )
					{
						// and the wildcard version is a pseudoatom
						Residue residue = protein.getResidue( mappedAddresses.get( 0 ) );
						if( getPseudoatoms().getPseudoatomName( residue.getAminoAcid(), tweakedAddress.getAtomName() ) != null )
						{
							// then add the wildcard to the address
							address.setAtomName( address.getAtomName() + "#" );
						}
					}
				}
				
				if( interpretAddress( protein, address ) > 0.0 )
				{
					numBuilt++;
				}
			}
		}
		return numBuilt;
	}
	
	
	/**************************
	 *   Functions
	 **************************/
	
	private void addPseudoatom( String pseudoatomName, ArrayList<String> referenceAtoms, Residue residue )
	{
		// if this residue already has this pseudoatom, bail
		if( residue.getAtomByName( pseudoatomName ) != null )
		{
			return;
		}
		
		// determine the average position of the reference atoms
		Vector3 avgPos = new Vector3();
		int numAtomsAdded = 0;
		for( String referenceAtom : referenceAtoms )
		{
			Atom atom = residue.getAtomByName( referenceAtom );
			assert( atom != null ) : "Didn't find atom: " + residue.getNumber() + ":" + referenceAtom;
			avgPos.add( atom.getPosition() );
			numAtomsAdded++;
		}
		assert( numAtomsAdded == referenceAtoms.size() );
		avgPos.scale( 1.0 / numAtomsAdded );
		
		// add a new atom to the residue
		Atom pseudoatom = new Atom();
		pseudoatom.setName( pseudoatomName );
		pseudoatom.setPosition( avgPos );
		pseudoatom.setId( residue.getAtoms().size() );
		residue.getAtoms().add( pseudoatom );
	}
	
	private double interpretAddress( Protein protein, AtomAddressReadable address )
	{
		// shortcut
		if( !address.isAmbiguous() )
		{
			return 0.0;
		}
		
		// get the amino acid for the residue
		List<AtomAddressInternal> addresses = AddressMapper.mapAddress( protein, address );
		if( addresses.isEmpty() )
		{
			return 0.0;
		}
		AminoAcid aminoAcid = protein.getResidue( addresses.get( 0 ) ).getAminoAcid();
		
		// is there a pseudoatom?
		String pseudoatomName = m_pseudoatoms.getPseudoatomName( aminoAcid, address.getAtomName() );
		if( pseudoatomName != null )
		{
			address.setAtomName( pseudoatomName );
			return m_pseudoatoms.getCorrection( aminoAcid, pseudoatomName );
		}
		
		return 0.0;
	}
}
