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

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;

import edu.duke.donaldLab.jdshot.grid.SearchSpace;
import edu.duke.donaldLab.jdshot.grid.Symmetry;
import edu.duke.donaldLab.share.geom.Sphere;
import edu.duke.donaldLab.share.geom.Vector3;
import edu.duke.donaldLab.share.mapping.NameMapper;
import edu.duke.donaldLab.share.mapping.NameScheme;
import edu.duke.donaldLab.share.nmr.DistanceRestraint;
import edu.duke.donaldLab.share.nmr.DistanceRestraintMapper;
import edu.duke.donaldLab.share.nmr.DistanceRestraintReader;
import edu.duke.donaldLab.share.pdb.ProteinReader;
import edu.duke.donaldLab.share.protein.Atom;
import edu.duke.donaldLab.share.protein.AtomAddressInternal;
import edu.duke.donaldLab.share.protein.AtomAddressReadable;
import edu.duke.donaldLab.share.protein.Protein;
import edu.duke.donaldLab.share.protein.Subunit;
import edu.duke.donaldLab.share.protein.SubunitOrder;
import edu.duke.donaldLab.share.protein.tools.MonomerCloner;
import edu.duke.donaldLab.share.pseudoatoms.PseudoatomBuilder;

public class SearchContext
{
	/**************************
	 *   Data Members
	 **************************/
	
	private Symmetry m_symmetry;
	private Subunit m_monomer;
	private ArrayList<DistanceRestraint<AtomAddressReadable>> m_readableRestraints;
	private ArrayList<DistanceRestraint<AtomAddressInternal>> m_internalRestraints;
	private SubunitOrder m_subunitOrder = null;
	private DistanceRestraint<AtomAddressInternal> m_maxDistanceRestraint;
	private Sphere m_boundingSphere;
	private SearchSpace m_searchSpace;
	
	
	/**************************
	 *   Constructors
	 **************************/
	
	public SearchContext(  )
	{
		m_symmetry = null;
		m_monomer = null;
		m_readableRestraints = null;
		m_internalRestraints = null;
		m_subunitOrder = null;
		m_maxDistanceRestraint = null;
		m_boundingSphere = null;
		m_searchSpace = null;
	}
	
	
	/**************************
	 *   Accessors
	 **************************/
	
	public Symmetry getSymmetry( )
	{
		return m_symmetry;
	}
	
	public Subunit getMonomer( )
	{
		return m_monomer;
	}
	
	public ArrayList<DistanceRestraint<AtomAddressReadable>> getNoes( )
	{
		return m_readableRestraints;
	}
	
	public ArrayList<DistanceRestraint<AtomAddressInternal>> getRestraints( )
	{
		return m_internalRestraints;
	}
	
	public SubunitOrder getSubunitOrder( )
	{
		return m_subunitOrder;
	}
	
	public int getNumSubunits( )
	{
		return m_subunitOrder.getNumSubunits();
	}
	
	public SearchSpace getSearchSpace( )
	{
		return m_searchSpace;
	}
	
	public DistanceRestraint<AtomAddressInternal> getMaxDistanceRestraint( )
	{
		return m_maxDistanceRestraint;
	}
	
	public Sphere getBoundingSphere( )
	{
		return m_boundingSphere;
	}
	
	
	/**************************
	 *   Methods
	 **************************/
	
	public void load( Symmetry symmetry, File monomerFile, File noesFile, SubunitOrder subunitOrder )
	throws IOException
	{
		// save parameters
		m_symmetry = symmetry;
		m_subunitOrder = subunitOrder;
		
		// read the protein
		ProteinReader proteinReader = new ProteinReader();
		Protein monomerProtein = proteinReader.read( monomerFile );
		NameMapper.ensureProtein( monomerProtein, NameScheme.New );
		m_monomer = monomerProtein.getSubunit( 0 );
		
		// read the noes if needed
		if( noesFile != null )
		{
			// read the restraints
			DistanceRestraintReader noeReader = new DistanceRestraintReader();
			ArrayList<DistanceRestraint<AtomAddressReadable>> readableRestraints = noeReader.read( noesFile );
			System.out.println( "SearchContext: read " + readableRestraints.size() + " NOEs" );
			
			// add pseudoatoms if needed
			if( PseudoatomBuilder.distanceRestraintsHavePseudoatoms( readableRestraints ) )
			{
				PseudoatomBuilder.getInstance().build( monomerProtein );
				System.out.println( "SearchContext: Built Pseudoatoms!" );
			}
			
			// make an oligomer protein for the mapping
			Protein oligomerProtein = MonomerCloner.clone( m_monomer, getNumSubunits() );
			NameMapper.ensureDistanceRestraints( oligomerProtein, readableRestraints, NameScheme.New );
			
			// map to internal addresses
			ArrayList<DistanceRestraint<AtomAddressInternal>> internalRestraints = DistanceRestraintMapper.mapReadableToInternal( readableRestraints, oligomerProtein );
			System.out.println( "SearchContext: " + internalRestraints.size() + " restraints" );
			m_internalRestraints = filterRestraints( internalRestraints );
			System.out.println( "SearchContext: " + m_internalRestraints.size() + " filtered restraints" );
			
			// map back to readable
			m_readableRestraints = DistanceRestraintMapper.mapInternalToReadable( m_internalRestraints, oligomerProtein );
			System.out.println( "SearchContext: using " + m_readableRestraints.size() + " NOEs" );
			
			checkRestraints();
			m_maxDistanceRestraint = calculateMaxDistanceRestraint();
			
			// just in case...
			assert( m_readableRestraints.size() == m_internalRestraints.size() );
		}
		
		calculateBoundingSphere();
		centerMonomer();
		
		// init the search space
		m_searchSpace = m_symmetry.newSearchSpace( m_boundingSphere.radius, getMaxDistance(), m_subunitOrder.getNumSubunits() );
	}
	
	// UNDONE: these two functions need to be rewritten/removed. We need to rethink how subunit mapping happens for Dn/Cn proteins
	public int getSubunitId( SubunitType type )
	{
		// TEMP
		throw new UnsupportedOperationException();
		//return m_subunitOrder.mapComputedToReference( type.ordinal() );
	}
	public SubunitType getSubunitType( int subunitId )
	{
		// TEMP
		throw new UnsupportedOperationException();
		//return SubunitType.values()[m_subunitOrder.mapComputedToReference( subunitId )];
	}
	
	public SubunitType getRestraintType( DistanceRestraint<AtomAddressInternal> restraint )
	{
		// HACKHACK: won't work for subunit-ambiguous restraints!
		return getRestraintType( restraint.getLefts().iterator().next(), restraint.getRights().iterator().next() );
	}
	
	public double getMaxDistance( )
	{
		double distance = 0.0;
		
		if( m_maxDistanceRestraint != null )
		{
			distance = m_maxDistanceRestraint.getMaxDistance();
		}
		
		return distance;
	}
	
	
	/**************************
	 *   Functions
	 **************************/

	private SubunitType getRestraintType( AtomAddressInternal left, AtomAddressInternal right )
	{
		// find the main subunit
		AtomAddressInternal other = null;
		if( left.getSubunitId() == getSubunitId( SubunitType.Main ) )
		{
			other = right;
		}
		else if( right.getSubunitId() == getSubunitId( SubunitType.Main ) )
		{
			other = left;
		}
		else
		{
			// the main subunit is not referenced, bail
			return SubunitType.Unknown;
		}
		
		return getSubunitType( other.getSubunitId() );
	}
		
	private ArrayList<DistanceRestraint<AtomAddressInternal>> filterRestraints( ArrayList<DistanceRestraint<AtomAddressInternal>> restraints )
	{
		// UNDONE: change to use NoeFilterer
		
		HashSet<DistanceRestraint<AtomAddressInternal>> filteredRestraints = new HashSet<DistanceRestraint<AtomAddressInternal>>();
		
		// for each restraint...
		Iterator<DistanceRestraint<AtomAddressInternal>> iterRestraint = restraints.iterator();
		while( iterRestraint.hasNext() )
		{
			DistanceRestraint<AtomAddressInternal> restraint = iterRestraint.next();
			
			// we only care about intersubunit restraints
			if( getRestraintType( restraint ) == SubunitType.Unknown )
			{
				continue;
			}
			
			// check for restraints with no assignments
			assert( restraint.getLefts().size() > 0 );
			assert( restraint.getRights().size() > 0 );
			
			// put the main subunit on the left if needed
			// NOTE: we assume all assignments per endpoint are to the same residue
			if( restraint.getLefts().iterator().next().getSubunitId() != getSubunitId( SubunitType.Main ) )
			{
				restraint.swap();
			}
			
			// just in case...
			assert( restraint.getLefts().iterator().next().getSubunitId() == getSubunitId( SubunitType.Main ) );
			
			// only add unique restraints
			if( !filteredRestraints.contains( restraint ) )
			{
				filteredRestraints.add( restraint );
			}
		}
		
		// return the set contents
		ArrayList<DistanceRestraint<AtomAddressInternal>> retval = new ArrayList<DistanceRestraint<AtomAddressInternal>>( filteredRestraints.size() );
		retval.addAll( filteredRestraints );
		return retval;
	}
	
	private void checkRestraints( )
	{
		// UNDONE: move to Dn-specific section
		
		// just in case, make sure we have intradimer and interdimer restraints
		// we assume that AC is the dimer
		int numIntradimer = 0;
		int numInterdimerStandard = 0;
		int numInterdimerGoofy = 0;
		for( DistanceRestraint<AtomAddressInternal> restraint : m_internalRestraints )
		{
			switch( getRestraintType( restraint ) )
			{
				case Intradimer:
					numIntradimer++;
				break;
				
				case InterdimerStandard:
					numInterdimerStandard++;
				break;

				case InterdimerGoofy:
					numInterdimerGoofy++;
				break;
			}
		}
		
		System.out.println( "SearchContext: Intradimer restraints: " + numIntradimer );
		System.out.println( "SearchContext: Interdimer Standard restraints: " + numInterdimerStandard );
		System.out.println( "SearchContext: Interdimer Goofy restraints: " + numInterdimerGoofy );
	}
	
	private DistanceRestraint<AtomAddressInternal> calculateMaxDistanceRestraint( )
	{
		DistanceRestraint<AtomAddressInternal> maxRestraint = null;
		
		for( DistanceRestraint<AtomAddressInternal> restraint : m_internalRestraints )
		{
			if( maxRestraint == null || restraint.getMaxDistance() > maxRestraint.getMaxDistance() )
			{
				maxRestraint = restraint;
			}
		}
		
		// just in case
		assert( maxRestraint != null );
		
		return maxRestraint;
	}
	
	private void calculateBoundingSphere( )
	{
		// grab the positions vectors from the atoms
		Vector3[] points = new Vector3[m_monomer.getAtomIndex().size()];
		int atomIndex = 0;
		for( AtomAddressInternal address : m_monomer.getAtomIndex() )
		{
			Atom atom = m_monomer.getAtom( address );
			points[atomIndex++] = atom.getPosition();
		}
		m_boundingSphere = new Sphere( points );
	}
	
	private void centerMonomer( )
	{
		for( AtomAddressInternal address : m_monomer.getAtomIndex() )
		{
			Atom atom = m_monomer.getAtom( address );
			atom.getPosition().subtract( m_boundingSphere.center );
		}
	}
}
