import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
%matplotlib notebook
import pandas as pd
import h5py
import afinesanalysis.afinesanalysis as aa

"""
Explanation: This notebook is meant to test the functions written in python to analyze AFiNES simulation output. It's going to be messy, but so it goes...
End of explanation
"""


# Change this line if on a different system
dataFolder = '/media/daniel/storage1/local_LLM_Danny/rupture/a-1.0-p-1.0/txt_stack/'
filename = 'actins.txt'
dt = 1
domainSize = 50
nbins = 10
minpts = 10
dr = 1
rbfFunc = 'gaussian'
rbfEps = 5
savestuff = False
nGridPts = np.ceil(domainSize/dr) + 1
txy = aa.readData(dataFolder, filename)

# Load in hdf5 file

actinData = h5py.File(os.path.join(dataFolder, 'actinsParsed.hdf5'), 'a')
actins = actinData.create_dataset('actinData', data=txy)

del actinData

"""
Explanation: First a test of the readData function that reads in .txt files output by afines. Right now (9/12/2017) this only works for actins.txt
End of explanation
"""


# Don't include the id data when getting the divergence
[xx, yy, uut, vvt] = aa.interpolateVelocity(txy[:,:,:2], dt, domainSize, nbins, minpts, dr, rbfFunc, rbfEps, savestuff)

divV = aa.massWeightedVelocityDivergence(txy[:,:,:2], uut, vvt, dt)

interpedXVel = actinData.create_dataset('interpedXVel', data=uut)
interpedYVel = actinData.create_dataset('interpedYVel', data=vvt)
xGrid = actinData.create_dataset('xGrid', data=xx)
yGrid = actinData.create_dataset('yGrid', data=yy)

actinData.close()

divV[::10].shape

fig, ax = plt.subplots()
ax.plot(np.linspace(1,99,100), divV[::10])
ax.set_xlabel('time (s)')
ax.set_ylabel(r'$\int \rho \langle \nabla \cdot \vec{v} \rangle$')

fig.savefig(os.path.join(dataFolder, 'divergence.png'))

fig, ax = plt.subplots()
ax.pcolor(xx,yy,massWeight.T)
ax.set_aspect('equal')

fig, ax = plt.subplots()
ax.quiver(xx, yy ,uut[-10,...], vvt[-10,...])

divergence = np.gradient(uut)[2] + np.gradient(vvt)[1]
curl = np.gradient(uut)[2] * np.gradient(vvt)[1] - np.gradient(uut)[1] + np.gradient(vvt)[2]

np.sum(divergence[200,...])

for frame, div in enumerate(divergence[::10,...]):
    fig, ax = plt.subplots()
    cax1 = ax.pcolor(xx,yy, div, cmap='cool', vmax=0.3, vmin=-0.3)
    ax.set_aspect('equal')
    ax.quiver(xx,yy, uut[frame,...], vvt[frame,...])
    cbar1 = fig.colorbar(cax1, ax=ax)
    ax.set_xlabel('x ($\mu$m)')
    ax.set_ylabel('y ($\mu$m)')
    ax.set_title('Divergence Field, sum=%2f' % np.sum(div))
    fig.tight_layout()
    fig.savefig(os.path.join(dataFolder, 'frame{0}.png'.format(frame)))


fig, ax = plt.subplots()
divergenceSeries = -np.nansum(np.nansum(divergence, axis=2), axis=1)

plt.plot(divergenceSeries)

#fig, ax = plt.subplots()
#cax1 = ax.pcolor(xx, yy, divergence, cmap='inferno')
#ax.set_aspect('equal')
#ax.quiver(xx,yy, uu, vv)
#cbar1 = fig.colorbar(cax1, ax=ax)
#ax.set_xlabel('x ($\mu$m)')
#ax.set_ylabel('y ($\mu$m)')
#ax.set_title('Divergence Field, sum=%2f' % np.sum(divergence))
#fig.tight_layout()

#fig, ax = plt.subplots()
#cax2 = ax.pcolor(xx, yy, curl, cmap='cool')
#ax.quiver(xx,yy,uu,vv)
#ax.set_aspect('equal')
#cbar2 = fig.colorbar(cax2, ax=ax)
#ax.set_xlabel('x ($\mu$m)')
#ax.set_ylabel('y ($\mu$m)')
#ax.set_title('Curl Field, sum=%2f' % np.sum(curl))
#fig.tight_layout()

"""
Explanation: Now calculate the velocity divergence. This is essentially stolen from here
End of explanation
"""


