#-------------------------------------------------------------------------------
# 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
#-------------------------------------------------------------------------------


import jvm, share
from figure import Figure
from matplotlib.patches import Circle
from matplotlib.patches import Ellipse

share.getJvm().start()


# UNDONE: refactor this, move into share, and call from GB1 script

def plotShifts( pathOut, shifts ):
	
	# render the plot
	fig = Figure()
	for shift in shifts:
		circle = Circle(
			( shift.getValue(), shift.getElement().ordinal() ),
			radius=shift.getError(),
			facecolor='none',
			linewidth=0.2
		)
		fig.plot.add_patch( circle )
	fig.plot.axis( "equal" )
	fig.plot.set_yticks( [] )
	fig.plot.set_yticklabels( [] )
	fig.plot.set_xlabel( "Hydrogen Chemical Shift (ppm)" )
	
	# save the figure
	fig.save( pathOut )
	print "\tWrote chemical shifts plot to:\n\t\t%s" % ( pathOut )


def plotPairs( pathOut, pairs, element ):
	
	# render the plot
	fig = Figure()
	for pair in pairs:
		ellipse = Ellipse(
			xy=( pair.getHydrogenShift().getValue(), pair.getHeavyShift().getValue() ),
			width=pair.getHydrogenShift().getError() * 2,
			height=pair.getHeavyShift().getError() * 2,
			facecolor='none',
			linewidth=0.4
		)
		fig.plot.add_patch( ellipse )
	fig.plot.autoscale_view()
	fig.plot.set_xlabel( "Hydrogen Chemical Shift (ppm)" )
	fig.plot.set_ylabel( "%s Chemical Shift (ppm)" % element )
	
	# save the figure
	fig.save( pathOut )
	print "\tWrote Hydrogen-%s chemical shifts plot to:\n\t\t%s" % ( element, pathOut )


def totalNumberOfAssignments( restraints ):
	sum = 0
	for restraint in restraints:
		sum += restraint.getNumAssignments()
	return sum

	
def avgNumberOfAssignments( restraints ):
	return float( totalNumberOfAssignments( restraints ) ) / float( restraints.size() )


def run( pathOutRelaxedNoes, pathInMonomer, pathInNoes, pathInShifts, numSubunits ):
	
	# load the protein and noe data
	monomer = share.loadProtein( pathInMonomer ).getSubunits().get( 0 )
	oligomer = share.cloneMonomer( monomer, numSubunits )
	noes = share.loadDistanceRestraints( pathInNoes )
	share.interpretDistanceRestraintsWithPseudoatoms( noes, oligomer )
	restraints = share.mapDistanceRestraintsToProtein( noes, oligomer )
	share.printRestraintsStats( restraints )
	share.reportMismatchedRestraints( restraints, numSubunits )
	restraints = share.filterRestraintsIntersubunit( restraints )
	
	# load and process the chemical shifts
	shifts = share.loadShifts( pathInShifts )
	hydrogenShifts = share.f.nmr.ChemicalShiftMapper.filter( shifts, share.f.protein.Element.Hydrogen )
	carbonShifts = share.f.nmr.ChemicalShiftMapper.filter( shifts, share.f.protein.Element.Carbon )
	nitrogenShifts = share.f.nmr.ChemicalShiftMapper.filter( shifts, share.f.protein.Element.Nitrogen )
	share.interpretShiftsWithPseudoatoms( hydrogenShifts, oligomer )
	mappedHydrogenShifts = share.mapShifts( hydrogenShifts, oligomer )
	mappedCarbonShifts = share.mapShifts( carbonShifts, oligomer )
	mappedNitrogenShifts = share.mapShifts( nitrogenShifts, oligomer )
	
	# combine the heavy shifts
	mappedHeavyShifts = jvm.f.java.util.ArrayList()
	mappedHeavyShifts.addAll( mappedCarbonShifts )
	mappedHeavyShifts.addAll( mappedNitrogenShifts )
	
	# TEMP
	print "Num Hydrogen Shifts", mappedHydrogenShifts.size()
	print "Num Carbon Shifts", mappedCarbonShifts.size()
	print "Num Nitrogen Shifts", mappedNitrogenShifts.size()
	
	carbonPairs = share.f.nmr.ChemicalShiftMapper.associatePairs( oligomer, mappedHydrogenShifts, mappedCarbonShifts, share.f.protein.Element.Carbon )
	nitrogenPairs = share.f.nmr.ChemicalShiftMapper.associatePairs( oligomer, mappedHydrogenShifts, mappedNitrogenShifts, share.f.protein.Element.Nitrogen )
	orphanedHydrogenShifts = share.f.nmr.ChemicalShiftMapper.getOrphanedHydrogenShifts( oligomer, mappedHydrogenShifts, mappedHeavyShifts )
	
	# TEMP
	print "Num Orphaned Hydrogen Shifts", orphanedHydrogenShifts.size()
	print "Num Carbon Pairs", carbonPairs.size()
	print "Num Nitrogen Pairs", nitrogenPairs.size()
	
	plotPairs( dirOut + "/carbonPairs.pdf", carbonPairs, "Carbon" )
	plotPairs( dirOut + "/nitrogenPairs.pdf", nitrogenPairs, "Nitrogen" )
	plotShifts( dirOut + "/hydrogenShifts.pdf", hydrogenShifts ) 
	
	# reassign the NOEs
	hydrogenWindowSize = 0.05 # in ppm
	carbonWindowSize = 0.5
	nitrogenWindowSize = 0.5
	print "Reassigning NOEs..."
	print "\tAssignments before: total=%i, avg=%f" % ( totalNumberOfAssignments( restraints ), avgNumberOfAssignments( restraints ) )
	#relaxedRestraints = share.f.nmr.DistanceRestraintReassigner.reassign1D( oligomer, restraints, mappedHydrogenShifts, hydrogenWindowSize )
	relaxedRestraints = share.f.nmr.DistanceRestraintReassigner.reassignDouble3D(
		oligomer,
		restraints,
		mappedHydrogenShifts,
		carbonPairs,
		nitrogenPairs,
		hydrogenWindowSize,
		carbonWindowSize,
		nitrogenWindowSize
	)
	print "\tAssignments after (expanded pseudoatoms): total=%i, avg=%f" % ( totalNumberOfAssignments( relaxedRestraints ), avgNumberOfAssignments( relaxedRestraints ) )
	
	# unmap and save
	relaxedNoes = share.unmapDistanceRestraintsFromProtein( relaxedRestraints, oligomer, True )
	share.writeDistanceRestraints( pathOutRelaxedNoes, relaxedNoes )
	
	print "\tAssignments after (collapsed pseudoatoms): total=%i, avg=%f" % ( totalNumberOfAssignments( relaxedNoes ), avgNumberOfAssignments( relaxedNoes ) )
	
	# count the number of assignments
	numAssignments = 0
	for restraint in relaxedNoes:
		numAssignments += restraint.getNumAssignments()
	print "Total number of assignments: %d" % numAssignments
	

# GB1
dirOut = "/home/jeff/duke/donaldLab/workspace/jdshot/output/gb1-reassigned"
dirInput = "/home/jeff/duke/donaldLab/dshot/input"
run(
	dirOut + "/1Q10.reassigned.noe",
	dirInput + "/1Q10.monomer.protein",
	dirInput + "/1Q10.experimental.fixed.noe",
	dirInput + "/1Q10.experimental.shift",
	2
)
