#!/usr/bin/env python

#########################################
import matplotlib.pyplot as plt         #
import numpy as np                      #
from math import sqrt, radians, cos	#
from mpl_toolkits.axes_grid.inset_locator import inset_axes
#########################################

plt.rcParams.update({'font.size': 12})


#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Raw K+W data
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

KW_expt  = []
KW_trace = []

outfile = open('KW_expt.dat','r')
data_exp = outfile.readlines()
outfile.close()

for i in range(len(data_exp)):
        dexp = ([float(x) for x in data_exp[i].split()])
        KW_expt.append(dexp)

outfile = open('KW_trace.dat','r')
data_trace = outfile.readlines()
outfile.close()

for i in range(len(data_trace)):
        dtra = ([float(x) for x in data_trace[i].split()])
        KW_trace.append(dtra)

KW_expt = np.array(KW_expt)
KW_trace = np.array(KW_trace)

KW_expt = np.transpose(KW_expt)
KW_trace = np.transpose(KW_trace)

ax = plt.subplot(2,2,1)
plt.plot(KW_expt[0]-0.38, KW_expt[1], 'k', lw = 1.5)
#plt.xlabel('Time (s)')
plt.ylabel('Partial pressure (Arb. units)')
plt.xlim(-130,100)
plt.yticks(())
plt.annotate('A', xy = (-120,0.05e-10), xytext = (-120,0.05e-10), fontsize =16)
inset_axes = inset_axes(ax,
                    width="50%", # width = 30% of parent_bbox
                    height=1.2, # height : 1 inch
                    loc=10)

plt.plot(KW_expt[0]-0.38, KW_expt[1], 'k', lw = 1.5)
plt.xlim(-5,15)
plt.ylim(1.32e-10,1.45e-10)
plt.yticks(())
plt.xticks(())


plt.subplot(2,2,3)
plt.plot(KW_expt[0]-0.38, KW_expt[2], 'k', lw = 1.5, label = 'Expt')
plt.plot(KW_trace[0], KW_trace[1], 'r--', lw = 3, label = 'Fit')
plt.xlabel('Time (s)')
plt.ylabel('Sticking coefficient')
plt.xlim(-5,20)
plt.ylim(0,0.09)
plt.yticks(np.arange(0.,0.1,0.02))
plt.annotate('B', xy = (-4,0.005), xytext = (-4,0.005), fontsize =16)
lg = plt.legend(loc = 'best', prop={'size': 12})
lg.draw_frame(False)
plt.subplots_adjust(left=0.1, bottom=0.08, right=0.95, top=0.99, wspace=0.15, hspace=0.15)
plt.savefig('Fig2_KW_data.png',bbox_inches='tight', dpi = 300)
plt.show()