#-------------------------------------------------------------------------------
# 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 share, math, numpy
from figure import Figure

jvm = share.getJvm()
jvm.start()

def loadRdcs( pathIn ):
	
	assigns = share.f.nmr.AssignReader().read( pathIn )
	rdcs = [assign.getNumbers().get( 0 ).floatValue() for assign in assigns]
	error = assigns.get( 0 ).getNumbers().get( 1 ).floatValue()
	return rdcs, error


def getPrefixAverage( rdcs, error ):

	currentRdcs = numpy.array( [ rdcs[0] ] )
	for i in range( 1, len( rdcs ) ):
		currentRdcs = numpy.append( currentRdcs, rdcs[i] )
		if currentRdcs.std() > error:
			return currentRdcs[0:i].mean()

	return currentRdcs.mean()

	
def analyze( pathOut, pathRdcs ):
	
	# read the RDCs
	( rdcs, error ) = loadRdcs( pathRdcs )
	numRdcs = len( rdcs )
	print "\tnum RDCs: %d" % numRdcs
	
	# most negative principal component (could be Dyy or Dzz)
	rdcs.sort()
	mostNegative = getPrefixAverage( rdcs, error )
	
	# most positive principal component
	rdcs.reverse()
	mostPositive = getPrefixAverage( rdcs, error )
	
	# render the plot
	fig = Figure()
	( counts, bins, _ ) = fig.plot.hist( rdcs, bins=25, facecolor="#004586" )
	fig.plot.set_xlabel( "RDC Value in Hz" )
	fig.plot.set_ylabel( "Count" )
	
	# compute middle principal component
	leftIndex = numpy.argmax( numpy.array( counts ) )
	rightIndex = leftIndex
	leftEdge = bins[leftIndex]
	while counts[rightIndex+1] == counts[leftIndex]:
		rightIndex += 1
	rightEdge = bins[rightIndex+1]
	set = numpy.array( [] )
	for rdc in rdcs:
		if rdc >= leftEdge and rdc <= rightEdge:
			set = numpy.append( set, rdc ) 
	middle = set.mean()
	
	# assign principal components
	components = numpy.array( [ mostNegative, middle, mostPositive ] )
	absComponents = numpy.abs( components )
	index = absComponents.argmax()
	Dzz = components[index] 
	absComponents[index] = 0
	index = absComponents.argmax()
	Dyy = components[index]
	absComponents[index] = 0
	index = absComponents.argmax()
	Dxx = components[index]
	pcSum = abs( Dxx + Dyy + Dzz )
	
	# compute rhombicity using linear least squares
	# Dzz = 2Da
	# Dyy = -Da - 3/2 Dr
	# Dxx = -Da + 3/2 Dr
	A = numpy.matrix( [[2, 0], [-1, -3/2], [-1, 3/2]] )
	b = numpy.matrix( [[Dzz], [Dyy], [Dxx]] )
	( x, residual, _, _ ) = numpy.linalg.lstsq( A, b )
	Da = x[0][0]
	Dr = x[1][0]
	R = Dr / Da
	
	print "\tPrincipal Components:\n\t\tDxx: %.2f\n\t\tDyy: %.2f\n\t\tDzz: %.2f\n\t\tSum: %.2f" % ( Dxx, Dyy, Dzz, pcSum )
	print "\tTensor Properties:\n\t\tDa: %.2f\n\t\tDr: %.2f\n\t\tR: %.2f\n\t\tResidual: %.2f" % ( Da, Dr, R, residual )
	
	# save the plot
	fig.save( pathOut )
	print "\tWrote RDC Histogram to:\n\t\t%s" % pathOut
	

# set up paths
dirData = "/home/jeff/duke/donaldLab/dshot/input"
dirResults = "/home/jeff/duke/donaldLab/noe+rdc/results"

print( "GB1" )
analyze(
	dirResults + "/GB1/rdcHistogram.png",
	dirData + "/1Q10.experimental.rdc"
)

print( "DAGK" )
analyze(
	dirResults + "/DAGK/rdcHistogram.png",
	dirData + "/2KDC.experimental.rdc"
)

print( "PLB" )
analyze(
	dirResults + "/PLB/rdcHistogram.png",
	dirData + "/1ZLL.experimental.rdc"
)

print( "M2" )
analyze(
	dirResults + "/M2/rdcHistogram.png",
	dirData + "/2RLF.experimental.rdc"
)

print( "BM2" )
analyze(
	dirResults + "/BM2/rdcHistogram.png",
	dirData + "/2KJ1.experimental.rdc"
)
