sys.path.insert(0,f"{script_dir}/python_lib")

fromplotsimport plotSpatialAll
fromplotsimport plotPhiHeatmap, plotFiForAd
fromutilsimport computFullGridVelocity, getGridPhi, getGridVelocityForSourceSink, getFiOfEmb, loadAddLayersVelocytoData

X_grid_full, V_grid_full, X_emb = computFullGridVelocity(ST, emb='spatial')
gridPhi = getGridPhi(X_grid_full, V_grid_full, R=2)
plotPhiHeatmap(gridPhi, X_grid_full, emb='spatial')

fig, axs = plt.subplots(1, 2, figsize=(10, 5))
params = dict(density=1, arrow_size=10, arrow_length=100, color=identity, size=45, alpha=0.75) 
ax = axs[0]
X_source, V_source = getGridVelocityForSourceSink(X_grid_full, V_grid_full, gridPhi, isSource=True)
scv.pl.velocity_embedding_grid(ST, basis='X_spatial', X_grid=X_source, V_grid=V_source, legend_loc='none', title='Source', ax=axs[0], show=False, **params)
#scv.pl.velocity_embedding_grid(ST, basis='X_spatial', legend_loc='none', title='Source', ax=axs[0], show=False,**params)

image = ST.uns['spatial']['ATC4']['images']['lowres'].copy()
ax.imshow(image, interpolation='none', extent=[0, image.shape[1]/sf, image.shape[0]/sf,0], zorder=-1)

ax = axs[1]
X_sink, V_sink = getGridVelocityForSourceSink(X_grid_full, V_grid_full, gridPhi, isSource=False)
scv.pl.velocity_embedding_grid(ST, basis='X_spatial', X_grid=X_sink, V_grid=V_sink, legend_loc='none', title='Sink', ax=axs[1], show=False, **params)

image = ST.uns['spatial']['ATC4']['images']['lowres'].copy()

ax.imshow(image, interpolation='none', extent=[0, image.shape[1]/sf, image.shape[0]/sf,0], zorder=-1)


#ax.legend(frameon=False, loc='upper right', fontsize=4, markerscale=1)
fig.legend(loc='right', fontsize=5, markerscale=1,frameon=False )
fig.subplots_adjust(left=0.15, right=0.85, bottom=0.15, top=0.85, wspace=0.2, hspace=0.2)

Fi_emb = getFiOfEmb(X_grid_full, gridPhi, X_emb)
ST.obs['Fi_spatial'] = Fi_emb
sc.pl.spatial(ST, color='Fi_spatial', cmap='bwr', vmin=-1, vmax=1)
sc.pl.violin(ST, keys='Fi_spatial', groupby='GraphST_clusters')