import numpy
from analysisNEW import *
import pylab

#
# this code is to see what are the minimum number of surfaces needed 
# to get a reasonable broken bond parameter fit
#

def fitSampleSurfaces(fileName,sampleSize,p0):

	indices, e_unrelaxed, e_relaxed, sizes = loadFileIntoLists(fileName)

	e_relaxed_sorted = sorted(e_relaxed)

	N_interval = len(e_relaxed)/sampleSize
	sample_indices=[]
	sample_energies=[]

	for i in range(0,sampleSize):
		e_rel = e_relaxed_sorted[i*N_interval]
		sample_energies.append(e_rel)
		loc = e_relaxed.index(e_rel)
		sample_indices.append(indices[loc])
	
	sample_indices = numpy.array(sample_indices)
	sample_energies = numpy.array(sample_energies)
	
	bfparams, flag = leastsq(residual,p0,args=(sample_indices,sample_energies))
	
	return bfparams, sample_indices	


def plotParamsVsSample(fileName,p0,bfstandard):

	indices, e_unrelaxed, e_relaxed, sizes = loadFileIntoLists(fileName)

	N_max = len(e_relaxed)-1

	sizelist = numpy.arange(5,N_max,10)	 

	percent_diff_p0=[]
	percent_diff_p1=[]

	for size in sizelist:
		bfparams, sample_indices = fitSampleSurfaces(fileName,size,p0)
		err_p0 = abs(bfparams[0]-bfstandard[0])/bfstandard[0]
		err_p1 = abs((bfparams[1]-bfstandard[1])/bfstandard[1])
		percent_diff_p0.append(err_p0)
		percent_diff_p1.append(err_p1)

	pylab.plot(sizelist,percent_diff_p0,'r.')
	pylab.plot(sizelist,percent_diff_p1,'b.')

	pylab.show()

def fitListSurfaces(fileName,sample_indices,p0,corrections=0):
	
	indices, e_unrelaxed, e_relaxed, sizes = loadFileIntoLists(fileName)

	sample_energies=[]
	
	for ind in sample_indices:
		curr_index = indices.index(ind)
		sample_energies.append(e_relaxed[curr_index])

	bfparams, flag = leastsq(residual,p0,args=(numpy.array(sample_indices),sample_energies,corrections))

	return bfparams