/*******************************************************************************
 * 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.List;
import java.util.Set;

import junit.framework.AssertionFailedError;

import edu.duke.donaldLab.share.mapping.NameMapper;
import edu.duke.donaldLab.share.mapping.NameScheme;
import edu.duke.donaldLab.share.nmr.ChemicalShift;
import edu.duke.donaldLab.share.nmr.ChemicalShiftMapper;
import edu.duke.donaldLab.share.nmr.ChemicalShiftReader;
import edu.duke.donaldLab.share.nmr.DistanceRestraint;
import edu.duke.donaldLab.share.nmr.DistanceRestraintFilterer;
import edu.duke.donaldLab.share.nmr.DistanceRestraintMapper;
import edu.duke.donaldLab.share.nmr.DistanceRestraintReassigner;
import edu.duke.donaldLab.share.nmr.MappedChemicalShift;
import edu.duke.donaldLab.share.nmr.MappedChemicalShiftPair;
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.Element;
import edu.duke.donaldLab.share.protein.Protein;
import edu.duke.donaldLab.share.protein.AtomAddressReadable;
import edu.duke.donaldLab.share.protein.tools.MonomerCloner;
import edu.duke.donaldLab.share.pseudoatoms.PseudoatomBuilder;
import edu.duke.donaldLab.share.test.ExtendedTestCase;

public class TestDistanceRestraintReassigner extends ExtendedTestCase
{
	final double HydrogenWindowSize = 0.05; // in ppm
	final double CarbonWindowSize = 0.5;
	final double NitrogenWindowSize = 0.5;
	
	private Protein m_monomer;
	private Protein m_oligomer;
	private List<MappedChemicalShift> m_mappedHydrogenShifts;
	private List<MappedChemicalShift> m_mappedCarbonShifts;
	private List<MappedChemicalShift> m_mappedNitrogenShifts;
	private List<MappedChemicalShiftPair> m_carbonPairs;
	private List<MappedChemicalShiftPair> m_nitrogenPairs;
	private List<DistanceRestraint<AtomAddressReadable>> m_noes;
	private List<DistanceRestraint<AtomAddressInternal>> m_restraints;
	
	@Override
	public void setUp( )
	throws Exception
	{
		// read the protein monomer
		m_monomer = new ProteinReader().read( "src/resources/test/1Q10.monomer.protein" );
		NameMapper.ensureProtein( m_monomer, NameScheme.New );
		PseudoatomBuilder.getInstance().build( m_monomer );
		
		// clone an oligomer
		m_oligomer = MonomerCloner.clone( m_monomer.getSubunit( 0 ), 2 );
		
		// read the shifts
		List<ChemicalShift> shifts = new ChemicalShiftReader().read( "src/resources/test/1Q10.experimental.shift" );
		NameMapper.ensureShifts( m_monomer, shifts, NameScheme.New );
		
		// filter the shifts
		List<ChemicalShift> hydrogenShifts = ChemicalShiftMapper.filter( shifts, Element.Hydrogen );
		List<ChemicalShift> carbonShifts = ChemicalShiftMapper.filter( shifts, Element.Carbon );
		List<ChemicalShift> nitrogenShifts = ChemicalShiftMapper.filter( shifts, Element.Nitrogen );
		
		// map the shifts
		PseudoatomBuilder.getInstance().buildShifts( m_monomer, hydrogenShifts );
		m_mappedHydrogenShifts = ChemicalShiftMapper.map( hydrogenShifts, m_monomer );
		m_mappedCarbonShifts = ChemicalShiftMapper.map( carbonShifts, m_monomer );
		m_mappedNitrogenShifts = ChemicalShiftMapper.map( nitrogenShifts, m_monomer );
		m_carbonPairs = ChemicalShiftMapper.associatePairs(
			m_monomer,
			m_mappedHydrogenShifts,
			m_mappedCarbonShifts,
			Element.Carbon
		);
		m_nitrogenPairs = ChemicalShiftMapper.associatePairs(
			m_monomer,
			m_mappedHydrogenShifts,
			m_mappedNitrogenShifts,
			Element.Nitrogen
		);
		
		// read the NOEs
		List<DistanceRestraint<AtomAddressReadable>> noes = new DistanceRestraintReader().read( "src/resources/test/1Q10.experimental.fixed.noe" );
		PseudoatomBuilder.getInstance().buildDistanceRestraints( m_oligomer, noes );
		m_restraints = DistanceRestraintMapper.mapReadableToInternal( noes, m_oligomer );
		m_restraints = DistanceRestraintFilterer.pickIntersubunit( m_restraints );
		m_restraints = DistanceRestraintFilterer.pickSubunitEitherSide( m_restraints, 0 );
		m_restraints = DistanceRestraintFilterer.pickUnique( m_restraints );
		m_restraints = DistanceRestraintFilterer.pickOneFromSymmetricGroup( m_restraints );
		m_noes = DistanceRestraintMapper.mapInternalToReadable( m_restraints, m_oligomer );
	}
	
	public void testNullReassignment1D( )
	throws Exception
	{
		// don't actually reassign any restraints
		List<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = DistanceRestraintReassigner.reassign1D(
			m_oligomer, m_restraints, m_mappedHydrogenShifts,
			-1
		);
		assertSameRestraints( m_restraints, reassignedRestraints );
	}
	
	public void testNullReassignment2x3D( )
	throws Exception
	{
		// don't actually reassign any restraints
		List<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = DistanceRestraintReassigner.reassignDouble3D(
			m_oligomer, m_restraints, m_mappedHydrogenShifts, m_carbonPairs, m_nitrogenPairs,
			-1, -1, -1
		);
		assertSameRestraints( m_restraints, reassignedRestraints );
	}
	
	public void testPseudoatomTranslation1( )
	throws Exception
	{
		// mutilated: ((0,0,9),(0,0,10))	((1,49,17),(1,49,18))		
		// original: ((0,0,21))	((1,49,22))
		
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.0 );
		restraint.setLefts( new AtomAddressInternal( 0, 0, 21 ) );
		restraint.setRights( new AtomAddressInternal( 1, 49, 22 ) );
		
		assertNullReassignments( restraint );
	}
	
	public void testPseudoatomTranslation2( )
	throws Exception
	{
		// mutilated: ((0,0,19),(0,5,10))	((1,49,10))
		// original: ((0,0,19))	((1,49,10))
		
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.0 );
		restraint.setLefts( new AtomAddressInternal( 0, 0, 19 ) );
		restraint.setRights( new AtomAddressInternal( 1, 49, 10 ) );
		
		assertNullReassignments( restraint );
	}
	
	public void testPseudoatomTranslation3( )
	throws Exception
	{
		// mutilated: ((0,4,18),(0,4,17))	((1,11,22),(1,11,21))
		// original: ((0,4,16))	((1,11,20))
		
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.0 );
		restraint.setLefts( new AtomAddressInternal( 0, 4, 16 ) );
		restraint.setRights( new AtomAddressInternal( 1, 11, 20 ) );
		
		assertNullReassignments( restraint );
	}
	
	public void testOriginalAssignmentsSurvived( )
	throws Exception
	{
		reassignAssertOriginal( m_restraints );
	}
	
	public void testOriginalAssignments1( )
	throws Exception
	{
		// original: ((0,0,21))	((1,49,22))
		// reassigned: ((0,38,9),(0,4,9),(0,20,9),(0,55,12),(0,0,9),(0,0,10),(0,55,13))
		//             ((1,2,14),(1,26,10),(1,44,14),(1,9,17),(1,7,10),(1,21,10),(1,12,17),(1,49,17),(1,49,18))
		
		DistanceRestraint<AtomAddressInternal> restraint = new DistanceRestraint<AtomAddressInternal>();
		restraint.setMinDistance( 1.8 );
		restraint.setMaxDistance( 7.0 );
		restraint.setLefts( new AtomAddressInternal( 0, 4, 16 ) );
		restraint.setRights( new AtomAddressInternal( 1, 11, 20 ) );
		
		reassignAssertOriginal( restraint );
	}
	
	public void testOriginalNoeAssignmentsSurvived( )
	throws Exception
	{
		List<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = null;
		
		// test 1D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassign1D(
			m_oligomer, m_restraints, m_mappedHydrogenShifts,
			HydrogenWindowSize
		);
		List<DistanceRestraint<AtomAddressReadable>> reassignedNoes = DistanceRestraintMapper.mapInternalToReadable( reassignedRestraints, m_oligomer, true );
		assertOriginalNoeAssignments( m_noes, reassignedNoes );
		
		// test 2x3D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassignDouble3D(
			m_oligomer, m_restraints, m_mappedHydrogenShifts, m_carbonPairs, m_nitrogenPairs,
			HydrogenWindowSize, CarbonWindowSize, NitrogenWindowSize
		);
		reassignedNoes = DistanceRestraintMapper.mapInternalToReadable( reassignedRestraints, m_oligomer, true );
		assertOriginalNoeAssignments( m_noes, reassignedNoes );
	}
	
	private void assertSameRestraints( List<DistanceRestraint<AtomAddressInternal>> expected, List<DistanceRestraint<AtomAddressInternal>> observed )
	{
		assertEquals( expected.size(), observed.size() );
		for( int i=0; i<expected.size(); i++ )
		{
			assertEquals( expected.get( i ), observed.get( i ) );
		}
	}
	
	private void assertNullReassignments( DistanceRestraint<AtomAddressInternal> restraint )
	throws Exception
	{
		List<DistanceRestraint<AtomAddressInternal>> restraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>();
		restraints.add( restraint );
		assertNullReassignments( restraints );
	}
	
	private void assertNullReassignments( List<DistanceRestraint<AtomAddressInternal>> restraints )
	throws Exception
	{
		List<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = null;
		
		// test 1D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassign1D(
			m_oligomer, restraints, m_mappedHydrogenShifts,
			-1
		);
		assertSameRestraints( restraints, reassignedRestraints );
		
		// test 2x3D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassignDouble3D(
			m_oligomer, restraints, m_mappedHydrogenShifts, m_carbonPairs, m_nitrogenPairs,
			-1, -1, -1
		);
		assertSameRestraints( restraints, reassignedRestraints );
	}
	
	private void assertOriginalAssignments( List<DistanceRestraint<AtomAddressInternal>> original, List<DistanceRestraint<AtomAddressInternal>> reassigned )
	{
		assertEquals( original.size(), reassigned.size() );
		for( int i=0; i<original.size(); i++ )
		{
			assertOriginalAssignments( original.get( i ), reassigned.get( i ) );
		}
	}
	
	private void assertOriginalAssignments( DistanceRestraint<AtomAddressInternal> original, DistanceRestraint<AtomAddressInternal> reassigned )
	{
		try
		{
			assertOriginalAssignments( original.getLefts(), reassigned.getLefts() );
			assertOriginalAssignments( original.getRights(), reassigned.getRights() );
		}
		catch( AssertionFailedError err )
		{
			throw new AssertionFailedError( "Original assignments missing!! original<" + original + "> reassigned<" + reassigned + ">" );
		}
	}
	
	private void assertOriginalAssignments( Set<AtomAddressInternal> original, Set<AtomAddressInternal> reassigned )
	{
		for( AtomAddressInternal address : original )
		{
			assertTrue( reassigned.contains( address ) );
		}
	}
	
	private void assertOriginalNoeAssignments( List<DistanceRestraint<AtomAddressReadable>> original, List<DistanceRestraint<AtomAddressReadable>> reassigned )
	{
		assertEquals( original.size(), reassigned.size() );
		for( int i=0; i<original.size(); i++ )
		{
			assertOriginalNoeAssignments( original.get( i ), reassigned.get( i ) );
		}
	}
	
	private void assertOriginalNoeAssignments( DistanceRestraint<AtomAddressReadable> original, DistanceRestraint<AtomAddressReadable> reassigned )
	{
		try
		{
			assertOriginalNoeAssignments( original.getLefts(), reassigned.getLefts() );
			assertOriginalNoeAssignments( original.getRights(), reassigned.getRights() );
		}
		catch( AssertionFailedError err )
		{
			throw new AssertionFailedError( "Original NOE assignments missing!! original<" + original + "> reassigned<" + reassigned + ">" );
		}
	}
	
	private void assertOriginalNoeAssignments( Set<AtomAddressReadable> original, Set<AtomAddressReadable> reassigned )
	{
		for( AtomAddressReadable address : original )
		{
			assertTrue( reassigned.contains( address ) );
		}
	}
	
	private void reassignAssertOriginal( DistanceRestraint<AtomAddressInternal> restraint )
	throws Exception
	{
		List<DistanceRestraint<AtomAddressInternal>> restraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>();
		restraints.add( restraint );
		reassignAssertOriginal( restraints );
	}
	
	private void reassignAssertOriginal( List<DistanceRestraint<AtomAddressInternal>> restraints )
	throws Exception
	{
		List<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = null;
		
		// test 1D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassign1D(
			m_oligomer, restraints, m_mappedHydrogenShifts,
			HydrogenWindowSize
		);
		assertOriginalAssignments( restraints, reassignedRestraints );
		
		// test 2x3D reassignment
		reassignedRestraints = DistanceRestraintReassigner.reassignDouble3D(
			m_oligomer, restraints, m_mappedHydrogenShifts, m_carbonPairs, m_nitrogenPairs,
			HydrogenWindowSize, CarbonWindowSize, NitrogenWindowSize
		);
		assertOriginalAssignments( restraints, reassignedRestraints );
	}
}
