import numpy as np
import matplotlib.pyplot as plt
import sys, os
import scipy.stats as ss
import scipy as sp
#import pylab as P


if len(sys.argv)>2:
	filenames = sys.argv[1:]
else:
	filenames = [sys.argv[1]]


endpoint = 0
j = 0
dp = np.zeros(len(filenames))
Fmax = np.zeros(len(filenames))
Wstar = np.zeros(len(filenames))
total_travel = np.zeros(len(filenames))
Fpoly = np.zeros((len(filenames),3))
kpeb = np.zeros(len(filenames))
err_peb = np.zeros(len(filenames))
E_rec = [[] for i in range(len(filenames))]
for loadname in filenames:
	# Open the file again to read the pebble diameter
	f = open(loadname)
	content = f.readlines()
	dp1 = content[1]
	dp[j] = float(dp1[15:(dp1.find("mm")-2)])/1000.
	j +=1

# plt.figure(4)
# n, bins, patches = plt.hist(dp*1000, histtype='stepfilled')
# plt.setp(patches, 'facecolor', 'g', 'alpha', 0.75)
# plt.title('KIT pebbles (0.2 ~ 0.6 mm) at room temperature')
# plt.xlabel('Pebble diameter (mm)')
# plt.ylabel('Count')


j = 0


k = np.linspace(0.01,1,1000)

for loadname in filenames:
	#print loadname
	if dp[j] < 0.45/1000:
		plotColor = 'g'
	elif dp[j] >= 0.45/1000 and dp[j] < 0.55/1000:
		plotColor = 'r'
	else:
		plotColor = 'k'
	filedata = np.loadtxt(loadname,skiprows=6,delimiter=',')
	s = filedata[:,0]
	F = filedata[:,1]
	
	# Search for the maximum force, discard the rest of the plot
	# after this point
	maxIndex = max(enumerate(F),key=lambda x: x[1])[0]
	F = F[0:maxIndex]	
	s = s[0:maxIndex]


	total_travel[j] = s[-1]
	# Make a polynomial fit
	sfit = np.linspace(0,s[-1],1001)
	Fpoly = np.poly1d(np.polyfit(s,F,3, full=False))
	Ffit = Fpoly(sfit)

	# Find crush force
	Fmax[j] = F[-1]


	# endpoint for the plot
	if s[-1] > endpoint:
		endpoint = s[-1]
	if s[-1] > 0.035:
		print "Warning, "+loadname+" has very high strain (perhaps a bad pebble)"
	# Integrate Force over displacement for strain energy
	W = np.trapz(F,s)

	
	# Find strain energy per pebble volume
	Wstar[j] = W/(4/3. * np.pi * (dp[j]/(2.))**3.)
	

	Estand = 220.e9
	nustand = 0.27
	Fn0 = 0.1
	nupeb = 0.24
	err = 1e5
	Epebbulk = 120.e9
	for i in k:
		# Using Hertz theory, find the predicted force for 
		# given displacement
		Epeb = i*Epebbulk
		Estar = 1./((1.-nupeb**2)/Epeb + (1.-nustand**2)/Estand)
		# in the experiment, there is overlap causing the pre-load
		# use that preload to back-out the initial overlap
		so = ((3*Fn0/Estar)**2*(1/dp[j]))**(1./3)

		# Hertz force (with corrected initial strain)
		Fhertz = (1./3)*Estar*np.sqrt(dp[j]*(so+sfit/1000.)**3)+Fn0
		diff = Fhertz - Ffit
		
		erri = np.linalg.norm(diff)
		
		if erri < err:
			err = erri
			kpeb[j] = i
			err_peb[j] = erri
	if s[-1] > 0.05:
		print loadname + " pebble has very long strain, review this data"
	Epeb = kpeb[j]*Epebbulk
	E_rec[j] = "E = "+ str(Epeb/10**9) + " GPa"
	Estar = 1./((1.-nupeb**2)/Epeb + (1.-nustand**2)/Estand)
	Fhertz = (1./3)*Estar*np.sqrt(dp[j]*(so+s/1000.)**3)+Fn0
	
	# Plot all forces over displacements
	labelstring = 'Pebble '+str(j)
	labelstringH = "Hertz theory (E = "+str(Epeb/10.**9)+"GPa )"
	plt.figure(1)
	plt.plot(s, F, color='k')
	plt.plot(s, Fhertz, color='k',linestyle='--')
	plt.legend(('Experimental',"Hertz theory"))
	plt.xlabel('Standard travel (mm)')
	plt.ylabel('Standard force (N)')
	plt.xlim((0, endpoint))
	plt.ylim((0,50))
	
	j+=1
plt.savefig('exp_v_hertz')



#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# CRUSH FORCE HISTOGRAM PLOT
#
#
plt.figure(100)
plt.hist(Fmax)
plt.xlabel('Crush force (N)')
plt.ylabel('Count')
print 'F_max-average = ' + str(np.average(Fmax)) + '\n'
plt.savefig('fmax')

print 'average total travel =' + str(np.average(total_travel)) + '\n'
print "Young moduli = " + str(kpeb*Epebbulk/10**9) + '\n'
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~





#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# PEBBLE DIAMETER VS. K SCATTER
#
#
plt.figure(101)
plt.scatter(dp*1000, kpeb, s=60, color='g')
plt.xlim(min(dp*1000), max(dp*1000))
plt.ylabel('Elasticity reduction factor')
plt.xlabel('Pebble diameter (mm)')
plt.savefig('f-v-d')
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~



#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# PEBBLE DIAMETER VS. FMAX SCATTER
#
#
plt.figure(102)
plt.scatter(dp*1000, Fmax, s=60, color='g')
plt.xlim(min(dp*1000), max(dp*1000))
plt.ylabel('Crush force (N)')
plt.xlabel('Pebble diameter (mm)')
plt.savefig('k-v-d')
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~





#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# K HISTOGRAM PLOT
#
#

plt.figure(103)

n, bins, patches = plt.hist(kpeb, bins=10, histtype='stepfilled', normed=True)
plt.xlim((0,1))
plt.setp(patches, 'facecolor', 'g', 'alpha', 0.75)
plt.xlabel('Elasticity reduction factor')
plt.ylabel('Count')
plt.savefig('k_hist')
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# These made an f-distribution to fit kpeb -- but i'm not using it
# from scipy.stats import f
# dfn, dfd, loc, other = f.fit(kpeb, loc=min(kpeb))
# x = np.linspace(0,1,100)
# pdf = f.pdf(x, dfn, dfd, loc=loc, scale=other)
# r = f.rvs(dfn, dfd, loc=loc,size=20)
# plt.hist(r, bins=np.arange(loc, 1.1, 0.05), normed=False, alpha=0.4, color='k')
# f.pdf(x, dfn, dfd, loc=loc, scale=other)
# plt.plot(x,pdf)

plt.show()