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

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;

import edu.duke.donaldLab.jdshot.context.DistanceRestraintsContext;
import edu.duke.donaldLab.jdshot.context.RdcsContext;
import edu.duke.donaldLab.jdshot.context.SymmetryContext;
import edu.duke.donaldLab.jdshot.disco.AnnulusCalculator;
import edu.duke.donaldLab.jdshot.disco.AnnulusIndex;
import edu.duke.donaldLab.jdshot.disco.OrientationCalculator;
import edu.duke.donaldLab.jdshot.disco.AnnulusIndex.AnnulusEntry;
import edu.duke.donaldLab.jdshot.grid.Symmetry;
import edu.duke.donaldLab.jdshot.test.ExtendedTestCase;
import edu.duke.donaldLab.share.geom.Annulus2;
import edu.duke.donaldLab.share.geom.Vector2;
import edu.duke.donaldLab.share.geom.Vector3;
import edu.duke.donaldLab.share.nmr.AlignmentTensor;
import edu.duke.donaldLab.share.nmr.Assignment;
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.protein.AtomAddressInternal;
import edu.duke.donaldLab.share.protein.AtomAddressReadable;
import edu.duke.donaldLab.share.protein.Protein;
import edu.duke.donaldLab.share.protein.tools.MonomerCloner;

public class TestAnnulusCalculator extends ExtendedTestCase
{
	private static final double Epsilon = 1e-4;
	
	public void testDagkAnnuliSome( )
	throws Exception
	{
		// read inputs
		SymmetryContext symmetryContext = new SymmetryContext( Symmetry.Cn, new File( "src/resources/test/2KDC.monomer.protein" ), 3 );
		RdcsContext rdcsContext = new RdcsContext( symmetryContext, "src/resources/test/2KDC.experimental.rdc" );
		
		// previous steps
		AlignmentTensor tensor = AlignmentTensor.compute( symmetryContext.getMonomer(), rdcsContext.getInternalRdcs() );
		Vector3 orientation = OrientationCalculator.computeOrientations( 3, tensor ).get( 0 );
		OrientationCalculator.normalizeProtein( symmetryContext.getMonomer(), orientation );
		
		String assign = null;
		List<Annulus2> annuli = null;
		
		// check 1
		assign = "assign (resid 33 and name CB and segid A)((resid 96 and name CB and segid B) or (resid 96 and name CB and segid C)) 3.0 0.0 4.0";
		annuli = getAnnuli( symmetryContext, assign );
		assertEquals( new Annulus2( new Vector2( -7.914469, -5.903309 ), 0.000000, 3.358382 ), annuli.get( 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( 5.779047, -5.968768 ), 0.000000, 3.358382 ), annuli.get( 1 ), Epsilon );
		
		// check 2
		assign = "assign (resid 70 and name CB and segid A)((resid 99 and name CB and segid B) or (resid 99 and name CB and segid C)) 3.0 0.0 4.0";
		annuli = getAnnuli( symmetryContext, assign );
		assertEquals( new Annulus2( new Vector2( -7.867316, -10.751873 ), 0.000000, 3.275446 ), annuli.get( 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.769900, -3.990126 ), 0.000000, 3.275446 ), annuli.get( 1 ), Epsilon );
		
		// check 3
		assign = "assign (resid 70 and name CB and segid C)((resid 96 and name CB and segid A) or (resid 96 and name CB and segid B)) 3.0 0.0 4.0";
		annuli = getAnnuli( symmetryContext, assign );
		assertEquals( new Annulus2( new Vector2( -10.432922, -12.122239 ), 1.731855, 4.041368 ), annuli.get( 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.865931, -6.897188 ), 1.731855, 4.041368 ), annuli.get( 1 ), Epsilon );
		
		// check 4
		assign = "assign (resid 52 and name CB and segid C)((resid 117 and name CB and segid A) or (resid 117 and name CB and segid B)) 3.0 0.0 4.0";
		annuli = getAnnuli( symmetryContext, assign );
		assertEquals( new Annulus2( new Vector2( -11.366702, -12.131778 ), 1.510863, 3.951714 ), annuli.get( 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -2.960984, -10.397802 ), 1.510863, 3.951714 ), annuli.get( 1 ), Epsilon );
		
		// check 5
		assign = "assign (resid 55 and name CB and segid C)((resid 114 and name CB and segid A) or (resid 114 and name CB and segid B)) 3.0 0.0 4.0";
		annuli = getAnnuli( symmetryContext, assign );
		assertEquals( new Annulus2( new Vector2( -10.589915, -9.708849 ), 0.000000, 3.435031 ), annuli.get( 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -1.043337, -8.149793 ), 0.000000, 3.435031 ), annuli.get( 1 ), Epsilon );
	}
	
	public void testDagkAnnuliAll( )
	throws Exception
	{
		// read inputs
		SymmetryContext symmetryContext = new SymmetryContext( Symmetry.Cn, new File( "src/resources/test/2KDC.monomer.protein" ), 3 );
		RdcsContext rdcsContext = new RdcsContext( symmetryContext, "src/resources/test/2KDC.experimental.rdc" );
		final DistanceRestraintsContext distanceRestraintsContext = new DistanceRestraintsContext( symmetryContext, "src/resources/test/2KDC.experimental.onlyDisulfideBonds.noe" );
		
		// previous steps
		AlignmentTensor tensor = AlignmentTensor.compute( symmetryContext.getMonomer(), rdcsContext.getInternalRdcs() );
		Vector3 orientation = OrientationCalculator.computeOrientations( 3, tensor ).get( 0 );
		OrientationCalculator.normalizeProtein( symmetryContext.getMonomer(), orientation );
		
		final AnnulusIndex annulusIndex = AnnulusCalculator.computeAnnuli(
			symmetryContext.getMonomer(),
			distanceRestraintsContext.getInternalRestraints(),
			3
		);
			
		// make sure all distance restraints are present
		for( DistanceRestraint<AtomAddressInternal> restraint : distanceRestraintsContext.getInternalRestraints() )
		{
			assertNotNull( annulusIndex.getEntries( restraint ) );
		}
		
		// check counts
		assertEquals( 24, annulusIndex.getNumRestraints() );
		int numNonDegenerateCircles = 0;
		for( Entry<DistanceRestraint<AtomAddressInternal>,HashMap<Assignment<AtomAddressInternal>,AnnulusEntry>> entry : annulusIndex )
		{
			Collection<AnnulusEntry> entries = entry.getValue().values();
			assertEquals( 2, entries.size() );
			for( AnnulusEntry annulusEntry : entries )
			{
				if( annulusEntry.innerCircle != null )
				{
					numNonDegenerateCircles++;
				}
				
				if( annulusEntry.outerCircle != null )
				{
					numNonDegenerateCircles++;
				}
			}
		}
		assertEquals( 80, numNonDegenerateCircles );
		
		// just for some syntactic sugar
		class Helper
		{
			public Annulus2 get( int restraintId, int annulusId )
			{
				DistanceRestraint<AtomAddressInternal> restraint = distanceRestraintsContext.getInternalRestraints().get( restraintId );
				return annulusIndex.getEntries( restraint ).get( annulusId ).annulus;
			}
		}
		Helper helper = new Helper();
		
		// check individual annuli
		assertEquals( new Annulus2( new Vector2( -7.530658, -4.749182 ), 2.309401, 5.773503 ), helper.get( 0, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.530658, -4.749182 ), 2.309401, 5.773503 ), helper.get( 0, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.665511, -3.985224 ), 1.732051, 4.041452 ), helper.get( 1, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.665511, -3.985224 ), 1.732051, 4.041452 ), helper.get( 1, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.407949, -5.036420 ), 2.309401, 5.196152 ), helper.get( 2, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.407949, -5.036420 ), 2.309401, 5.196152 ), helper.get( 2, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.653707, -6.554081 ), 1.732051, 4.041452 ), helper.get( 3, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.653707, -6.554081 ), 1.732051, 4.041452 ), helper.get( 3, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.653477, -3.079197 ), 2.309401, 5.773503 ), helper.get( 4, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.653477, -3.079197 ), 2.309401, 5.773503 ), helper.get( 4, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.410036, -4.199001 ), 1.732051, 4.041452 ), helper.get( 5, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -8.410036, -4.199001 ), 1.732051, 4.041452 ), helper.get( 5, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -11.409821, -4.594971 ), 2.309401, 5.773503 ), helper.get( 6, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -11.409821, -4.594971 ), 2.309401, 5.773503 ), helper.get( 6, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -12.879686, -4.771069 ), 1.732051, 4.041452 ), helper.get( 7, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -12.879686, -4.771069 ), 1.732051, 4.041452 ), helper.get( 7, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -6.195186, -5.859333 ), 1.058862, 5.396405 ), helper.get( 8, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( 6.292265, -5.953565 ), 1.058862, 5.396405 ), helper.get( 8, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.914469, -5.903309 ), 0.000000, 3.358382 ), helper.get( 9, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( 5.779047, -5.968768 ), 0.000000, 3.358382 ), helper.get( 9, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -5.310060, -7.625882 ), 2.032637, 5.668475 ), helper.get( 10, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( 8.264704, -6.070298 ), 2.032637, 5.668475 ), helper.get( 10, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -6.173619, -7.447862 ), 1.642427, 4.764196 ), helper.get( 11, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( 7.987094, -5.233424 ), 1.642427, 4.764196 ), helper.get( 11, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -9.709254, -12.474115 ), 2.017107, 5.662925 ), helper.get( 12, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -1.929968, -10.498367 ), 2.017107, 5.662925 ), helper.get( 12, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.379860, -11.803229 ), 0.000000, 3.428434 ), helper.get( 13, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -2.752094, -9.378896 ), 0.000000, 3.428434 ), helper.get( 13, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -11.132174, -12.863588 ), 1.835327, 5.600752 ), helper.get( 14, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -2.304134, -11.925390 ), 1.835327, 5.600752 ), helper.get( 14, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -11.366702, -12.131778 ), 1.510863, 3.951714 ), helper.get( 15, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -2.960984, -10.397802 ), 1.510863, 3.951714 ), helper.get( 15, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -11.418031, -10.846686 ), 0.000000, 5.120569 ), helper.get( 16, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -1.374962, -8.204809 ), 0.000000, 5.120569 ), helper.get( 16, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.589915, -9.708849 ), 0.000000, 3.435031 ), helper.get( 17, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -1.043337, -8.149793 ), 0.000000, 3.435031 ), helper.get( 17, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.140299, -11.815900 ), 2.178678, 5.722468 ), helper.get( 18, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.838830, -5.515281 ), 2.178678, 5.722468 ), helper.get( 18, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.432922, -12.122239 ), 1.731855, 4.041368 ), helper.get( 19, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.865931, -6.897188 ), 1.731855, 4.041368 ), helper.get( 19, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.877685, -11.346833 ), 0.000000, 4.830742 ), helper.get( 20, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.113746, -3.321266 ), 0.000000, 4.830742 ), helper.get( 20, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -7.867316, -10.751873 ), 0.000000, 3.275446 ), helper.get( 21, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.769900, -3.990126 ), 0.000000, 3.275446 ), helper.get( 21, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.181789, -11.367792 ), 0.000000, 4.626335 ), helper.get( 22, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.471502, -5.255295 ), 0.000000, 4.626335 ), helper.get( 22, 1 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -10.813595, -11.023819 ), 0.000000, 3.011387 ), helper.get( 23, 0 ), Epsilon );
		assertEquals( new Annulus2( new Vector2( -0.105008, -6.018306 ), 0.000000, 3.011387 ), helper.get( 23, 1 ), Epsilon );
	}
	
	public void testGb1AnnuliAll( )
	throws Exception
	{
		// read inputs
		SymmetryContext symmetryContext = new SymmetryContext( Symmetry.Cn, new File( "src/resources/test/1Q10.monomer.protein" ), 2 );
		RdcsContext rdcsContext = new RdcsContext( symmetryContext, "src/resources/test/1Q10.experimental.rdc" );
		final DistanceRestraintsContext distanceRestraintsContext = new DistanceRestraintsContext( symmetryContext, "src/resources/test/1Q10.experimental.fixed.noe" );
		
		// previous steps
		AlignmentTensor tensor = AlignmentTensor.compute( symmetryContext.getPseudoMonomer(), rdcsContext.getInternalRdcs() );
		Vector3 orientation = OrientationCalculator.computeOrientations( 2, tensor ).get( 0 );
		OrientationCalculator.normalizeProtein( symmetryContext.getPseudoMonomer(), orientation );
		
		final AnnulusIndex annulusIndex = AnnulusCalculator.computeAnnuli(
			symmetryContext.getPseudoMonomer(),
			distanceRestraintsContext.getInternalRestraints(),
			2
		);
		
		// make sure all distance restraints are present
		for( DistanceRestraint<AtomAddressInternal> restraint : distanceRestraintsContext.getInternalRestraints() )
		{
			assertNotNull( annulusIndex.getEntries( restraint ) );
		}
		
		// check counts
		assertEquals( 296, annulusIndex.getNumRestraints() );
		int numNonDegenerateCircles = 0;
		for( Entry<DistanceRestraint<AtomAddressInternal>,HashMap<Assignment<AtomAddressInternal>,AnnulusEntry>> entry : annulusIndex )
		{
			Collection<AnnulusEntry> entries = entry.getValue().values();
			assertEquals( 1, entries.size() );
			for( AnnulusEntry annulusEntry : entries )
			{
				if( annulusEntry.innerCircle != null )
				{
					numNonDegenerateCircles++;
				}
				
				if( annulusEntry.outerCircle != null )
				{
					numNonDegenerateCircles++;
				}
			}
		}
		assertEquals( 414, numNonDegenerateCircles );
	}
	
	private List<Annulus2> getAnnuli( SymmetryContext symmetryContext, String assign )
	throws Exception
	{
		// read the distance restraint
		InputStream stringIn = new ByteArrayInputStream( assign.getBytes() );
		List<DistanceRestraint<AtomAddressReadable>> restraintsReadable = new DistanceRestraintReader().read( stringIn );
		Protein clonedOligomer = MonomerCloner.clone( symmetryContext.getPseudoMonomer(), symmetryContext.getNumSubunits() );
		List<DistanceRestraint<AtomAddressInternal>> restraints = DistanceRestraintMapper.mapReadableToInternal( restraintsReadable, clonedOligomer );
		
		// compute the annulus
		AnnulusIndex annulusIndex = AnnulusCalculator.computeAnnuli( symmetryContext.getMonomer(), restraints, 3 );
		List<AnnulusEntry> entries = annulusIndex.getEntries( restraints.get( 0 ) );
		ArrayList<Annulus2> annuli = new ArrayList<Annulus2>( entries.size() );
		for( AnnulusEntry entry : entries )
		{
			annuli.add( entry.annulus );
		}
		return annuli;
	}
}
