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

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import edu.duke.donaldLab.share.bond.Bond;
import edu.duke.donaldLab.share.bond.BondGraph;
import edu.duke.donaldLab.share.bond.BondGraphBuilder;
import edu.duke.donaldLab.share.io.Transformer;
import edu.duke.donaldLab.share.mapping.AddressMapper;
import edu.duke.donaldLab.share.protein.AminoAcid;
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.Residue;
import edu.duke.donaldLab.share.pseudoatoms.PseudoatomBuilder;

public class DistanceRestraintReassigner
{
	/**************************
	 *   Static Methods
	 **************************/
	
	public static List<DistanceRestraint<AtomAddressInternal>> reassign1D( Protein protein, List<DistanceRestraint<AtomAddressInternal>> restraints, List<MappedChemicalShift> hydrogenShifts, double hydrogenWindowSize )
	{
		ArrayList<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>( restraints.size() );
		for( DistanceRestraint<AtomAddressInternal> restraint : restraints )
		{
			DistanceRestraint<AtomAddressInternal> reassignedRestraint = new DistanceRestraint<AtomAddressInternal>( restraint );
			reassignedRestraint.setLefts( getNearbyAddresses1D( protein, restraint.getLefts(), hydrogenShifts, hydrogenWindowSize ) );
			reassignedRestraint.setRights( getNearbyAddresses1D( protein, restraint.getRights(), hydrogenShifts, hydrogenWindowSize ) );
			reassignedRestraints.add( reassignedRestraint );
			
			restorePseudoatoms( protein, restraint, reassignedRestraint );
		}
		
		return reassignedRestraints;
	}
	
	public static List<DistanceRestraint<AtomAddressInternal>> reassignDouble3D( Protein protein, List<DistanceRestraint<AtomAddressInternal>> restraints, List<MappedChemicalShift> hydrogenShifts, List<MappedChemicalShiftPair> carbonPairs, List<MappedChemicalShiftPair> nitrogenPairs, double hydrogenWindowSize, double carbonWindowSize, double nitrogenWindowSize )
	throws IOException
	{
		// build the bond graphs for the protein
		ArrayList<BondGraph> bondGraphs = BondGraphBuilder.getInstance().build( protein );
		
		ArrayList<DistanceRestraint<AtomAddressInternal>> reassignedRestraints = new ArrayList<DistanceRestraint<AtomAddressInternal>>( restraints.size() );
		for( DistanceRestraint<AtomAddressInternal> restraint : restraints )
		{
			DistanceRestraint<AtomAddressInternal> reassignedRestraint = new DistanceRestraint<AtomAddressInternal>( restraint );
			
			// classify the restraints
			Element leftElement = getHeavyElement( protein, bondGraphs, restraint.getLefts() );
			
			// left side
			if( leftElement == Element.Carbon )
			{
				reassignedRestraint.setLefts( getNearbyAddresses2D(
					protein,
					restraint.getLefts(),
					hydrogenShifts,
					carbonPairs,
					hydrogenWindowSize,
					carbonWindowSize
				) );
			}
			else if( leftElement == Element.Nitrogen )
			{
				reassignedRestraint.setLefts( getNearbyAddresses2D(
					protein,
					restraint.getLefts(),
					hydrogenShifts,
					nitrogenPairs,
					hydrogenWindowSize,
					nitrogenWindowSize
				) );
			}
			else
			{
				assert( false ) : "Unknown element: " + leftElement;
			}
			
			// right side
			reassignedRestraint.setRights( getNearbyAddresses1D(
				protein,
				restraint.getRights(),
				hydrogenShifts,
				hydrogenWindowSize
			) );
			
			reassignedRestraints.add( reassignedRestraint );
			
			restorePseudoatoms( protein, restraint, reassignedRestraint );
		}
		
		return reassignedRestraints;
	}
	
	
	/**************************
	 *   Static Functions
	 **************************/
	
	private static Set<AtomAddressInternal> getNearbyAddresses1D( Protein protein, Set<AtomAddressInternal> addresses, List<MappedChemicalShift> hydrogenShifts, double hydrogenWindowSize )
	{
		HashSet<AtomAddressInternal> relaxedAddresses = new HashSet<AtomAddressInternal>();
		for( AtomAddressInternal address : addresses )
		{
			// look up the target atom
			MappedChemicalShift target = getShiftFromHydrogen( address, hydrogenShifts );
			if( target != null )
			{
				// add the shift's addresses
				relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), target.getAddresses() ) );
				
				// add additional nearby addresses
				relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), getAddressesNearShift( target, hydrogenShifts, hydrogenWindowSize ) ) );
			}
			else if( isPseudoatom( protein, address ) )
			{
				// search the sub-atoms of the pseudoatom
				boolean foundAny = false;
				for( AtomAddressInternal subAddress : getSubAddresses( protein, address ) )
				{
					// re-add the original address
					relaxedAddresses.add( changeToSubunit( address.getSubunitId(), subAddress ) );
					
					// add additional nearby addresses if any
					target = getShiftFromHydrogen( subAddress, hydrogenShifts );
					if( target != null )
					{
						Set<AtomAddressInternal> nearbyAddresses = getAddressesNearShift( target, hydrogenShifts, hydrogenWindowSize );
						relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), nearbyAddresses ) );
						foundAny = true;
					}
				}
				if( !foundAny )
				{
					System.err.println( "No Shift for sub-atoms under atom: " + AddressMapper.mapAddress( protein, address ) );
				}
			}
			else
			{
				System.err.println( "No Shift for atom: " + AddressMapper.mapAddress( protein, address ) + " and no sub-atoms." );
			}
		}
		return relaxedAddresses;
	}
	
	private static Set<AtomAddressInternal> getNearbyAddresses2D( Protein protein, Set<AtomAddressInternal> addresses, List<MappedChemicalShift> hydrogenShifts, List<MappedChemicalShiftPair> heavyPairs, double hydrogenWindowSize, double heavyWindowSize )
	{
		HashSet<AtomAddressInternal> relaxedAddresses = new HashSet<AtomAddressInternal>();
		for( AtomAddressInternal address : addresses )
		{
			// look up hydrogen in the heavy pairs
			MappedChemicalShiftPair target = getPairFromHydrogen( address, heavyPairs );
			if( target != null )
			{
				// add the shift's addresses
				relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), target.getHydrogenShift().getAddresses() ) );
				
				// add additional nearby addresses
				relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), getAddressesNearPair( target, heavyPairs, hydrogenWindowSize, heavyWindowSize ) ) );
			}
			else if( isPseudoatom( protein, address ) )
			{
				// search the sub-atoms of the pseudoatom, if any
				boolean foundAny = false;
				for( AtomAddressInternal subAddress : getSubAddresses( protein, address ) )
				{
					// re-add the original address
					relaxedAddresses.add( changeToSubunit( address.getSubunitId(), subAddress ) );
					
					// add additional nearby addresses if any
					target = getPairFromHydrogen( subAddress, heavyPairs );
					if( target != null )
					{
						Set<AtomAddressInternal> nearbyAddresses = getAddressesNearPair( target, heavyPairs, hydrogenWindowSize, heavyWindowSize );
						relaxedAddresses.addAll( changeToSubunit( address.getSubunitId(), nearbyAddresses ) );
						foundAny = true;
					}
				}
				if( !foundAny )
				{
					System.err.println( "No Pair for sub-atoms under atom: " + AddressMapper.mapAddress( protein, address ) );
				}
			}
			else
			{
				// do a 1D search
				return getNearbyAddresses1D( protein, addresses, hydrogenShifts, hydrogenWindowSize );
			}
		}
		return relaxedAddresses;
	}
	
	private static MappedChemicalShift getShiftFromHydrogen( AtomAddressInternal address, List<MappedChemicalShift> hydrogenShifts )
	{
		// a simple linear scan should be fast enough for now
		// perhaps a hash table would make sense for larger lists
		for( MappedChemicalShift shift : hydrogenShifts )
		{
			for( AtomAddressInternal shiftAddress : shift.getAddresses() )
			{
				if( shiftAddress.equalsResidueAtom( address ) )
				{
					return shift;
				}
			}
		}
		
		return null;
	}
	
	private static MappedChemicalShiftPair getPairFromHydrogen( AtomAddressInternal address, List<MappedChemicalShiftPair> pairs )
	{
		for( MappedChemicalShiftPair pair : pairs )
		{
			for( AtomAddressInternal shiftAddress : pair.getHydrogenShift().getAddresses() )
			{
				if( shiftAddress.equalsResidueAtom( address ) )
				{
					return pair;
				}
			}
		}
		return null;
	}
	
	private static HashSet<AtomAddressInternal> getAddressesNearShift( MappedChemicalShift target, List<MappedChemicalShift> hydrogenShifts, double hydrogenWindowSize )
	{
		HashSet<AtomAddressInternal> nearbyAddresses = new HashSet<AtomAddressInternal>();
				
		// NOTE: this could be sped up to O(logn) using a geometric algorithm
		// but it's probably not worth the trouble to implement
		for( MappedChemicalShift shift : hydrogenShifts )
		{
			// don't add yourself
			if( shift == target )
			{
				continue;
			}
			
			if( isNearby1D( shift, target, hydrogenWindowSize ) )
			{
				nearbyAddresses.addAll( shift.getAddresses() );
			}
		}
		return nearbyAddresses;
	}
	
	private static HashSet<AtomAddressInternal> getAddressesNearPair( MappedChemicalShiftPair target, List<MappedChemicalShiftPair> heavyPairs, double hydrogenWindowSize, double heavyWindowSize )
	{
		HashSet<AtomAddressInternal> nearbyAddresses = new HashSet<AtomAddressInternal>();
		for( MappedChemicalShiftPair pair : heavyPairs )
		{
			// don't add yourself
			if( pair == target )
			{
				continue;
			}
			
			if( isNearby2D( target, pair, hydrogenWindowSize, heavyWindowSize ) )
			{
				nearbyAddresses.addAll( pair.getHydrogenShift().getAddresses() );
			}
		}
		return nearbyAddresses;
	}
	
	private static boolean isNearby1D( MappedChemicalShift left, MappedChemicalShift right, double hydrogenWindowSize )
	{
		return getDist( left, right ) <= hydrogenWindowSize;
	}
	
	private static boolean isNearby2D( MappedChemicalShiftPair left, MappedChemicalShiftPair right, double hydrogenWindowSize, double heavyWindowSize )
	{
		return getDist( left.getHydrogenShift(), right.getHydrogenShift() ) <= hydrogenWindowSize
			&& getDist( left.getHeavyShift(), right.getHeavyShift() ) <= heavyWindowSize;
	}
	
	private static double getDist( MappedChemicalShift left, MappedChemicalShift right )
	{
		return Math.abs( left.getValue() - right.getValue() );
	}
	
	private static boolean isPseudoatom( Protein protein, AtomAddressInternal address )
	{
		return PseudoatomBuilder.getPseudoatoms().isPseudoatom(
			protein.getResidue( address ).getAminoAcid(),
			protein.getAtom( address ).getName()
		);
	}
	
	private static ArrayList<AtomAddressInternal> getSubAddresses( Protein protein, AtomAddressInternal address )
	{
		// try to get the sub-atoms if this is a pseudoatom
		Residue residue = protein.getResidue( address );
		String atomName = protein.getAtom( address ).getName();
		
		// build a list of additional atoms to search for
		ArrayList<String> subatomNames = PseudoatomBuilder.getPseudoatoms().getAtoms( residue.getAminoAcid(), atomName );
		
		/* HACKHACK:
			For valine and leucine, we could have combined methyls.
			We need to search for both methyl groups before looking at each sub-atom
		*/
		if( residue.getAminoAcid() == AminoAcid.Leucine && atomName.equalsIgnoreCase( "QD" ) )
		{
			subatomNames = Transformer.toArrayList( "md1", "md2" );
		}
		else if( residue.getAminoAcid() == AminoAcid.Valine && atomName.equalsIgnoreCase( "QG" ) )
		{
			subatomNames = Transformer.toArrayList( "mg1", "mg2" );
		}
		
		// convert to an atom list
		ArrayList<AtomAddressInternal> subAddresses = new ArrayList<AtomAddressInternal>();
		for( String name : subatomNames )
		{
			subAddresses.add( new AtomAddressInternal(
				address.getSubunitId(),
				residue.getId(),
				residue.getAtomByName( name ).getId()
			) );
		}
		return subAddresses;
	}
	
	private static HashSet<AtomAddressInternal> changeToSubunit( int subunitId, Set<AtomAddressInternal> addresses )
	{
		HashSet<AtomAddressInternal> changedAddresses = new HashSet<AtomAddressInternal>( addresses.size() );
		for( AtomAddressInternal address : addresses )
		{
			changedAddresses.add( changeToSubunit( subunitId, address ) );
		}
		return changedAddresses;
	}
	
	private static AtomAddressInternal changeToSubunit( int subunitId, AtomAddressInternal address )
	{
		AtomAddressInternal changedAddress = new AtomAddressInternal( address );
		changedAddress.setSubunitId( subunitId );
		return changedAddress;
	}
	
	private static Element getHeavyElement( Protein protein, List<BondGraph> bondGraphs, Collection<AtomAddressInternal> hydrogenAddresses )
	{
		// NOTE: assume all addresses in this list are bound to the same heavy atom
		AtomAddressInternal hydrogenAddress = hydrogenAddresses.iterator().next();
		ArrayList<Bond> bonds = bondGraphs.get( hydrogenAddress.getSubunitId() ).getBonds( hydrogenAddress );
		assert( bonds != null && bonds.size() == 1 )
			: "Not single-bonded hydrogen? : " + AddressMapper.mapAddress( protein, hydrogenAddress );
		return protein.getAtom( bonds.get( 0 ).getOtherAddress( hydrogenAddress ) ).getElement();
	}
	
	private static Set<AtomAddressInternal> getPseudoatomAddresses( Protein protein, Set<AtomAddressInternal> addresses )
	{
		HashSet<AtomAddressInternal> pseudoatomAddresses = new HashSet<AtomAddressInternal>();
		for( AtomAddressInternal address : addresses )
		{
			if( isPseudoatom( protein, address ) )
			{
				pseudoatomAddresses.add( address );
			}
		}
		return pseudoatomAddresses;
	}
	
	private static void restorePseudoatoms( Protein protein, DistanceRestraint<AtomAddressInternal> restraint, DistanceRestraint<AtomAddressInternal> reassignedRestraint )
	{
		restorePseudoatom( protein, restraint.getLefts(), reassignedRestraint.getLefts() );
		restorePseudoatom( protein, restraint.getRights(), reassignedRestraint.getRights() );
	}
	
	private static void restorePseudoatom( Protein protein, Set<AtomAddressInternal> addresses, Set<AtomAddressInternal> reassignedAddresses )
	{
		for( AtomAddressInternal pseudoatomAddress : getPseudoatomAddresses( protein, addresses ) )
		{
			Residue residue = protein.getResidue( pseudoatomAddress );
			String atomName = protein.getAtom( pseudoatomAddress ).getName();
			
			// get the sub-addresses for this pseudoatom
			List<String> subatomNames = PseudoatomBuilder.getPseudoatoms().getAtoms( residue.getAminoAcid(), atomName );
			
			// HACKHACK: combined methyls need to be handled specially
			if( residue.getAminoAcid() == AminoAcid.Leucine && atomName.equalsIgnoreCase( "QD" ) )
			{
				subatomNames = Transformer.toArrayList( "md1", "md2" );
			}
			else if( residue.getAminoAcid() == AminoAcid.Valine && atomName.equalsIgnoreCase( "QG" ) )
			{
				subatomNames = Transformer.toArrayList( "mg1", "mg2" );
			}
			
			// find all the sub-atoms for this pseudoatom
			ArrayList<AtomAddressInternal> foundAddresses = new ArrayList<AtomAddressInternal>();
			for( AtomAddressInternal address : reassignedAddresses )
			{
				// skip atoms not in this residue
				if( address.getResidueId() != pseudoatomAddress.getResidueId() )
				{
					continue;
				}
				
				for( String subatomName : subatomNames )
				{
					if( subatomName.equalsIgnoreCase( protein.getAtom( address ).getName() ) )
					{
						foundAddresses.add( address );
						break;
					}
				}
			}
			
			// did we find them all?
			if( foundAddresses.size() == subatomNames.size() )
			{
				// restore the pseudoatom
				reassignedAddresses.removeAll( foundAddresses );
				reassignedAddresses.add( changeToSubunit( foundAddresses.get( 0 ).getSubunitId(), pseudoatomAddress ) );
			}
		}
	}
}
