import numpy as np
import sys
from scipy.interpolate import griddata
import matplotlib.pyplot as plt

x = []
y = []
z = []
s = []

with open('strainmapXYZ_R10mm.txt') as f:
    file = f.readlines()[1:]
    
for line in file:
    a = line.split()
    x.append(float(a[1]))
    y.append(float(a[2]))
    z.append(float(a[3]))
    s.append(float(a[4]))

xmin=min(x)
xmax=max(x)
ymin=min(y)
ymax=max(y)

# create grid based on raw data of X,Y locations and strain data
grid_x, grid_y = np.mgrid[xmin:xmax:3000j, ymin:ymax:3000j]
x_test=np.linspace(xmin,xmax,3000)
y_test=np.linspace(ymin,ymax,3000)
x1=min(range(len(x_test)), key=lambda i: abs(x_test[i]+1e-6))
x2=min(range(len(x_test)), key=lambda i: abs(x_test[i]-1e-6))
y1=min(range(len(y_test)), key=lambda i: abs(y_test[i]-6e-6))
y2=min(range(len(y_test)), key=lambda i: abs(y_test[i]-7e-6))
x0=np.linspace(x1,x2,x2-x1+1, dtype=int)
y0=np.linspace(y1,y2,y2-y1+1, dtype=int)

# Interpolate Strain data surface map
s = griddata((x,y), s, (grid_x, grid_y), method='cubic')
# Remove grating pitch section where is air - Strain is set to 0
for i in x0:
    for j in y0:
        if -1e-6 <= x_test[i] <= 1e-6 and 6e-6 <= y_test[j] <= 7e-6:
            s[i,j]=0
s[np.isnan(s)] = 0

#save interpolated strain data
with open('interpolated_data.txt', 'w') as fd:
    for row in s.T:
        fd.write(' '.join([str(a) for a in row]) + '\n')

# plot strain distribution XY map
plt.figure("strain")
#plt.imshow(s.T, extent=(xmin*1e6, xmax*1e6, ymin*1e6, ymax*1e6), origin='lower')
plt.imshow(s.T, extent=(xmin*1e6, xmax*1e6, ymin*1e6, ymax*1e6), origin='lower')
plt.colorbar(label="microstrain", orientation="vertical")
plt.clim(1e-6, 1e-4)
plt.title('interpolated Strain map')
plt.xlabel('X-axis position (μm)')
plt.ylabel('Y-axis position (μm)')
plt.show()