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

import java.util.ArrayList;
import java.util.Iterator;

import edu.duke.donaldLab.share.io.Transformer;
import edu.duke.donaldLab.share.mapping.NameMapper;
import edu.duke.donaldLab.share.mapping.NameScheme;
import edu.duke.donaldLab.share.math.CompareReal;
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.AtomAddressInternal;
import edu.duke.donaldLab.share.protein.Protein;
import edu.duke.donaldLab.share.protein.AtomAddressReadable;
import edu.duke.donaldLab.share.pseudoatoms.PseudoatomBuilder;
import edu.duke.donaldLab.share.test.ExtendedTestCase;

public class TestDistanceRestraintMapper extends ExtendedTestCase
{
	public void testNoeToRestraintAtoms( )
	throws Exception
	{
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		NameMapper.ensureProtein( protein, NameScheme.New );
		
		DistanceRestraintReader noeReader = new DistanceRestraintReader();
		ArrayList<DistanceRestraint<AtomAddressReadable>> readableRestraints = noeReader.read( "src/resources/test/large.noe" );
		NameMapper.ensureDistanceRestraints( protein, readableRestraints, NameScheme.New );
		
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = DistanceRestraintMapper.mapReadableToInternal( readableRestraints, protein, true );
		
		// make sure we have the right number of restraints
		assertEquals( readableRestraints.size(), restraints.size() );
		
		// for each restraint...
		Iterator<DistanceRestraint<AtomAddressInternal>> iterRestraint = restraints.iterator();
		while( iterRestraint.hasNext() )
		{
			DistanceRestraint<AtomAddressInternal> restraint = iterRestraint.next();
			
			// skip null restraints
			if( restraint == null )
			{
				continue;
			}
			
			// make sure the atoms exist
			for( AtomAddressInternal endpoint : restraint.getLefts() )
			{
				assertNotNull( protein.getAtom( endpoint ) );
			}
			for( AtomAddressInternal endpoint : restraint.getRights() )
			{
				assertNotNull( protein.getAtom( endpoint ) );
			}
		}
	}
	
	public void testNoeToRestraintPound( )
	throws Exception
	{
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		NameMapper.ensureProtein( protein, NameScheme.New );
		
		// add our test NOE
		// assign (resid 330 and name hg and segid B)(resid 332 and name hd1# and segid D) 4.0 2.2 3.5
		DistanceRestraint<AtomAddressReadable> noe = new DistanceRestraint<AtomAddressReadable>();
		noe.setLefts( new AtomAddressReadable( 'B', 330, "hg" ) );
		noe.setRights( new AtomAddressReadable( 'D', 332, "hd1#" ) );
		noe.setMinDistance( 1.8 );
		noe.setMaxDistance( 7.5 );
		ArrayList<DistanceRestraint<AtomAddressReadable>> noes = new ArrayList<DistanceRestraint<AtomAddressReadable>>();
		noes.add( noe );
		
		// get our restraints
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = DistanceRestraintMapper.mapReadableToInternal( noes, protein, true );
		assertEquals( 1, restraints.size() );
		
		// make sure this restraint was mapped correctly
		// should map to 1:11:12, (3:13:16 3:13:17 3:13:18), 1.8, 7.5
		DistanceRestraint<AtomAddressInternal> restraint = restraints.get( 0 );
		assertNotNull( restraint );
		assertEquals( 1, restraint.getLefts().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressInternal( 1, 11, 12 )
		), restraint.getLefts() );
		assertEquals( 3, restraint.getRights().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressInternal( 3, 13, 16 ),
			new AtomAddressInternal( 3, 13, 17 ),
			new AtomAddressInternal( 3, 13, 18 )
		), restraint.getRights() );
		assertEquals( 1.8, restraint.getMinDistance(), CompareReal.getEpsilon() );
		assertEquals( 7.5, restraint.getMaxDistance(), CompareReal.getEpsilon() );
	}
	
	public void testNoeToRestraintStar( )
	throws Exception
	{
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		
		// add our test NOE
		// assign (resid 350 and name hd* and segid D)(resid 354 and name he22 and segid C) 2.5 0.7 2.0
		DistanceRestraint<AtomAddressReadable> noe = new DistanceRestraint<AtomAddressReadable>();
		noe.setLefts( new AtomAddressReadable( 'D', 350, "hd*" ) );
		noe.setRights( new AtomAddressReadable( 'C', 354, "he22" ) );
		noe.setMinDistance( 1.8 );
		noe.setMaxDistance( 4.5 );
		ArrayList<DistanceRestraint<AtomAddressReadable>> noes = new ArrayList<DistanceRestraint<AtomAddressReadable>>();
		noes.add( noe );
		
		// get our restraints
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = DistanceRestraintMapper.mapReadableToInternal( noes, protein, true );
		assertEquals( 1, restraints.size() );
		
		// make sure this restraint was mapped correctly
		// should map to (3:31:13, 3:31:14, 3:31:15, 3:31:16, 3:31:17, 3:31:18), 2:35:16, 1.8, 4.5
		DistanceRestraint<AtomAddressInternal> restraint = restraints.get( 0 );
		assertNotNull( restraint );
		assertEquals( 6, restraint.getLefts().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressInternal( 3, 31, 13 ),
			new AtomAddressInternal( 3, 31, 14 ),
			new AtomAddressInternal( 3, 31, 15 ),
			new AtomAddressInternal( 3, 31, 16 ),
			new AtomAddressInternal( 3, 31, 17 ),
			new AtomAddressInternal( 3, 31, 18 )
		), restraint.getLefts() );
		assertEquals( 1, restraint.getRights().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressInternal( 2, 35, 16 )
		), restraint.getRights() );
		assertEquals( 1.8, restraint.getMinDistance(), CompareReal.getEpsilon() );
		assertEquals( 4.5, restraint.getMaxDistance(), CompareReal.getEpsilon() );
	}
	
	public void testUnambiguousRestraintToNoe( )
	throws Exception
	{
		// read in stuff
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		
		// build our restraint
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setLefts( new AtomAddressInternal( 1, 11, 12 ) );
		restraint.setRights( new AtomAddressInternal( 3, 13, 16 ) );
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.5 );
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>();
		restraints.add( restraint );
		
		// get our noes
		ArrayList<DistanceRestraint<AtomAddressReadable>> noes = DistanceRestraintMapper.mapInternalToReadable( restraints, protein );
		assertEquals( 1, noes.size() );
		
		// check the NOE, should be B:12:hg D:14:hd11 1.8 0.0 5.7
		DistanceRestraint<AtomAddressReadable> noe = noes.get( 0 );
		assertNotNull( noe );
		int firstResidueNumberm = protein.getSubunit( 0 ).getFirstResidueNumber() - 1;
		assertEquals( 1, noe.getLefts().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressReadable( 'B', firstResidueNumberm + 12, "hg" )
		), noe.getLefts() );
		assertEquals( 1, noe.getRights().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressReadable( 'D', firstResidueNumberm + 14, "hd11" )
		), noe.getRights() );
		assertEquals( 1.8, noe.getMinDistance() );
		assertEquals( 7.5, noe.getMaxDistance() );
	}

	public void testAmbiguousRestraintToNoe( )
	throws Exception
	{
		// read in stuff
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		PseudoatomBuilder.getInstance().build( protein );
		
		// build our restraint
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setLefts(
			new AtomAddressInternal( 1, 11, 10),
			new AtomAddressInternal( 1, 11, 11 )
		);
		restraint.setRights( new AtomAddressInternal( 3, 13, 16 ) );
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.5 );
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>();
		restraints.add( restraint );
		
		// get our noes
		ArrayList<DistanceRestraint<AtomAddressReadable>> noes = DistanceRestraintMapper.mapInternalToReadable( restraints, protein, true );
		assertEquals( 1, noes.size() );
		
		// check the NOE, should be (B:12:qb) D:14:hd11 1.8 0.0 5.7
		DistanceRestraint<AtomAddressReadable> noe = noes.get( 0 );
		assertNotNull( noe );
		int firstResidueNumberm = protein.getSubunit( 0 ).getFirstResidueNumber() - 1;
		assertEquals( 1, noe.getLefts().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressReadable( 'B', firstResidueNumberm + 12, "qb" )
		), noe.getLefts() );
		assertEquals( 1, noe.getRights().size() );
		assertEquals( Transformer.toTreeSet(
			new AtomAddressReadable( 'D', firstResidueNumberm + 14, "hd11" )
		), noe.getRights() );
		assertEquals( 1.8, noe.getMinDistance() );
		assertEquals( 7.5, noe.getMaxDistance() );
	}
	
	public void testMappingCycle( )
	throws Exception
	{
		/* Jeff: 6/09/2008
			The idea here is that we should be able to map back and forth from noes to
			distance restraints without losing any information.
		*/
		
		// read in stuff
		ProteinReader proteinReader = new ProteinReader();
		Protein protein = proteinReader.read( "src/resources/test/largeProtein.pdb" );
		NameMapper.ensureProtein( protein, NameScheme.New );
		DistanceRestraintReader noeReader = new DistanceRestraintReader();
		ArrayList<DistanceRestraint<AtomAddressReadable>> noes = noeReader.read( "src/resources/test/large.noe" );
		
		// HACKHACK: these NOEs are sadly all named with the wrong residue numbers, so fix them
		for( DistanceRestraint<AtomAddressReadable> noe : noes )
		{
			for( AtomAddressReadable address : noe.getLefts() )
			{
				address.setResidueNumber( address.getResidueNumber() + protein.getSubunit( address.getSubunitName() ).getFirstResidueNumber() - 1 );
			}
			for( AtomAddressReadable address : noe.getRights() )
			{
				address.setResidueNumber( address.getResidueNumber() + protein.getSubunit( address.getSubunitName() ).getFirstResidueNumber() - 1 );
			}
		}
		
		NameMapper.ensureDistanceRestraints( protein, noes, NameScheme.New );
		
		// map to distance restraints
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraints = DistanceRestraintMapper.mapReadableToInternal( noes, protein, false );
		
		// DEBUG: show the non-mappers
		if( noes.size() != restraints.size() )
		{
			ArrayList<DistanceRestraint<AtomAddressInternal>> restraintsWithNulls = DistanceRestraintMapper.mapReadableToInternal( noes, protein, true );
			assertEquals( noes.size(), restraintsWithNulls.size() );
			for( int i=0; i<noes.size(); i++ )
			{
				if( restraintsWithNulls.get( i ) == null )
				{
					System.err.println( noes.get( i ) );
				}
			}
			assertEquals( noes.size(), restraints.size() );
		}
		
		// map back to noes
		ArrayList<DistanceRestraint<AtomAddressReadable>> noesAgain = DistanceRestraintMapper.mapInternalToReadable( restraints, protein );
		
		// and map back to distance restraints again
		ArrayList<DistanceRestraint<AtomAddressInternal>> restraintsAgain = DistanceRestraintMapper.mapReadableToInternal( noesAgain, protein );
		
		// now make sure our restraints match!
		assertEquals( restraints.size(), restraintsAgain.size() );
		for( int i=0; i<restraints.size(); i++ )
		{
			assertEquals( restraints.get( i ), restraintsAgain.get( i ) );
		}
	}	
}
